Implemented SSD train action

Signed-off-by: Jim Martens <github@2martens.de>

src/twomartens/masterthesis/cli.py

@@ -28,10 +28,58 @@ import argparse
 
 
 def train(args: argparse.Namespace) -> None:
-    if args.network == "auto_encoder":
+    if args.network == "ssd" or args.network == "bayesian_ssd":
+        _ssd_train(args)
+    elif args.network == "auto_encoder":
         _auto_encoder_train(args)
 
 
+def _ssd_train(args: argparse.Namespace) -> None:
+    import os
+    import pickle
+    
+    import tensorflow as tf
+    from tensorflow.python.ops import summary_ops_v2
+
+    from twomartens.masterthesis import data
+    from twomartens.masterthesis import ssd
+    
+    tf.enable_eager_execution()
+    
+    batch_size = 16
+    image_size = 300
+    use_dropout = False if args.network == "ssd" else True
+    
+    pre_trained_weights_file = f"{args.weights_path}/VGG_coco_SSD_300x300_iter_400000.h5"
+    weights_path = f"{args.weights_path}/train/{args.network}/"
+    os.makedirs(weights_path, exist_ok=True)
+    
+    # load prepared ground truth
+    with open(f"{args.ground_truth_path}/photo_paths.bin", "rb") as file:
+        file_names_photos = pickle.load(file)
+    with open(f"{args.ground_truth_path}/instances.bin", "rb") as file:
+        instances = pickle.load(file)
+    
+    scenenet_data, nr_digits = data.load_scenenet_data(file_names_photos, instances, args.coco_path,
+                                                       batch_size=batch_size,
+                                                       resized_shape=(image_size, image_size))
+    del file_names_photos, instances
+
+    use_summary_writer = summary_ops_v2.create_file_writer(
+        f"{args.summary_path}/val/{args.network}/{args.iteration}"
+    )
+    
+    if args.debug:
+        with use_summary_writer.as_default():
+            ssd.train(scenenet_data, args.iteration, use_dropout, weights_prefix=weights_path,
+                      weights_path=pre_trained_weights_file, batch_size=batch_size,
+                      nr_epochs=args.num_epochs)
+    else:
+        ssd.train(scenenet_data, args.iteration, use_dropout, weights_prefix=weights_path,
+                  weights_path=pre_trained_weights_file, batch_size=batch_size,
+                  nr_epochs=args.num_epochs)
+
+
 def _auto_encoder_train(args: argparse.Namespace) -> None:
     from twomartens.masterthesis import data
     from twomartens.masterthesis.aae import train
@@ -197,7 +245,7 @@ def _ssd_val(args: argparse.Namespace) -> None:
     with open(f"{args.ground_truth_path}/instances.bin", "rb") as file:
         instances = pickle.load(file)
     
-    scenenet_data, nr_digits = data.load_scenenet_val(file_names_photos, instances, args.coco_path,
+    scenenet_data, nr_digits = data.load_scenenet_data(file_names_photos, instances, args.coco_path,
                                                        batch_size=batch_size,
                                                        resized_shape=(image_size, image_size))
     del file_names_photos, instances

src/twomartens/masterthesis/data.py

@@ -227,7 +227,7 @@ def _load_images_callback(resized_shape: Sequence[int]) -> Callable[
     return _load_images
 
 
-def load_scenenet_val(photo_paths: Sequence[Sequence[str]],
+def load_scenenet_data(photo_paths: Sequence[Sequence[str]],
                        instances: Sequence[Sequence[Sequence[dict]]],
                        coco_path: str,
                        num_epochs: int = 1, batch_size: int = 32,
@@ -244,7 +244,7 @@ def load_scenenet_val(photo_paths: Sequence[Sequence[str]],
         resized_shape: shape of input images to SSD
 
     Returns:
-        scenenet val data set
+        scenenet data set
         number of digits required to print largest batch number
     """
     trajectories = zip(photo_paths, instances)

src/twomartens/masterthesis/main.py

@@ -72,14 +72,25 @@ def _build_train(parser: argparse.ArgumentParser) -> None:
     sub_parsers = parser.add_subparsers(dest="network")
     sub_parsers.required = True
     
+    ssd_parser = sub_parsers.add_parser("ssd", help="SSD")
     # ssd_bayesian_parser = sub_parsers.add_parser("bayesian_ssd", help="SSD with dropout layers")
     auto_encoder_parser = sub_parsers.add_parser("auto_encoder", help="Auto-encoder network")
     
     # build sub parsers
+    _build_ssd_train(ssd_parser)
     # _build_bayesian_ssd(ssd_bayesian_parser)
     _build_auto_encoder_train(auto_encoder_parser)
  
 
+def _build_ssd_train(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--coco_path", type=str, help="the path to the COCO data set")
+    parser.add_argument("--weights_path", type=str, help="path to the weights directory")
+    parser.add_argument("--ground_truth_path", type=str, help="path to the prepared ground truth directory")
+    parser.add_argument("--summary_path", type=str, help="path to the summaries directory")
+    parser.add_argument("num_epochs", type=int, help="the number of epochs to train", default=80)
+    parser.add_argument("iteration", type=int, help="the training iteration")
+    
+
 def _build_auto_encoder_train(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("--coco_path", type=str, help="the path to the COCO data set")
     parser.add_argument("--weights_path", type=str, help="path to the weights directory")