Merge branch '103-add-disparity-margins-cv' into 'release'

Resolve "Ajout d'une marge sur la dimension disparité du cost volume"

See merge request 3d/PandoraBox/pandora2d!97

pandora2d/disparity/disparity.py

@@ -26,6 +26,7 @@
 from typing import Dict, Tuple, Callable
 from json_checker import Or, And, Checker
 from pandora.margins.descriptors import NullMargins
+from pandora.margins import Margins
 
 import numpy as np
 import xarray as xr
@@ -200,36 +201,70 @@ def compute_disp_maps(self, cost_volumes: xr.Dataset) -> Tuple[np.ndarray, np.nd
         :rtype: tuple (numpy.ndarray, numpy.ndarray, numpy.ndarray)
         """
 
-        indices_nan = np.isnan(cost_volumes["cost_volumes"].data)
+        disparity_margins = cost_volumes.attrs["disparity_margins"]
+
+        # Check margins presence
+        if disparity_margins is not None and disparity_margins != Margins(0, 0, 0, 0):
+            margins = disparity_margins.asdict()
+            for key in margins.keys():
+                margins[key] *= cost_volumes.attrs["subpixel"]
+
+            # Get the right index when right and down margins are equals to 0
+            if margins["right"] == 0:
+                margins["right"] = -cost_volumes.sizes["disp_col"]
+            if margins["down"] == 0:
+                margins["down"] = -cost_volumes.sizes["disp_row"]
+
+            cost_volumes_user = xr.Dataset(
+                {
+                    "cost_volumes": (
+                        ["row", "col", "disp_col", "disp_row"],
+                        cost_volumes["cost_volumes"].data[
+                            :, :, margins["left"] : -margins["right"], margins["up"] : -margins["down"]
+                        ],
+                    )
+                },
+                coords={
+                    "row": cost_volumes.coords["row"],
+                    "col": cost_volumes.coords["col"],
+                    "disp_col": cost_volumes.coords["disp_col"][margins["left"] : -margins["right"]],
+                    "disp_row": cost_volumes.coords["disp_row"][margins["up"] : -margins["down"]],
+                },
+            )
+        else:
+            cost_volumes_user = cost_volumes.copy(deep=True)
+
+        indices_nan = np.isnan(cost_volumes_user["cost_volumes"].data)
 
         # Winner Takes All strategy
         if cost_volumes.attrs["type_measure"] == "max":
-            cost_volumes["cost_volumes"].data[indices_nan] = -np.inf
+
+            cost_volumes_user["cost_volumes"].data[indices_nan] = -np.inf
             # -------compute disp_map row---------
             # process of maximum for dispx
-            maps_max_col = self.extrema_split(cost_volumes, 2, np.max)
+            maps_max_col = self.extrema_split(cost_volumes_user, 2, np.max)
             # process of argmax for dispy
-            disp_map_row = cost_volumes["disp_row"].data[self.arg_split(maps_max_col, 2, np.argmax)]
+            disp_map_row = cost_volumes_user["disp_row"].data[self.arg_split(maps_max_col, 2, np.argmax)]
             # -------compute disp_map col---------
             # process of maximum for dispy
-            maps_max_row = self.extrema_split(cost_volumes, 3, np.max)
+            maps_max_row = self.extrema_split(cost_volumes_user, 3, np.max)
             # process of argmax for dispx
-            disp_map_col = cost_volumes["disp_col"].data[self.arg_split(maps_max_row, 2, np.argmax)]
+            disp_map_col = cost_volumes_user["disp_col"].data[self.arg_split(maps_max_row, 2, np.argmax)]
             # --------compute correlation score----
             score_map = self.get_score(maps_max_row, np.max)
 
         else:
             # -------compute disp_map row---------
-            cost_volumes["cost_volumes"].data[indices_nan] = np.inf
+            cost_volumes_user["cost_volumes"].data[indices_nan] = np.inf
             # process of minimum for dispx
-            maps_min_col = self.extrema_split(cost_volumes, 2, np.min)
+            maps_min_col = self.extrema_split(cost_volumes_user, 2, np.min)
             # process of argmin for disp
-            disp_map_row = cost_volumes["disp_row"].data[self.arg_split(maps_min_col, 2, np.argmin)]
+            disp_map_row = cost_volumes_user["disp_row"].data[self.arg_split(maps_min_col, 2, np.argmin)]
             # -------compute disp_map col---------
             # process of maximum for dispy
-            maps_min_row = self.extrema_split(cost_volumes, 3, np.min)
+            maps_min_row = self.extrema_split(cost_volumes_user, 3, np.min)
             # process of argmin for dispx
-            disp_map_col = cost_volumes["disp_col"].data[self.arg_split(maps_min_row, 2, np.argmin)]
+            disp_map_col = cost_volumes_user["disp_col"].data[self.arg_split(maps_min_row, 2, np.argmin)]
             # --------compute correlation score----
             score_map = self.get_score(maps_min_row, np.min)
 

pandora2d/matching_cost/matching_cost.py

@@ -33,6 +33,7 @@
 from pandora import matching_cost
 from pandora.criteria import validity_mask
 from pandora.margins.descriptors import HalfWindowMargins
+from pandora.margins import Margins
 
 
 from pandora2d import img_tools
@@ -142,7 +143,13 @@ def allocate_cost_volumes(
         return cost_volumes
 
     def allocate_cost_volume_pandora(
-        self, img_left: xr.Dataset, img_right: xr.Dataset, grid_min_col: np.ndarray, grid_max_col: np.ndarray, cfg: Dict
+        self,
+        img_left: xr.Dataset,
+        img_right: xr.Dataset,
+        grid_min_col: np.ndarray,
+        grid_max_col: np.ndarray,
+        cfg: Dict,
+        margins: Margins = None,
     ) -> None:
         """
 
@@ -158,13 +165,20 @@ def allocate_cost_volume_pandora(
         :type grid_max_col: np.ndarray
         :param cfg: matching_cost computation configuration
         :type cfg: Dict
+        :param margins: refinement margins
+        :type margins: Margins
         :return: None
         """
         # Adapt Pandora matching cost configuration
         copy_matching_cost_cfg_with_step = copy.deepcopy(cfg["pipeline"]["matching_cost"])
         copy_matching_cost_cfg_with_step["step"] = self._step_col
         img_left.attrs["disparity_source"] = img_left.attrs["col_disparity_source"]
 
+        if margins is not None:
+
+            grid_min_col -= margins.left
+            grid_max_col += margins.right
+
         # Initialize Pandora matching cost
         self.pandora_matching_cost_ = matching_cost.AbstractMatchingCost(**copy_matching_cost_cfg_with_step)
 
@@ -189,6 +203,7 @@ def compute_cost_volumes(
         grid_max_col: np.ndarray,
         grid_min_row: np.ndarray,
         grid_max_row: np.ndarray,
+        margins: Margins = None,
     ) -> xr.Dataset:
         """
 
@@ -210,6 +225,8 @@ def compute_cost_volumes(
         :type grid_min_row: np.ndarray
         :param grid_max_row: grid containing max disparities for rows.
         :type grid_max_row: np.ndarray
+        :param margins: refinement margins
+        :type margins: Margins
         :return: cost_volumes: 4D Dataset containing the cost_volumes
         :rtype: cost_volumes: xr.Dataset
         """
@@ -219,6 +236,11 @@ def compute_cost_volumes(
         # Adapt Pandora matching cost configuration
         img_left.attrs["disparity_source"] = img_left.attrs["col_disparity_source"]
 
+        if margins is not None:
+
+            grid_min_row -= margins.up
+            grid_max_row += margins.down
+
         # Obtain absolute min and max disparities
         min_row, max_row = self.pandora_matching_cost_.get_min_max_from_grid(grid_min_row, grid_max_row)
         min_col, max_col = self.pandora_matching_cost_.get_min_max_from_grid(grid_min_col, grid_max_col)
@@ -294,6 +316,7 @@ def compute_cost_volumes(
         del cost_volumes.attrs["disparity_source"]
         cost_volumes.attrs["col_disparity_source"] = img_left.attrs["col_disparity_source"]
         cost_volumes.attrs["row_disparity_source"] = img_left.attrs["row_disparity_source"]
+        cost_volumes.attrs["disparity_margins"] = margins
         cost_volumes.attrs["step"] = self.cfg["step"]
 
         # Delete ROI_margins attributes which we used to calculate the row coordinates in the cost_volumes

pandora2d/state_machine.py

@@ -318,8 +318,9 @@ def matching_cost_prepare(self, cfg: Dict[str, dict], input_step: str) -> None:
         :return: None
         """
         self.matching_cost_ = matching_cost.MatchingCost(cfg["pipeline"][input_step])
+
         self.matching_cost_.allocate_cost_volume_pandora(
-            self.left_img, self.right_img, self.disp_min_col, self.disp_max_col, cfg
+            self.left_img, self.right_img, self.disp_min_col, self.disp_max_col, cfg, self.margins.get("refinement")
         )
 
     def estimation_run(self, cfg: Dict[str, dict], input_step: str) -> None:
@@ -366,6 +367,7 @@ def matching_cost_run(self, _, __) -> None:
             self.disp_max_col,
             self.disp_min_row,
             self.disp_max_row,
+            self.margins.get("refinement"),
         )
 
     def disp_maps_run(self, cfg: Dict[str, dict], input_step: str) -> None:

tests/functional_tests/matching_cost/test_disparity_margins.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python
+#
+# Copyright (c) 2024 Centre National d'Etudes Spatiales (CNES).
+#
+# This file is part of PANDORA2D
+#
+#     https://github.com/CNES/Pandora2D
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Tests disparity margins in matching cost step
+"""
+
+import pytest
+import numpy as np
+import xarray as xr
+
+from pandora2d.img_tools import add_disparity_grid
+from pandora2d.state_machine import Pandora2DMachine
+
+
+class TestDisparityMargins:
+    """
+    Test disparity margins in the cost volume
+    """
+
+    @pytest.fixture()
+    def create_datasets(self):
+        """
+        Creates left and right datasets
+        """
+
+        data = np.full((10, 10), 1)
+        left = xr.Dataset(
+            {"im": (["row", "col"], data)},
+            coords={"row": np.arange(data.shape[0]), "col": np.arange(data.shape[1])},
+        )
+
+        add_disparity_grid(left, [1, 3], [-2, 2])
+
+        left.attrs = {
+            "no_data_img": -9999,
+            "valid_pixels": 0,
+            "no_data_mask": 1,
+            "crs": None,
+            "col_disparity_source": [1, 3],
+            "row_disparity_source": [-2, 2],
+        }
+
+        data = np.full((10, 10), 1)
+        right = xr.Dataset(
+            {"im": (["row", "col"], data)},
+            coords={"row": np.arange(data.shape[0]), "col": np.arange(data.shape[1])},
+        )
+
+        right.attrs = {
+            "no_data_img": -9999,
+            "valid_pixels": 0,
+            "no_data_mask": 1,
+            "crs": None,
+            "col_disparity_source": [1, 3],
+            "row_disparity_source": [-2, 2],
+        }
+
+        return left, right
+
+    @pytest.fixture()
+    def config(self, subpix, refinement_config, matching_cost_method):
+        return {
+            # "input": input_config,
+            "pipeline": {
+                "matching_cost": {
+                    "matching_cost_method": matching_cost_method,
+                    "window_size": 1,
+                    "step": [1, 1],
+                    "subpix": subpix,
+                },
+                "disparity": {"disparity_method": "wta", "invalid_disparity": -6},
+                "refinement": refinement_config,
+            }
+        }
+
+    @pytest.mark.parametrize("matching_cost_method", ["sad", "ssd", "zncc"])
+    @pytest.mark.parametrize(
+        ["subpix", "refinement_config", "cv_shape_expected", "disp_col_expected", "disp_row_expected"],
+        [
+            pytest.param(
+                1,
+                {"refinement_method": "dichotomy", "iterations": 1, "filter": "bicubic"},
+                (10, 10, 7, 9),
+                [-1, 0, 1, 2, 3, 4, 5],
+                [-4, -3, -2, -1, 0, 1, 2, 3, 4],
+                id="Subpix=1 and refinement_method=dichotomy",
+            ),
+            pytest.param(
+                1,
+                {
+                    "refinement_method": "interpolation",
+                },
+                (10, 10, 9, 11),
+                [-2, -1, 0, 1, 2, 3, 4, 5, 6],
+                [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5],
+                id="Subpix=1 and refinement_method=interpolation",
+            ),
+            pytest.param(
+                1,
+                {
+                    "refinement_method": "optical_flow",
+                },
+                (10, 10, 3, 5),
+                [1, 2, 3],
+                [-2, -1, 0, 1, 2],
+                id="Subpix=1 and refinement_method=optical_flow",
+            ),
+            pytest.param(
+                2,
+                {"refinement_method": "dichotomy", "iterations": 1, "filter": "bicubic"},
+                (10, 10, 13, 17),
+                [-1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5],
+                [-4, -3.5, -3, -2.5, -2, -1.5, -1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4],
+                id="Subpix=2 and refinement_method=dichotomy",
+            ),
+            pytest.param(
+                2,
+                {
+                    "refinement_method": "interpolation",
+                },
+                (10, 10, 17, 21),
+                [-2, -1.5, -1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5, 6],
+                [-5, -4.5, -4, -3.5, -3, -2.5, -2, -1.5, -1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5],
+                id="Subpix=2 and refinement_method=interpolation",
+            ),
+            pytest.param(
+                2,
+                {
+                    "refinement_method": "optical_flow",
+                },
+                (10, 10, 5, 9),
+                [1, 1.5, 2, 2.5, 3],
+                [-2, -1.5, -1, -0.5, 0, 0.5, 1, 1.5, 2],
+                id="Subpix=2 and refinement_method=optical_flow",
+            ),
+            pytest.param(
+                4,
+                {"refinement_method": "dichotomy", "iterations": 1, "filter": "bicubic"},
+                (10, 10, 25, 33),
+                np.arange(-1, 5.25, 0.25),
+                np.arange(-4, 4.25, 0.25),
+                id="Subpix=4 and refinement_method=dichotomy",
+            ),
+            pytest.param(
+                4,
+                {
+                    "refinement_method": "interpolation",
+                },
+                (10, 10, 33, 41),
+                np.arange(-2, 6.25, 0.25),
+                np.arange(-5, 5.25, 0.25),
+                id="Subpix=4 and refinement_method=interpolation",
+            ),
+            pytest.param(
+                4,
+                {
+                    "refinement_method": "optical_flow",
+                },
+                (10, 10, 9, 17),
+                [1, 1.25, 1.5, 1.75, 2, 2.25, 2.5, 2.75, 3],
+                [-2, -1.75, -1.5, -1.25, -1, -0.75, -0.5, -0.25, 0, 0.25, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2],
+                id="Subpix=4 and refinement_method=optical_flow",
+            ),
+        ],
+    )
+    def test_disparity_margins_in_cost_volumes(
+        self, cv_shape_expected, disp_col_expected, disp_row_expected, config, create_datasets
+    ):
+        """
+        Test that the disparity margins are correctly added in the cost volumes
+        according to the refinement margins.
+        """
+
+        pandora2d_machine = Pandora2DMachine()
+
+        img_left, img_right = create_datasets
+
+        pandora2d_machine.check_conf(config)
+        pandora2d_machine.run_prepare(img_left, img_right, config)
+        pandora2d_machine.run("matching_cost", config)
+
+        cost_volumes = pandora2d_machine.cost_volumes["cost_volumes"]
+
+        np.testing.assert_array_equal(cost_volumes.shape, cv_shape_expected)
+        np.testing.assert_array_equal(cost_volumes.disp_col, disp_col_expected)
+        np.testing.assert_array_equal(cost_volumes.disp_row, disp_row_expected)