Enrich model test, remove useless __main__

.pylintrc

@@ -60,7 +60,12 @@ confidence=
 # --enable=similarities". If you want to run only the classes checker, but have
 # no Warning level messages displayed, use "--disable=all --enable=classes
 # --disable=W".
-disable=bad-continuation,no-value-for-parameter,unexpected-keyword-arg,redundant-keyword-arg
+disable=bad-continuation,
+		no-value-for-parameter,
+		unexpected-keyword-arg,
+		redundant-keyword-arg,
+		too-many-arguments,
+		too-many-locals
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option

tests/test_model.py

@@ -1,6 +1,9 @@
+from pathlib import Path
+
 import pytest
 import tensorflow as tf
 
+from tf2_yolov4.anchors import YOLOV4_ANCHORS
 from tf2_yolov4.model import YOLOv4
 
 
@@ -41,3 +44,85 @@ def test_model_instanciation_should_fail_with_input_shapes_not_multiple_of_32(
 ):
     with pytest.raises(ValueError):
         YOLOv4(input_shape, 80, [])
+
+
+def test_should_raise_error_if_weights_argument_is_unknown():
+    with pytest.raises(ValueError):
+        YOLOv4(input_shape=(416, 416, 3), num_classes=80, anchors=[], weights="unknown")
+
+
+def test_should_download_pretrained_weight_if_not_available(mocker):
+    mocker.patch("tf2_yolov4.model.csp_darknet53")
+    mocker.patch("tf2_yolov4.model.compute_normalized_anchors")
+    mocker.patch("tf2_yolov4.model.yolov4_neck")
+    mocker.patch("tf2_yolov4.model.yolov3_head")
+    mocker.patch(
+        "tf2_yolov4.model.tf.keras.Model"
+    ).return_value = mocker.sentinel.yolov4
+    mocker.sentinel.yolov4.load_weights = mocker.MagicMock()
+
+    mock_is_darknet_weights_called = mocker.patch(
+        "tf2_yolov4.model.is_darknet_weights_available"
+    )
+    mock_is_darknet_weights_called.return_value = False
+    mock_download_darknet_weights = mocker.patch(
+        "tf2_yolov4.model.download_darknet_weights"
+    )
+
+    YOLOv4(
+        input_shape=(416, 416, 3),
+        num_classes=80,
+        anchors=YOLOV4_ANCHORS,
+        weights="darknet",
+    )
+    mock_is_darknet_weights_called.assert_called_once_with()
+    mock_download_darknet_weights.assert_called_once_with(mocker.sentinel.yolov4)
+
+
+def test_should_load_pretrained_weights_if_available(mocker):
+    mocker.patch("tf2_yolov4.model.csp_darknet53")
+    mocker.patch("tf2_yolov4.model.compute_normalized_anchors")
+    mocker.patch("tf2_yolov4.model.yolov4_neck")
+    mocker.patch("tf2_yolov4.model.yolov3_head")
+    mocker.patch(
+        "tf2_yolov4.model.tf.keras.Model"
+    ).return_value = mocker.sentinel.yolov4
+    mocker.sentinel.yolov4.load_weights = mocker.MagicMock()
+
+    mock_is_darknet_weights_called = mocker.patch(
+        "tf2_yolov4.model.is_darknet_weights_available"
+    )
+    mock_is_darknet_weights_called.return_value = True
+    mock_download_darknet_weights = mocker.patch(
+        "tf2_yolov4.model.download_darknet_weights"
+    )
+
+    YOLOv4(
+        input_shape=(416, 416, 3),
+        num_classes=80,
+        anchors=YOLOV4_ANCHORS,
+        weights="darknet",
+    )
+    mock_is_darknet_weights_called.assert_called_once_with()
+    assert mock_download_darknet_weights.call_count == 0
+
+
+def test_should_load_weights_from_file_if_path_exists(mocker):
+    mocker.patch("tf2_yolov4.model.csp_darknet53")
+    mocker.patch("tf2_yolov4.model.compute_normalized_anchors")
+    mocker.patch("tf2_yolov4.model.yolov4_neck")
+    mocker.patch("tf2_yolov4.model.yolov3_head")
+    mocker.patch(
+        "tf2_yolov4.model.tf.keras.Model"
+    ).return_value = mocker.sentinel.yolov4
+    mocker.sentinel.yolov4.load_weights = mocker.MagicMock()
+
+    YOLOv4(
+        input_shape=(416, 416, 3),
+        num_classes=80,
+        anchors=YOLOV4_ANCHORS,
+        weights=Path(__file__),
+    )
+    mocker.sentinel.yolov4.load_weights.assert_called_once_with(
+        Path(__file__), by_name=True, skip_mismatch=True,
+    )

tf2_yolov4/backbones/csp_darknet53.py

@@ -112,8 +112,3 @@ def csp_darknet53(input_shape):
     output_3 = csp_block(output_2, filters=1024, num_blocks=4)
 
     return tf.keras.Model(inputs, [output_1, output_2, output_3], name="CSPDarknet53")
-
-
-if __name__ == "__main__":
-    cspdarknet53 = csp_darknet53((416, 416, 3))
-    cspdarknet53.summary()

tf2_yolov4/heads/yolov3_head.py

@@ -5,7 +5,6 @@
 """
 import tensorflow as tf
 
-from tf2_yolov4.anchors import YOLOV3_ANCHORS, compute_normalized_anchors
 from tf2_yolov4.layers import conv_bn
 
 
@@ -249,17 +248,3 @@ def yolo_nms(yolo_feats, yolo_max_boxes, yolo_iou_threshold, yolo_score_threshol
     )
 
     return [boxes, scores, classes, valid_detections]
-
-
-if __name__ == "__main__":
-
-    model = yolov3_head(
-        [(52, 52, 256), (26, 26, 512), (13, 13, 1024)],
-        anchors=compute_normalized_anchors(YOLOV3_ANCHORS, (416, 416, 3)),
-        num_classes=80,
-        training=True,
-        yolo_max_boxes=50,
-        yolo_iou_threshold=0.5,
-        yolo_score_threshold=0.8,
-    )
-    model.summary()

tf2_yolov4/model.py

@@ -5,7 +5,7 @@
 
 import tensorflow as tf
 
-from tf2_yolov4.anchors import YOLOV4_ANCHORS, compute_normalized_anchors
+from tf2_yolov4.anchors import compute_normalized_anchors
 from tf2_yolov4.backbones.csp_darknet53 import csp_darknet53
 from tf2_yolov4.heads.yolov3_head import yolov3_head
 from tf2_yolov4.necks.yolov4_neck import yolov4_neck
@@ -89,24 +89,10 @@ def YOLOv4(
 
     if weights == "darknet":
         if not is_darknet_weights_available():
-            download_darknet_weights(model)
+            download_darknet_weights(yolov4)
 
-        model.load_weights(DARKNET_AS_H5_PATH, by_name=True, skip_mismatch=True)
+        yolov4.load_weights(DARKNET_AS_H5_PATH, by_name=True, skip_mismatch=True)
     elif Path(weights).is_file():
         yolov4.load_weights(weights, by_name=True, skip_mismatch=True)
 
     return yolov4
-
-
-if __name__ == "__main__":
-    model = YOLOv4(
-        input_shape=(608, 416, 3),
-        weights="darknet",
-        num_classes=80,
-        anchors=YOLOV4_ANCHORS,
-    )
-
-    outputs = model.predict(tf.random.uniform((16, 608, 416, 3)), steps=1)
-    model.summary()
-    for output in outputs:
-        print(output.shape)

tf2_yolov4/necks/yolov4_neck.py

@@ -70,8 +70,3 @@ def yolov4_neck(input_shapes):
     return tf.keras.Model(
         [input_1, input_2, input_3], [output_1, output_2, output_3], name="YOLOv4_neck"
     )
-
-
-if __name__ == "__main__":
-    model = yolov4_neck([(52, 52, 256), (26, 26, 512), (13, 13, 1024)])
-    model.summary()