wip

sicara · Jul 2, 2020 · e8d8fa8 · e8d8fa8
1 parent 3a4f70c
commit e8d8fa8
Show file tree

Hide file tree

Showing 3 changed files with 97 additions and 85 deletions.
diff --git a/scripts/test.py b/scripts/test.py
@@ -1,10 +1,14 @@
 import matplotlib.pyplot as plt
+import numpy as np
 import tensorflow as tf
 
-from tf2_yolov4.anchors import YOLOV4_ANCHORS
+from tf2_yolov4.anchors import YOLOV4_ANCHORS, compute_normalized_anchors
 from tf2_yolov4.model import YOLOv4
 
+
 HEIGHT, WIDTH = (608, 608)
+INPUT_SHAPE = (HEIGHT, WIDTH, 3)
+YOLOV4_ANCHORS_NORMALIZED = compute_normalized_anchors(YOLOV4_ANCHORS, INPUT_SHAPE)
 
 image = tf.io.read_file("../notebooks/images/cars.jpg")
 image = tf.image.decode_image(image)
@@ -17,18 +21,22 @@
     num_classes=20,
     training=False,
     yolo_max_boxes=100,
-    yolo_iou_threshold=0.5,
-    yolo_score_threshold=0.15,
+    yolo_iou_threshold=0.4,
+    yolo_score_threshold=0.1,
 )
-for layer in (
-        model.get_layer("CSPDarknet53").layers + model.get_layer("YOLOv4_neck").layers
-):
-    layer.trainable = False
+# for layer in (
+#         model.get_layer("CSPDarknet53").layers #+ model.get_layer("YOLOv4_neck").layers
+# ):
+#     layer.trainable = False
 
-model.load_weights("../yolov4_all_frozen.h5")
+model.load_weights("../yolov4_full.h5")
 model.summary()
 
+import time
+t0 = time.time()
 boxes, scores, classes, valid_detections = model.predict(images)
+t1 = time.time()
+print(boxes)
 
 CLASSES = [
     "aeroplane",

diff --git a/scripts/train.py b/scripts/train.py
@@ -12,7 +12,7 @@
 from tf2_yolov4.heads.yolov3_head import yolov3_boxes_regression
 from tf2_yolov4.model import YOLOv4
 
-INPUT_SHAPE = (416, 416, 3)
+INPUT_SHAPE = (608, 608, 3)
 BATCH_SIZE = 8
 BOUNDING_BOXES_FIXED_NUMBER = 60
 PASCAL_VOC_NUM_CLASSES = 20
@@ -185,7 +185,7 @@ def transform_targets(y_train, anchors, anchor_masks, size):
         y_outs.append(transform_targets_for_output(y_train, grid_size, anchor_idxs))
         grid_size *= 2
 
-    return tuple(reversed(y_outs))
+    return tuple(y_outs)
 
 
 def pad_bounding_boxes_to_fixed_number_of_bounding_boxes(bounding_boxes, pad_number):
@@ -298,7 +298,7 @@ def prepare_dataset(
 
 
 if __name__ == "__main__":
-    voc_dataset, infos = tfds.load("voc/2012", with_info=True, shuffle_files=True)
+    voc_dataset, infos = tfds.load("voc", with_info=True, shuffle_files=True)
     ds_train, ds_test = voc_dataset["train"], voc_dataset["validation"]
     ds_train = prepare_dataset(
         ds_train,
@@ -322,11 +322,14 @@ def prepare_dataset(
         num_classes=PASCAL_VOC_NUM_CLASSES,
         training=True,
     )
-    darknet_weights = Path("./yolov4.h5")
-    if darknet_weights.exists():
-        model.load_weights(str(darknet_weights), by_name=True, skip_mismatch=True)
-        print("Darknet weights loaded.")
 
+    print([output.shape for output in model.outputs])
+
+    # darknet_weights = Path("./yolov4.h5")
+    # if darknet_weights.exists():
+    #     model.load_weights(str(darknet_weights), by_name=True, skip_mismatch=True)
+    #     print("Darknet weights loaded.")
+    #
     optimizer = tf.keras.optimizers.Adam(1e-4)
     loss = [
         YoloLoss(
@@ -335,72 +338,73 @@ def prepare_dataset(
         for mask in YOLOV4_ANCHORS_MASKS
     ]
 
-    model.summary()
-    # Start training: 5 epochs with backbone + neck frozen
-    for layer in (
-        model.get_layer("CSPDarknet53").layers + model.get_layer("YOLOv4_neck").layers
-    ):
-        layer.trainable = False
-    model.compile(optimizer=optimizer, loss=loss)
-    history = model.fit(
-        ds_train,
-        steps_per_epoch=steps_per_epoch,
-        validation_data=ds_test,
-        validation_steps=validation_steps,
-        epochs=ALL_FROZEN_EPOCH_NUMBER,
-        callbacks=[
-            tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR),
-            tf.keras.callbacks.ModelCheckpoint(
-                str(LOG_DIR / "yolov4_all_frozen.h5"),
-                save_best_only=True,
-                save_weights_only=True,
-            ),
-        ],
-    )
-    # Keep training: 10 epochs with backbone frozen -- unfreeze neck
-    for layer in model.get_layer("YOLOv4_neck").layers:
-        layer.trainable = True
-    model.compile(optimizer=optimizer, loss=loss)
-    history = model.fit(
-        ds_train,
-        steps_per_epoch=steps_per_epoch,
-        validation_data=ds_test,
-        validation_steps=validation_steps,
-        epochs=BACKBONE_FROZEN_EPOCH_NUMBER + ALL_FROZEN_EPOCH_NUMBER,
-        initial_epoch=ALL_FROZEN_EPOCH_NUMBER,
-        callbacks=[
-            tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR),
-            tf.keras.callbacks.ModelCheckpoint(
-                str(LOG_DIR / "yolov4_backbone_frozen.h5"),
-                save_best_only=True,
-                save_weights_only=True,
-                verbose=True,
-            ),
-        ],
-    )
-    # Final training: 35 epochs with all weights unfrozen
-    for layer in model.get_layer("CSPDarknet53").layers:
-        layer.trainable = True
-    model.compile(optimizer=optimizer, loss=loss)
-    history = model.fit(
-        ds_train,
-        steps_per_epoch=steps_per_epoch,
-        validation_data=ds_test,
-        validation_steps=validation_steps,
-        epochs=TOTAL_NUMBER_OF_EPOCHS,
-        initial_epoch=ALL_FROZEN_EPOCH_NUMBER + BACKBONE_FROZEN_EPOCH_NUMBER,
-        callbacks=[
-            tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR),
-            tf.keras.callbacks.ModelCheckpoint(
-                str(LOG_DIR / "yolov4_full.h5"),
-                save_best_only=True,
-                save_weights_only=True,
-            ),
-            tf.keras.callbacks.ModelCheckpoint(
-                str(LOG_DIR / "yolov4_train_loss.h5"),
-                save_best_only=True,
-                save_weights_only=True,
-                monitor="loss",
-            ),
-        ],
-    )
+    print([YOLOV4_ANCHORS_NORMALIZED[mask] for mask in YOLOV4_ANCHORS_MASKS])
+    #
+    # # Start training: 5 epochs with backbone + neck frozen
+    # for layer in (
+    #     model.get_layer("CSPDarknet53").layers + model.get_layer("YOLOv4_neck").layers
+    # ):
+    #     layer.trainable = False
+    # model.compile(optimizer=optimizer, loss=loss)
+    # history = model.fit(
+    #     ds_train,
+    #     steps_per_epoch=steps_per_epoch,
+    #     validation_data=ds_test,
+    #     validation_steps=validation_steps,
+    #     epochs=ALL_FROZEN_EPOCH_NUMBER,
+    #     callbacks=[
+    #         tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR),
+    #         tf.keras.callbacks.ModelCheckpoint(
+    #             str(LOG_DIR / "yolov4_all_frozen.h5"),
+    #             save_best_only=True,
+    #             save_weights_only=True,
+    #         ),
+    #     ],
+    # )
+    # # Keep training: 10 epochs with backbone frozen -- unfreeze neck
+    # for layer in model.get_layer("YOLOv4_neck").layers:
+    #     layer.trainable = True
+    # model.compile(optimizer=optimizer, loss=loss)
+    # history = model.fit(
+    #     ds_train,
+    #     steps_per_epoch=steps_per_epoch,
+    #     validation_data=ds_test,
+    #     validation_steps=validation_steps,
+    #     epochs=BACKBONE_FROZEN_EPOCH_NUMBER + ALL_FROZEN_EPOCH_NUMBER,
+    #     initial_epoch=ALL_FROZEN_EPOCH_NUMBER,
+    #     callbacks=[
+    #         tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR),
+    #         tf.keras.callbacks.ModelCheckpoint(
+    #             str(LOG_DIR / "yolov4_backbone_frozen.h5"),
+    #             save_best_only=True,
+    #             save_weights_only=True,
+    #             verbose=True,
+    #         ),
+    #     ],
+    # )
+    # # Final training: 35 epochs with all weights unfrozen
+    # for layer in model.get_layer("CSPDarknet53").layers:
+    #     layer.trainable = True
+    # model.compile(optimizer=optimizer, loss=loss)
+    # history = model.fit(
+    #     ds_train,
+    #     steps_per_epoch=steps_per_epoch,
+    #     validation_data=ds_test,
+    #     validation_steps=validation_steps,
+    #     epochs=TOTAL_NUMBER_OF_EPOCHS,
+    #     initial_epoch=ALL_FROZEN_EPOCH_NUMBER + BACKBONE_FROZEN_EPOCH_NUMBER,
+    #     callbacks=[
+    #         tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR),
+    #         tf.keras.callbacks.ModelCheckpoint(
+    #             str(LOG_DIR / "yolov4_full.h5"),
+    #             save_best_only=True,
+    #             save_weights_only=True,
+    #         ),
+    #         tf.keras.callbacks.ModelCheckpoint(
+    #             str(LOG_DIR / "yolov4_train_loss.h5"),
+    #             save_best_only=True,
+    #             save_weights_only=True,
+    #             monitor="loss",
+    #         ),
+    #     ],
+    # )
diff --git a/tf2_yolov4/heads/yolov3_head.py b/tf2_yolov4/heads/yolov3_head.py
@@ -95,7 +95,7 @@ def yolov3_head(
     if training:
         return tf.keras.Model(
             [input_1, input_2, input_3],
-            [output_1, output_2, output_3],
+            [output_3, output_2, output_1],
             name="YOLOv3_head",
         )