tests fix

phborba · Sep 21, 2020 · 69d38ff · 69d38ff
1 parent 8ce8889
commit 69d38ff
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/segmentation_models_trainer/experiment_builder/experiment.py b/segmentation_models_trainer/experiment_builder/experiment.py
@@ -179,7 +179,7 @@ def test_and_create_folder(path):
 
             'model' : json.loads('{"description": "test case", "backbone": "resnet18", "architecture": "Unet", "activation": "sigmoid", "use_imagenet_weights": true}'),
             'loss' : json.loads('{"class_name": "bce_dice_loss", "config": {}, "framework": "sm"}'),
-            'callbacks' : json.loads('{"items": [{"name": "ReduceLROnPlateau", "config": {"monitor": "val_loss", "factor": 0.2, "patience": 5, "min_lr": 0.001}}, {"name": "ModelCheckpoint", "config": {"filepath": "/data/teste"}}]}'),
+            'callbacks' : json.loads('{"items": [{"name": "ReduceLROnPlateau", "config": {"monitor": "val_loss", "factor": 0.2, "patience": 5, "min_lr": 0.001}}, {"name": "ModelCheckpoint", "config": {"filepath": "/data/teste/checkpoint.hdf5"}}]}'),
             'metrics' : json.loads('{"items": [{"class_name": "iou_score", "config": {}, "framework": "sm"}, {"class_name": "precision", "config": {}, "framework": "sm"}, {"class_name": "recall", "config": {}, "framework": "sm"}, {"class_name": "f1_score", "config": {}, "framework": "sm"}, {"class_name": "f2_score", "config": {}, "framework": "sm"}, {"class_name": "LogCoshError", "config": {}, "framework": "tf.keras"}, {"class_name": "KLDivergence", "config": {}, "framework": "tf.keras"}, {"class_name": "MeanIoU", "config": {"num_classes": 2}, "framework": "tf.keras"}]}'),
         }
     )

diff --git a/segmentation_models_trainer/experiment_builder/metric.py b/segmentation_models_trainer/experiment_builder/metric.py
@@ -31,14 +31,18 @@ class Metric(JsonSchemaMixin):
     framework: str
 
     def __post_init__(self):
+        if self.framework not in ['sm', 'tf.keras']:
+            raise ValueError("Metric not implemented")
+
+    def get_metric(self):
         if self.framework == 'sm':
-            self.metric_obj = self.get_sm_metric(self.class_name)
+            return self.get_sm_metric(self.class_name)
         elif self.framework == 'tf.keras':
             identifier = {
                 "class_name" : self.class_name,
                 "config" : self.config
             }
-            self.metric_obj = tf.keras.metrics.get(identifier)
+            return tf.keras.metrics.get(identifier)
         else:
             raise ValueError("Metric not implemented")
 
@@ -62,7 +66,7 @@ class MetricList(JsonSchemaMixin):
 
     def get_tf_objects(self):
         return [
-            i.metric_obj for i in self.items
+            i.get_metric() for i in self.items
         ]
 
 if __name__ == "__main__":

diff --git a/tests/test_training_script.py b/tests/test_training_script.py
@@ -98,7 +98,10 @@ def setUp(self):
                 '''{"name": "test_ds", "file_path": "'''+self.csv_test_ds_file+'''", "n_classes": 1, "dataset_size": 1, "augmentation_list": [{"name": "random_crop", "parameters": {"crop_width": 256, "crop_height": 256}}, {"name": "per_image_standardization", "parameters": {}}], "cache": false, "shuffle": false, "shuffle_buffer_size": 1, "shuffle_csv": true, "ignore_errors": true, "num_paralel_reads": 1, "img_dtype": "float32", "img_format": "png", "img_width": 256, "img_length": 256, "img_bands": 3, "mask_bands": 1, "use_ds_width_len": false, "autotune": -1, "distributed_training": false}'''
             ),
 
-            'model' : json.loads('{"description": "test case", "backbone": "resnet18", "architecture": "Unet", "activation": "sigmoid", "use_imagenet_weights": true}')
+            'model' : json.loads('{"description": "test case", "backbone": "resnet18", "architecture": "Unet", "activation": "sigmoid", "use_imagenet_weights": true}'),
+            'loss' : json.loads('{"class_name": "bce_dice_loss", "config": {}, "framework": "sm"}'),
+            'callbacks' : json.loads('{"items": [{"name": "ReduceLROnPlateau", "config": {"monitor": "val_loss", "factor": 0.2, "patience": 5, "min_lr": 0.001}}, {"name": "ModelCheckpoint", "config": {"filepath": "/data/teste/checkpoint.hdf5"}}]}'),
+            'metrics' : json.loads('{"items": [{"class_name": "iou_score", "config": {}, "framework": "sm"}, {"class_name": "precision", "config": {}, "framework": "sm"}, {"class_name": "recall", "config": {}, "framework": "sm"}, {"class_name": "f1_score", "config": {}, "framework": "sm"}, {"class_name": "f2_score", "config": {}, "framework": "sm"}, {"class_name": "LogCoshError", "config": {}, "framework": "tf.keras"}, {"class_name": "KLDivergence", "config": {}, "framework": "tf.keras"}, {"class_name": "MeanIoU", "config": {"num_classes": 2}, "framework": "tf.keras"}]}'),
         }
         self.settings_json = os.path.join(
             current_dir, 'testing_data', 'settings.json'