diff --git a/src/twomartens/masterthesis/cli.py b/src/twomartens/masterthesis/cli.py index 6af4813..4858f51 100644 --- a/src/twomartens/masterthesis/cli.py +++ b/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,9 +245,9 @@ 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, - batch_size=batch_size, - resized_shape=(image_size, image_size)) + 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( diff --git a/src/twomartens/masterthesis/data.py b/src/twomartens/masterthesis/data.py index 03937ec..b0068e2 100644 --- a/src/twomartens/masterthesis/data.py +++ b/src/twomartens/masterthesis/data.py @@ -227,11 +227,11 @@ def _load_images_callback(resized_shape: Sequence[int]) -> Callable[ return _load_images -def load_scenenet_val(photo_paths: Sequence[Sequence[str]], - instances: Sequence[Sequence[Sequence[dict]]], - coco_path: str, - num_epochs: int = 1, batch_size: int = 32, - resized_shape: Sequence[int] = (256, 256)) -> Tuple[tf.data.Dataset, int]: +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, + resized_shape: Sequence[int] = (256, 256)) -> Tuple[tf.data.Dataset, int]: """ Loads the SceneNet RGB-D data and returns a data set. @@ -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) diff --git a/src/twomartens/masterthesis/main.py b/src/twomartens/masterthesis/main.py index 476115a..5afb272 100644 --- a/src/twomartens/masterthesis/main.py +++ b/src/twomartens/masterthesis/main.py @@ -72,12 +72,23 @@ 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: