Fixes single instance model visualization during training (talmolab#620)

Fixes talmolab#604
ShreyashKad · Dec 29, 2021 · c75e143 · c75e143
1 parent 507b71a
commit c75e143
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/sleap/nn/data/pipelines.py b/sleap/nn/data/pipelines.py
@@ -454,9 +454,14 @@ def make_viz_pipeline(
             A `Pipeline` instance configured to fetch data and run inference to generate
             predictions useful for visualization during training.
         """
-        pipeline = self.make_base_pipeline(data_provider=data_provider)
-        pipeline += Prefetcher()
+        pipeline = Pipeline(providers=data_provider)
+        if self.data_config.preprocessing.resize_and_pad_to_target:
+            pipeline += SizeMatcher.from_config(
+                config=self.data_config.preprocessing,
+                provider=data_provider,
+            )
         pipeline += Repeater()
+        pipeline += Prefetcher()
         return pipeline
 
 

diff --git a/sleap/nn/training.py b/sleap/nn/training.py
@@ -1033,6 +1033,7 @@ def _setup_visualization(self):
         inference_layer = SingleInstanceInferenceLayer(
             keras_model=self.keras_model,
             input_scale=self.config.data.preprocessing.input_scaling,
+            pad_to_stride=self.config.data.preprocessing.pad_to_stride,
             peak_threshold=0.2,
             return_confmaps=True,
         )