Merge pull request #157 from PUTvision/sigmoid_output_fix_and_tests

Fixes after testing with single sigmoid

src/deepness/processing/map_processor/map_processor_segmentation.py

@@ -47,10 +47,10 @@ def _run(self) -> MapProcessingResult:
                     full_result_img=full_result_img)
 
         blur_size = int(self.segmentation_parameters.postprocessing_dilate_erode_size // 2) * 2 + 1  # needs to be odd
-        
+
         for i in range(full_result_img.shape[0]):
             full_result_img[i] = cv2.medianBlur(full_result_img[i], blur_size)
-        
+
         full_result_img = self.limit_extended_extent_image_to_base_extent_with_mask(full_img=full_result_img)
 
         self.set_results_img(full_result_img)
@@ -62,48 +62,42 @@ def _run(self) -> MapProcessingResult:
             message=result_message,
             gui_delegate=gui_delegate,
         )
-        
+
     def _check_output_layer_is_sigmoid_and_has_more_than_one_name(self, output_id: int) -> bool:
         if self.model.outputs_names is None or self.model.outputs_are_sigmoid is None:
             return False
-        
+
         return len(self.model.outputs_names[output_id]) > 1 and self.model.outputs_are_sigmoid[output_id]
 
     def _create_result_message(self, result_img: np.ndarray) -> str:
-        
+
         txt = f'Segmentation done, with the following statistics:\n'
-        
+
         for output_id, layer_sizes in enumerate(self._get_indexes_of_model_output_channels_to_create()):
-            
+
             txt += f'Channels for output {output_id}:\n'
-            
+
             unique, counts = np.unique(result_img[output_id], return_counts=True)
             counts_map = {}
             for i in range(len(unique)):
                 counts_map[unique[i]] = counts[i]
-                
+
             # # we cannot simply take image dimensions, because we may have irregular processing area from polygon
             number_of_pixels_in_processing_area = np.sum([counts_map[k] for k in counts_map.keys()])
             total_area = number_of_pixels_in_processing_area * self.params.resolution_m_per_px**2
-            
+
             for channel_id in range(layer_sizes):
                 # See note in the class description why are we adding/subtracting 1 here
                 pixels_count = counts_map.get(channel_id, 0)
                 area = pixels_count * self.params.resolution_m_per_px**2
-                
+
                 if total_area > 0 and not np.isnan(total_area) and not np.isinf(total_area):
                     area_percentage = area / total_area * 100
                 else:
                     area_percentage = 0.0
                     # TODO
-
-                # hardcode if someone add "background class" for sigmoid output, we need to skip it
-                if self._check_output_layer_is_sigmoid_and_has_more_than_one_name(output_id):
-                    channel_id_name = channel_id
-                else:
-                    channel_id_name = channel_id - 1
-
-                txt += f'\t- {self.model.get_channel_name(output_id, channel_id_name)}: area = {area:.2f} m^2 ({area_percentage:.2f} %)\n'
+
+                txt += f'\t- {self.model.get_channel_name(output_id, channel_id)}: area = {area:.2f} m^2 ({area_percentage:.2f} %)\n'
 
         return txt
 
@@ -136,7 +130,7 @@ def _create_vlayer_from_mask_for_base_extent(self, mask_img) -> Callable:
                         current_contour_index=0)
                 else:
                     pass  # just nothing, we already have an empty list of features
-                
+
                 layer_name = self.model.get_channel_name(output_id, channel_id)
                 vlayer = QgsVectorLayer("multipolygon", layer_name, "memory")
                 vlayer.setCrs(self.rlayer.crs())
@@ -152,13 +146,13 @@ def _create_vlayer_from_mask_for_base_extent(self, mask_img) -> Callable:
                 vlayer.updateExtents()
 
                 output_vlayers.append(vlayer)
-                
+
             vlayers.append(output_vlayers)
 
         # accessing GUI from non-GUI thread is not safe, so we need to delegate it to the GUI thread
         def add_to_gui():
             group = QgsProject.instance().layerTreeRoot().insertGroup(0, 'model_output')
-            
+
             if len(vlayers) == 1:
                 for vlayer in vlayers[0]:
                     QgsProject.instance().addMapLayer(vlayer, False)
@@ -194,5 +188,5 @@ def _process_tile(self, tile_img_batched: np.ndarray) -> np.ndarray:
             many_outputs.append(result[:, 0])
 
         many_outputs = np.array(many_outputs).transpose((1, 0, 2, 3))
-        
+
         return many_outputs

src/deepness/processing/models/segmentor.py

@@ -22,13 +22,13 @@ def __init__(self, model_file_path: str):
             Path to the model file
         """
         super(Segmentor, self).__init__(model_file_path)
-        
+
         self.outputs_are_sigmoid = self.check_loaded_model_outputs()
-        
+
         for idx in range(len(self.outputs_layers)):
             if self.outputs_names is None:
                 continue
-            
+
             if len(self.outputs_names[idx]) == 1 and self.outputs_are_sigmoid[idx]:
                 self.outputs_names[idx] = ['background', self.outputs_names[idx][0]]
 
@@ -82,41 +82,23 @@ def get_class_display_name(cls):
         """
         return cls.__name__
 
-    def check_loaded_model_outputs(self):
-        """ Checks if the model outputs are valid
-
-        Valid means that:
-        - the model has at least one output
-        - the output is 4D (N,C,H,W) or 3D (N,H,W)
-        - the batch size is 1 or dynamic
-        - model resolution is equal to TILE_SIZE (is square)
-
-        """
-        if len(self.outputs_layers) == 0:
-            raise Exception('Model has no output layers')
-
-        for layer in self.outputs_layers:
-            if len(layer.shape) != 4 and len(layer.shape) != 3:
-                raise Exception(f'Segmentation model output should have 4 dimensions: (B,C,H,W) or 3 dimensions: (B,H,W). Has {layer.shape}')
-
-
     def check_loaded_model_outputs(self) -> List[bool]:
         """ Check if the model outputs are sigmoid (for segmentation)
-        
+
         Parameters
         ----------
-        
+
         Returns
         -------
         List[bool]
             List of booleans indicating if the model outputs are sigmoid
         """
         outputs = []
-        
+
         for output in self.outputs_layers:
             if len(output.shape) == 3:
                 outputs.append(True)
             else:
                 outputs.append(output.shape[-3] == 1)
-                
+
         return outputs

test/data/dummy_model/GenerateDummy.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 4,
    "id": "7ad5306f-18fc-4499-b764-ab6902ac301f",
    "metadata": {},
    "outputs": [],
@@ -33,7 +33,7 @@
     "        super().__init__()\n",
     "\n",
     "        self.conv = nn.Conv2d(3, 1, 1)\n",
-    "        \n",
+    "\n",
     "        self.softmax = nn.Softmax(dim=1)\n",
     "        self.sigmoid = nn.Sigmoid()\n",
     "\n",
@@ -62,11 +62,63 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 11,
    "id": "79c0bfcb-47c4-44c9-9c9b-c4c75b2c3b34",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([1, 3, 512, 512])\n",
+      "torch.Size([1, 1, 512, 512])\n",
+      "torch.Size([1, 3, 512, 512])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/przemek/Projects/qgis-plugin-deepness/.venv/lib/python3.10/site-packages/torch/onnx/utils.py:2095: UserWarning: Provided key red for dynamic axes is not a valid input/output name\n",
+      "  warnings.warn(\n"
+     ]
+    }
+   ],
+   "source": [
+    "x = torch.rand(1, 3, 512, 512)\n",
+    "\n",
+    "class Model(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        self.sigmoid = nn.Sigmoid()\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        print(x.shape)  # torch.Size([1, 3, 512, 512])\n",
+    "        x = (x[:,0:1]-0.5)\n",
+    "\n",
+    "        # unsqueeze to add channel dimension\n",
+    "        # x = x.unsqueeze(1)\n",
+    "\n",
+    "        return self.sigmoid(x) #, self.sigmoid(x)\n",
+    "\n",
+    "m = Model()\n",
+    "# print(m(x)[0].shape, m(x)[1].shape)\n",
+    "print(m(x).shape)\n",
+    "\n",
+    "torch.onnx.export(m,\n",
+    "                  x,\n",
+    "                  \"one_output_sigmoid_red_detector.onnx\",\n",
+    "                  export_params=True,\n",
+    "                  opset_version=12,\n",
+    "                  do_constant_folding=True,\n",
+    "                  input_names = ['input'],\n",
+    "                  output_names = ['output'],\n",
+    "                  dynamic_axes={'input' : {0 : 'batch_size'},\n",
+    "                                'output' : {0 : 'batch_size'},\n",
+    "                               # 'output2' : {0 : 'batch_size'},\n",
+    "                               })"
+   ]
   }
  ],
  "metadata": {

test/test_map_processor_segmentation_sigmoid.py

@@ -0,0 +1,72 @@
+from test.test_utils import (create_default_input_channels_mapping_for_rgba_bands, create_rlayer_from_file,
+                             create_vlayer_from_file, get_dummy_fotomap_area_crs3857_path, get_dummy_fotomap_area_path,
+                             get_dummy_fotomap_small_path, get_dummy_segmentation_model_path, get_dummy_sigmoid_model_path, init_qgis)
+from unittest.mock import MagicMock
+
+import matplotlib.pyplot as plt
+import numpy as np
+from qgis.core import QgsCoordinateReferenceSystem, QgsRectangle
+
+from deepness.common.processing_overlap import ProcessingOverlap, ProcessingOverlapOptions
+from deepness.common.processing_parameters.map_processing_parameters import ProcessedAreaType
+from deepness.common.processing_parameters.segmentation_parameters import SegmentationParameters
+from deepness.processing.map_processor.map_processor_segmentation import MapProcessorSegmentation
+from deepness.processing.models.segmentor import Segmentor
+
+RASTER_FILE_PATH = get_dummy_fotomap_small_path()
+
+VLAYER_MASK_FILE_PATH = get_dummy_fotomap_area_path()
+
+VLAYER_MASK_CRS3857_FILE_PATH = get_dummy_fotomap_area_crs3857_path()
+
+MODEL_FILE_PATH = get_dummy_sigmoid_model_path()
+
+INPUT_CHANNELS_MAPPING = create_default_input_channels_mapping_for_rgba_bands()
+
+
+def test_sigmoid_model_processing__entire_file():
+    qgs = init_qgis()
+
+    rlayer = create_rlayer_from_file(RASTER_FILE_PATH)
+    model = Segmentor(MODEL_FILE_PATH)
+
+    params = SegmentationParameters(
+        resolution_cm_per_px=3,
+        tile_size_px=model.get_input_size_in_pixels()[0],  # same x and y dimensions, so take x
+        batch_size=1,
+        local_cache=False,
+        processed_area_type=ProcessedAreaType.ENTIRE_LAYER,
+        mask_layer_id=None,
+        input_layer_id=rlayer.id(),
+        input_channels_mapping=INPUT_CHANNELS_MAPPING,
+        postprocessing_dilate_erode_size=5,
+        processing_overlap=ProcessingOverlap(ProcessingOverlapOptions.OVERLAP_IN_PERCENT, percentage=20),
+        pixel_classification__probability_threshold=0.6,
+        model=model,
+    )
+
+    map_processor = MapProcessorSegmentation(
+        rlayer=rlayer,
+        vlayer_mask=None,
+        map_canvas=MagicMock(),
+        params=params,
+    )
+
+    map_processor.run()
+    result_img = map_processor.get_result_img()
+
+    assert result_img.shape == (1, 561, 829)
+    non_zero_pixels = np.count_nonzero(result_img)
+    assert abs(non_zero_pixels - 25002) < 50  # number of RED pixels in the image
+
+    # you should see only the part of RASTER_FILE_PATH that is pure red pixels. Use snippet below for debugging
+    # from matplotlib import pyplot as plt
+    # plt.imshow(result_img[0])
+    # plt.show()
+
+    # TODO - add detailed check for pixel values once we have output channels mapping with thresholding
+
+
+if __name__ == '__main__':
+    test_sigmoid_model_processing__entire_file()
+    print('Done')

test/test_utils.py

@@ -19,13 +19,22 @@ def get_dummy_segmentation_model_path():
     return os.path.join(TEST_DATA_DIR, 'dummy_model', 'dummy_segmentation_models', 'dummy_model.onnx')
 
 
+def get_dummy_sigmoid_model_path():
+    """
+    Get path of a dummy onnx model. See details in README in model directory.
+    Model used for unit tests processing purposes
+    """
+    return os.path.join(TEST_DATA_DIR, 'dummy_model', 'one_output_sigmoid_red_detector.onnx')
+
+
 def get_dummy_segmentation_model_different_output_size_path():
     """
     Get path of a dummy onnx model. See details in README in model directory.
     Model used for unit tests processing purposes. Its output size is different than input size.
     """
     return os.path.join(TEST_DATA_DIR, 'dummy_model', 'dummy_segmentation_models', 'different_output_size_512_to_484.onnx')
 
+
 def get_dummy_segmentation_models_dict():
     """
     Get dictionary with dummy segmentation models paths. See details in README in model directory.
@@ -82,6 +91,7 @@ def get_dummy_regression_model_path_batched():
     """
     return os.path.join(TEST_DATA_DIR, 'dummy_model', 'dummy_regression_models', 'dummy_regression_model_batched.onnx')
 
+
 def get_dummy_regression_models_dict():
     """
     Get dictionary with dummy regression models paths. See details in README in model directory.