Refactor data manager to use same API

oarriaga · oarriaga · Jun 25, 2021 · Jun 18, 2021 · Jun 18, 2021 · Jun 18, 2021
commit 165f1b1debe27933bbcea6b02ad2ad6261c52e2f
diff --git a/examples/fine-tuning_object_detection/data_manager.py b/examples/fine-tuning_object_detection/data_manager.py
@@ -17,7 +17,7 @@ def build_class_to_arg(self, class_names):
         args = list(range(len(class_names)))
         return dict(zip(args, class_names))
 
-    def load(self):
+    def load_data(self):
         file = open(self.path, 'r')
         csv_reader = csv.reader(file, delimiter=',')
         dataset = []
@@ -41,4 +41,4 @@ def load(self):
     path = 'datasets/solar_panel/BoundingBox.txt'
     class_names = ['background', 'solar_panel']
     data_manager = CSVLoader(path, class_names)
-    dataset = data_manager.load()
+    dataset = data_manager.load_data()
diff --git a/examples/fine-tuning_object_detection/pipelines.py b/examples/fine-tuning_object_detection/pipelines.py
@@ -212,7 +212,7 @@ def __init__(self, prior_boxes, bkg_paths, split=pr.TRAIN, num_classes=2,
     path = 'datasets/solar_panel/BoundingBox.txt'
     class_names = ['background', 'solar_panel']
     data_manager = CSVLoader(path, class_names)
-    dataset = data_manager.load()
+    dataset = data_manager.load_data()
     home = os.path.expanduser('~')
     bkg_path = os.path.join(home, '.keras/paz/datasets/voc-backgrounds/')
     wild_card = bkg_path + '*.png'

diff --git a/examples/fine-tuning_object_detection/train.py b/examples/fine-tuning_object_detection/train.py
@@ -1,4 +1,5 @@
 import os
+import glob
 import argparse
 import tensorflow as tf
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -11,9 +12,9 @@
 from tensorflow.keras.optimizers import SGD
 from tensorflow.keras.callbacks import CSVLogger, ModelCheckpoint
 from paz.optimization.callbacks import LearningRateScheduler
-from detection import AugmentDetection
+from pipelines import AugmentDetection
 from paz.models import SSD300
-from paz.datasets import VOC
+from data_manager import CSVLoader
 from paz.optimization import MultiBoxLoss
 from paz.abstract import ProcessingSequence
 from paz.optimization.callbacks import EvaluateMAP
@@ -40,6 +41,11 @@
                     type=str, help='Path for writing model weights and logs')
 parser.add_argument('-dp', '--data_path', default='VOCdevkit/',
                     type=str, help='Path for writing model weights and logs')
+parser.add_argument('-bk', '--bkg_path',
+                    default=os.path.join(
+                        os.path.expanduser('~'),
+                        '.keras/paz/datasets/voc-backgrounds/'),
+                    type=str, help='Path for writing model weights and logs')
 parser.add_argument('-se', '--scheduled_epochs', nargs='+', type=int,
                     default=[110, 152], help='Epoch learning rate reduction')
 parser.add_argument('-mp', '--multiprocessing', default=False, type=bool,
@@ -50,23 +56,20 @@
 
 optimizer = SGD(args.learning_rate, args.momentum)
 
-data_splits = [['trainval', 'trainval'], 'test']
-data_names = [['VOC2007', 'VOC2012'], 'VOC2007']
-
+data_splits = ['train']
 # loading datasets
-data_managers, datasets, evaluation_data_managers = [], [], []
-for data_name, data_split in zip(data_names, data_splits):
-    data_manager = VOC(args.data_path, data_split, name=data_name)
+data_managers, datasets = [], []
+for data_split in zip(data_splits):
+    path = 'datasets/solar_panel/BoundingBox.txt'
+    class_names = ['background', 'solar_panel']
+    data_manager = CSVLoader(path, class_names, split=data_split)
     data_managers.append(data_manager)
     datasets.append(data_manager.load_data())
-    if data_split == 'test':
-        eval_data_manager = VOC(
-            args.data_path, data_split, name=data_name, evaluate=True)
-        evaluation_data_managers.append(eval_data_manager)
 
 # instantiating model
 num_classes = data_managers[0].num_classes
-model = SSD300(num_classes, base_weights='VGG', head_weights=None)
+model = SSD300(num_classes, base_weights='VGG', head_weights=None,
+               trainable_base=False)
 model.summary()
 
 # Instantiating loss and metrics
@@ -77,9 +80,12 @@
 model.compile(optimizer, loss.compute_loss, metrics)
 
 # setting data augmentation pipeline
+bkg_paths = glob.glob(args.bkg_path + '*.png')
+if len(bkg_paths) == 0:
+    raise ValueError('No background png files were found in', args.bkg_path)
 augmentators = []
 for split in [TRAIN, VAL]:
-    augmentator = AugmentDetection(model.prior_boxes, split)
+    augmentator = AugmentDetection(model.prior_boxes, bkg_paths, split)
     augmentators.append(augmentator)
 
 # setting sequencers
@@ -97,19 +103,14 @@
 checkpoint = ModelCheckpoint(save_path, verbose=1, save_weights_only=True)
 schedule = LearningRateScheduler(
     args.learning_rate, args.gamma_decay, args.scheduled_epochs)
-evaluate = EvaluateMAP(
-    evaluation_data_managers[0],
-    DetectSingleShot(model, data_managers[0].class_names, 0.01, 0.45),
-    args.evaluation_period,
-    args.save_path,
-    args.AP_IOU)
 
 # training
-model.fit(
+model.fit_generator(
     sequencers[0],
     epochs=args.num_epochs,
+    steps_per_epoch=100,
     verbose=1,
-    callbacks=[checkpoint, log, schedule, evaluate],
-    validation_data=sequencers[1],
+    callbacks=[checkpoint, log, schedule],
+    # validation_data=sequencers[1],
     use_multiprocessing=args.multiprocessing,
     workers=args.workers)