pipelines

talmolab · Feb 2, 2021 · 9408f9d · 9408f9d
1 parent acaa740
commit 9408f9d
Showing 1 changed file with 42 additions and 1 deletion.
diff --git a/sleap/nn/data/pipelines.py b/sleap/nn/data/pipelines.py
@@ -22,7 +22,7 @@
     RandomCropper,
 )
 from sleap.nn.data.normalization import Normalizer
-from sleap.nn.data.resizing import Resizer, PointsRescaler
+from sleap.nn.data.resizing import Resizer, PointsRescaler, SizeMatcher
 from sleap.nn.data.instance_centroids import InstanceCentroidFinder
 from sleap.nn.data.instance_cropping import InstanceCropper, PredictedInstanceCropper
 from sleap.nn.data.confidence_maps import (
@@ -68,6 +68,7 @@
     RandomCropper,
     Normalizer,
     Resizer,
+    SizeMatcher,
     InstanceCentroidFinder,
     InstanceCropper,
     MultiConfidenceMapGenerator,
@@ -353,6 +354,11 @@ def make_base_pipeline(self, data_provider: Provider) -> Pipeline:
         pipeline = Pipeline(providers=data_provider)
         pipeline += Normalizer.from_config(self.data_config.preprocessing)
         pipeline += Resizer.from_config(self.data_config.preprocessing)
+        if self.data_config.preprocessing.resize_and_pad_to_target:
+            pipeline += SizeMatcher.from_config(
+                config=self.data_config.preprocessing,
+                provider=data_provider,
+            )
         if self.optimization_config.augmentation_config.random_crop:
             pipeline += RandomCropper(
                 crop_height=self.optimization_config.augmentation_config.random_crop_height,
@@ -393,6 +399,11 @@ def make_training_pipeline(self, data_provider: Provider) -> Pipeline:
             )
         pipeline += Normalizer.from_config(self.data_config.preprocessing)
         pipeline += Resizer.from_config(self.data_config.preprocessing)
+        if self.data_config.preprocessing.resize_and_pad_to_target:
+            pipeline += SizeMatcher.from_config(
+                config=self.data_config.preprocessing,
+                provider=data_provider,
+            )
 
         pipeline += SingleInstanceConfidenceMapGenerator(
             sigma=self.single_instance_confmap_head.sigma,
@@ -483,6 +494,11 @@ def make_base_pipeline(self, data_provider: Provider) -> Pipeline:
         pipeline = Pipeline(providers=data_provider)
         pipeline += Normalizer.from_config(self.data_config.preprocessing)
         pipeline += Resizer.from_config(self.data_config.preprocessing)
+        if self.data_config.preprocessing.resize_and_pad_to_target:
+            pipeline += SizeMatcher.from_config(
+                config=self.data_config.preprocessing,
+                provider=data_provider,
+            )
         if self.optimization_config.augmentation_config.random_crop:
             pipeline += RandomCropper(
                 crop_height=self.optimization_config.augmentation_config.random_crop_height,
@@ -529,6 +545,11 @@ def make_training_pipeline(self, data_provider: Provider) -> Pipeline:
             )
         pipeline += Normalizer.from_config(self.data_config.preprocessing)
         pipeline += Resizer.from_config(self.data_config.preprocessing)
+        if self.data_config.preprocessing.resize_and_pad_to_target:
+            pipeline += SizeMatcher.from_config(
+                config=self.data_config.preprocessing,
+                provider=data_provider,
+            )
 
         pipeline += InstanceCentroidFinder.from_config(
             self.data_config.instance_cropping,
@@ -635,6 +656,11 @@ def make_base_pipeline(self, data_provider: Provider) -> Pipeline:
         pipeline = Pipeline(providers=data_provider)
         pipeline += Normalizer.from_config(self.data_config.preprocessing)
         pipeline += Resizer.from_config(self.data_config.preprocessing)
+        if self.data_config.preprocessing.resize_and_pad_to_target:
+            pipeline += SizeMatcher.from_config(
+                config=self.data_config.preprocessing,
+                provider=data_provider,
+            )
         pipeline += InstanceCentroidFinder.from_config(
             self.data_config.instance_cropping,
             skeletons=self.data_config.labels.skeletons,
@@ -672,6 +698,11 @@ def make_training_pipeline(self, data_provider: Provider) -> Pipeline:
         )
         pipeline += Normalizer.from_config(self.data_config.preprocessing)
         pipeline += Resizer.from_config(self.data_config.preprocessing)
+        if self.data_config.preprocessing.resize_and_pad_to_target:
+            pipeline += SizeMatcher.from_config(
+                config=self.data_config.preprocessing,
+                provider=data_provider,
+            )
 
         pipeline += InstanceCentroidFinder.from_config(
             self.data_config.instance_cropping,
@@ -764,6 +795,11 @@ def make_base_pipeline(self, data_provider: Provider) -> Pipeline:
         pipeline = Pipeline(providers=data_provider)
         pipeline += Normalizer.from_config(self.data_config.preprocessing)
         pipeline += Resizer.from_config(self.data_config.preprocessing)
+        if self.data_config.preprocessing.resize_and_pad_to_target:
+            pipeline += SizeMatcher.from_config(
+                config=self.data_config.preprocessing,
+                provider=data_provider,
+            )
         if self.optimization_config.augmentation_config.random_crop:
             pipeline += RandomCropper(
                 crop_height=self.optimization_config.augmentation_config.random_crop_height,
@@ -805,6 +841,11 @@ def make_training_pipeline(self, data_provider: Provider) -> Pipeline:
             )
         pipeline += Normalizer.from_config(self.data_config.preprocessing)
         pipeline += Resizer.from_config(self.data_config.preprocessing)
+        if self.data_config.preprocessing.resize_and_pad_to_target:
+            pipeline += SizeMatcher.from_config(
+                config=self.data_config.preprocessing,
+                provider=data_provider,
+            )
 
         pipeline += MultiConfidenceMapGenerator(
             sigma=self.confmaps_head.sigma,