in save rater_mask adjusted type logic and logging

DARPA-CRITICALMAAS · May 7, 2024 · 73dc6ff · 73dc6ff
1 parent 26bef9d
commit 73dc6ff
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/src/pipeline_steps.py b/src/pipeline_steps.py
@@ -227,22 +227,25 @@ def save_output(data_id, map_data: CMAAS_Map, output_dir, feedback_dir, output_t
     if 'raster_masks' in output_types:
         legend_index = 1
         for feature in map_data.legend.features:
-            pipeline_manager.log(logging.DEBUG, f'{map_data.name}, {feature.label} {feature.type}')
+            pipeline_manager.log(logging.DEBUG, f'Saving raster_mask for {map_data.name}, {feature.label} {feature.type}')
+            if feature.type == MapUnitType.LINE:
+                continue
             if feature.type == MapUnitType.POINT:
                 if map_data.point_segmentation_mask is None:
+                    pipeline_manager.log(logging.WARNING, f"Can\'t save feature {feature.label}. No predicted point_segmentation mask present.")
                     continue
                 feature_mask = np.zeros_like(map_data.point_segmentation_mask, dtype=np.uint8)
                 feature_mask[map_data.poly_segmentation_mask == legend_index] = 1
-                filepath = os.path.join(output_dir, sanitize_filename(f'{map_data.name}_{feature.label}.tif'))
-                io.saveGeoTiff(filepath, feature_mask, map_data.georef.crs, map_data.georef.transform)
+                #filepath = os.path.join(output_dir, sanitize_filename(f'{map_data.name}_{feature.label}.tif'))
+                #io.saveGeoTiff(filepath, feature_mask, map_data.georef.crs, map_data.georef.transform)
             if feature.type == MapUnitType.POLYGON:
                 if map_data.poly_segmentation_mask is None:
-                    #pipeline_manager.log(logging.WARNING, f'Can\'t validate feature {feature.label}. No predicted segmentation mask present.')
+                    pipeline_manager.log(logging.WARNING, f"Can\'t save feature {feature.label}. No predicted poly_segmentation mask present.")
                     continue
                 feature_mask = np.zeros_like(map_data.poly_segmentation_mask, dtype=np.uint8)
                 feature_mask[map_data.poly_segmentation_mask == legend_index] = 1
-                filepath = os.path.join(output_dir, sanitize_filename(f'{map_data.name}_{feature.label}.tif'))
-                io.saveGeoTiff(filepath, feature_mask, map_data.georef.crs, map_data.georef.transform)
+            filepath = os.path.join(output_dir, sanitize_filename(f'{map_data.name}_{feature.label}.tif'))
+            io.saveGeoTiff(filepath, feature_mask, map_data.georef.crs, map_data.georef.transform)
             legend_index += 1
     return map_data.name