Merge branch '160-masques-ajout-des-criteres-left_no_data-et-right_no…

…_data' into 'release'

Resolve "Masques: ajout des critères left_no_data et right_no_data"

See merge request 3d/PandoraBox/pandora2d!150

pandora2d/common.py

@@ -34,7 +34,7 @@
     from xarray import Coordinate as Coordinates
 
 import os
-from typing import Dict, Union, Tuple
+from typing import Dict, Union, Tuple, List
 import xarray as xr
 import numpy as np
 from numpy.typing import NDArray
@@ -196,102 +196,95 @@ def dataset_disp_maps(
 
 
 def set_out_of_row_disparity_range_to_other_value(
-    data: xr.Dataset,
+    data: xr.DataArray,
     min_disp_grid: NDArray[np.floating],
     max_disp_grid: NDArray[np.floating],
     value: Union[int, float, Criteria],
-    data_var_name: str,
-):
+    global_disparity_range: Union[None, List[int]] = None,
+) -> None:
     """
-    Put special value in data (cost_volumes or criteria_dataset) where the disparity is out of the range defined
-    by disparity grids.
+    Put special value in data  where the disparity is out of the range defined by disparity grids.
 
     The operation is done inplace.
 
-    :param data: cost_volumes or criteria_dataset to modify.
-    :type data: xr.Dataset 4D
+    :param data: cost_volumes or criteria_dataarray to modify.
+    :type data: xr.DataArray 4D
     :param min_disp_grid: grid of min disparity.
     :type min_disp_grid: NDArray[np.floating]
     :param max_disp_grid: grid of max disparity.
     :type max_disp_grid: NDArray[np.floating]
     :param value: value to set on data.
     :type value: Union[int, float, Criteria]
-    :param data_var_name: name of xarray.DataArray to set new value.
-    :type data_var_name: str
+    :param global_disparity_range:
+    :type global_disparity_range:
     """
     # WARNING: if one day we switch disp_row with disp_col index should be -2
-    ndisp_row = data[data_var_name].shape[-1]
+    ndisp_row = data.shape[-1]
 
     # We want to put special value on points that are not in the global disparity range (row_disparity_source)
-    if data_var_name == "cost_volumes":
-        for disp_row in range(ndisp_row):
+    for disp_row in range(ndisp_row):
+        if global_disparity_range is not None:  # Case we are working with cost volume
             masking = np.nonzero(
                 np.logical_or(
                     (data.coords["disp_row"].data[disp_row] < min_disp_grid)
-                    & (data.coords["disp_row"].data[disp_row] >= data.attrs["row_disparity_source"][0]),
+                    & (data.coords["disp_row"].data[disp_row] >= global_disparity_range[0]),
                     (data.coords["disp_row"].data[disp_row] > max_disp_grid)
-                    & (data.coords["disp_row"].data[disp_row] <= data.attrs["row_disparity_source"][1]),
+                    & (data.coords["disp_row"].data[disp_row] <= global_disparity_range[1]),
                 )
             )
-            data[data_var_name].data[masking[0], masking[1], :, disp_row] = value
-
-    else:
-        for disp_row in range(ndisp_row):
+        else:
             masking = np.nonzero(
                 np.logical_or(
                     data.coords["disp_row"].data[disp_row] < min_disp_grid,
                     data.coords["disp_row"].data[disp_row] > max_disp_grid,
                 )
             )
-            data[data_var_name].data[masking[0], masking[1], :, disp_row] = value
+        data.data[masking[0], masking[1], :, disp_row] = value
 
 
 def set_out_of_col_disparity_range_to_other_value(
-    data: xr.Dataset,
+    data: xr.DataArray,
     min_disp_grid: NDArray[np.floating],
     max_disp_grid: NDArray[np.floating],
     value: Union[int, float, Criteria],
-    data_var_name: str,
-):
+    global_disparity_range: Union[None, List[int]] = None,
+) -> None:
     """
-    Put special value in data (cost_volumes or criteria_dataset) where the disparity is out of the range defined
+    Put special value in data (cost_volumes or criteria_dataarray) where the disparity is out of the range defined
     by disparity grids.
 
     The operation is done inplace.
 
-    :param data: cost_volumes or criteria_dataset to modify.
-    :type data: xr.Dataset 4D
+    :param data: cost_volumes or criteria_dataarray to modify.
+    :type data: xr.DataArray 4D
     :param min_disp_grid: grid of min disparity.
     :type min_disp_grid: NDArray[np.floating]
     :param max_disp_grid: grid of max disparity.
     :type max_disp_grid: NDArray[np.floating]
     :param value: value to set on data.
     :type value: Union[int, float, Criteria]
-    :param data_var_name: name of xarray.DataArray to set new value.
-    :type data_var_name: str
+    :param global_disparity_range:
+    :type global_disparity_range:
     """
     # WARNING: if one day we switch disp_row with disp_col index should be -1
-    ndisp_col = data[data_var_name].shape[-2]
+    ndisp_col = data.shape[-2]
 
     # We want to put special value on points that are not in the global disparity range (col_disparity_source)
-    if data_var_name == "cost_volumes":
-        for disp_col in range(ndisp_col):
+    for disp_col in range(ndisp_col):
+        if global_disparity_range is not None:  # Case we are working with cost volume
             masking = np.nonzero(
                 np.logical_or(
                     (data.coords["disp_col"].data[disp_col] < min_disp_grid)
-                    & (data.coords["disp_col"].data[disp_col] >= data.attrs["col_disparity_source"][0]),
+                    & (data.coords["disp_col"].data[disp_col] >= global_disparity_range[0]),
                     (data.coords["disp_col"].data[disp_col] > max_disp_grid)
-                    & (data.coords["disp_col"].data[disp_col] <= data.attrs["col_disparity_source"][1]),
+                    & (data.coords["disp_col"].data[disp_col] <= global_disparity_range[1]),
                 )
             )
-            data[data_var_name].data[masking[0], masking[1], disp_col, :] = value
-
-    else:
-        for disp_col in range(ndisp_col):
+        else:
             masking = np.nonzero(
                 np.logical_or(
                     data.coords["disp_col"].data[disp_col] < min_disp_grid,
                     data.coords["disp_col"].data[disp_col] > max_disp_grid,
                 )
             )
-            data[data_var_name].data[masking[0], masking[1], disp_col, :] = value
+        data.data[masking[0], masking[1], disp_col, :] = value

pandora2d/constants.py

@@ -30,20 +30,22 @@ class Criteria(Flag):
 
     VALID = 0
 
-    # The pixel is invalid : border of left image according to window size
     PANDORA2D_MSK_PIXEL_LEFT_BORDER = auto()
-    # The pixel is invalid : nodata in left mask
+    """The pixel is invalid : border of left image according to window size."""
     PANDORA2D_MSK_PIXEL_LEFT_NODATA = auto()
-    # The pixel is invalid : nodata in right mask
+    """The pixel is invalid : nodata in left mask."""
     PANDORA2D_MSK_PIXEL_RIGHT_NODATA = auto()
-    # The pixel is invalid : disparity is out the right image
+    """The pixel is invalid : nodata in right mask."""
     PANDORA2D_MSK_PIXEL_RIGHT_DISPARITY_OUTSIDE = auto()
-    # The pixel is invalid : invalidated by validity mask of left image
+    """The pixel is invalid : disparity is out the right image."""
     PANDORA2D_MSK_PIXEL_INVALIDITY_MASK_LEFT = auto()
-    # The pixel is invalid : invalidated by validity mask of right image
+    """The pixel is invalid : invalidated by validity mask of left image."""
     PANDORA2D_MSK_PIXEL_INVALIDITY_MASK_RIGHT = auto()
-    # The pixel is invalid : The correlation peak is at the edge of disparity range.
-    # The calculations stopped at the pixellic stage.
+    """The pixel is invalid : invalidated by validity mask of right image."""
     PANDORA2D_MSK_PIXEL_PEAK_ON_EDGE = auto()
-    # The disparity is not processed because not included in the disparity range of the current point
+    """
+    The pixel is invalid : The correlation peak is at the edge of disparity range.
+    The calculations stopped at the pixellic stage.
+    """
     PANDORA2D_MSK_PIXEL_DISPARITY_UNPROCESSED = auto()
+    """The disparity is not processed because not included in the disparity range of the current point."""

pandora2d/criteria.py

@@ -17,48 +17,47 @@
 # limitations under the License.
 #
 """
-This module contains functions associated to the validity mask and criteria dataset created in the cost volume step.
+This module contains functions associated to the validity mask and criteria dataarray created in the cost volume step.
 """
-
+import itertools
 from typing import Union
 import xarray as xr
 import numpy as np
 from numpy.typing import NDArray
 
-
+from pandora.criteria import binary_dilation_msk
 from pandora2d.constants import Criteria
 from pandora2d.common import (
     set_out_of_col_disparity_range_to_other_value,
     set_out_of_row_disparity_range_to_other_value,
 )
 
 
-def allocate_criteria_dataset(
+def allocate_criteria_dataarray(
     cv: xr.Dataset, value: Union[int, float, Criteria] = Criteria.VALID, data_type: Union[np.dtype, None] = None
-) -> xr.Dataset:
+) -> xr.DataArray:
     """
-    This method creates the criteria_dataset with the same dimensions as cost_volumes (cv).
+    This method creates the criteria_dataarray with the same dimensions as cost_volumes (cv).
     Initially, all points are considered valid and have the value XX.
 
     :param cv: cost_volumes
     :type cv: 4D xarray.Dataset
     :param value: value representing the valid criteria, by default Criteria.VALID = 0
     :type value: Union[int, float, Criteria]
-    :param data_type: the desired data-type for the criteria_dataset.
+    :param data_type: the desired data-type for the criteria_dataarray.
     :type data_type: Union[np.dtype, None], by default None
-    :return: criteria_dataset: 4D Dataset containing the criteria
-    :rtype: criteria_dataset: xr.Dataset
+    :return: criteria_dataarray: 4D DataArray containing the criteria
+    :rtype: criteria_dataarray: xr.DataArray
     """
-    return xr.Dataset(
-        {
-            "criteria": (["row", "col", "disp_col", "disp_row"], np.full(cv.cost_volumes.shape, value, data_type)),
-        },
+    return xr.DataArray(
+        np.full(cv.cost_volumes.shape, value, data_type),
         coords={"row": cv.row.data, "col": cv.col.data, "disp_col": cv.disp_col.data, "disp_row": cv.disp_row.data},
+        dims=["row", "col", "disp_col", "disp_row"],
     )
 
 
 def set_unprocessed_disp(
-    criteria_dataset: xr.Dataset,
+    criteria_dataarray: xr.DataArray,
     min_grid_col: NDArray[np.floating],
     max_grid_col: NDArray[np.floating],
     min_grid_row: NDArray[np.floating],
@@ -68,8 +67,8 @@ def set_unprocessed_disp(
     This method sets PANDORA2D_MSK_PIXEL_DISPARITY_UNPROCESSED to points for disparities that will not be processed,
     based on the disparity grids provided.
 
-    :param criteria_dataset: 4D Dataset containing the criteria
-    :type criteria_dataset: xr.Dataset 4D
+    :param criteria_dataarray: 4D DataArray containing the criteria
+    :type criteria_dataarray: xr.DataArray 4D
     :param min_grid_col: grid of min disparity col
     :type min_grid_col: NDArray[np.floating]
     :param max_grid_col: grid of max disparity col
@@ -81,63 +80,111 @@ def set_unprocessed_disp(
     """
     # Check col disparity
     set_out_of_col_disparity_range_to_other_value(
-        criteria_dataset, min_grid_col, max_grid_col, Criteria.PANDORA2D_MSK_PIXEL_DISPARITY_UNPROCESSED, "criteria"
+        criteria_dataarray, min_grid_col, max_grid_col, Criteria.PANDORA2D_MSK_PIXEL_DISPARITY_UNPROCESSED
     )
     # Check row disparity
     set_out_of_row_disparity_range_to_other_value(
-        criteria_dataset, min_grid_row, max_grid_row, Criteria.PANDORA2D_MSK_PIXEL_DISPARITY_UNPROCESSED, "criteria"
+        criteria_dataarray, min_grid_row, max_grid_row, Criteria.PANDORA2D_MSK_PIXEL_DISPARITY_UNPROCESSED
     )
 
 
-def mask_border(cost_volumes: xr.Dataset, criteria_dataset: xr.Dataset):
+def mask_border(offset: int, criteria_dataarray: xr.DataArray) -> None:
     """
-    This method raises PANDORA2D_MSK_PIXEL_LEFT_BORDER criteria on the edges of the criteria_dataset
+    This method raises PANDORA2D_MSK_PIXEL_LEFT_BORDER criteria on the edges of the criteria_dataarray
     for each of the disparities.
 
     PANDORA2D_MSK_PIXEL_LEFT_BORDER criteria is non-cumulative, so this method will be called last.
 
-    :param cost_volumes: 4D xarray.Dataset
-    :type cost_volumes: 4D xarray.Dataset
-    :param criteria_dataset: 4D xarray.Dataset with all criteria
-    :type criteria_dataset: 4D xarray.Dataset
+    :param offset: offset
+    :type offset: int
+    :param criteria_dataarray: 4D xarray.DataArray with all criteria
+    :type criteria_dataarray: 4D xarray.DataArray
     """
 
-    offset = cost_volumes.offset_row_col
     if offset > 0:
 
         # Raise criteria 0 on border of criteria_disp_col according to offset value
-        criteria_dataset.criteria.data[:offset, :, :, :] = Criteria.PANDORA2D_MSK_PIXEL_LEFT_BORDER
-        criteria_dataset.criteria.data[-offset:, :, :, :] = Criteria.PANDORA2D_MSK_PIXEL_LEFT_BORDER
-        criteria_dataset.criteria.data[:, :offset, :, :] = Criteria.PANDORA2D_MSK_PIXEL_LEFT_BORDER
-        criteria_dataset.criteria.data[:, -offset:, :, :] = Criteria.PANDORA2D_MSK_PIXEL_LEFT_BORDER
+        criteria_dataarray.data[:offset, :, :, :] = Criteria.PANDORA2D_MSK_PIXEL_LEFT_BORDER
+        criteria_dataarray.data[-offset:, :, :, :] = Criteria.PANDORA2D_MSK_PIXEL_LEFT_BORDER
+        criteria_dataarray.data[:, :offset, :, :] = Criteria.PANDORA2D_MSK_PIXEL_LEFT_BORDER
+        criteria_dataarray.data[:, -offset:, :, :] = Criteria.PANDORA2D_MSK_PIXEL_LEFT_BORDER
 
 
-def mask_disparity_outside_right_image(cost_volumes: xr.Dataset, criteria_dataset: xr.Dataset):
+def mask_disparity_outside_right_image(offset: int, criteria_dataarray: xr.DataArray) -> None:
     """
     This method raises PANDORA2D_MSK_PIXEL_RIGHT_DISPARITY_OUTSIDE criteria for points with disparity dimension outside
     the right image
 
-    :param cost_volumes: 4D xarray.Dataset
-    :type cost_volumes: 4D xarray.Dataset
-    :param criteria_dataset: 4D xarray.Dataset with all criteria
-    :type criteria_dataset: 4D xarray.Dataset
+    :param offset: offset
+    :type offset: int
+    :param criteria_dataarray: 4D xarray.DataArray with all criteria
+    :type criteria_dataarray: 4D xarray.DataArray
     """
-    offset = cost_volumes.offset_row_col
-    col_coords = criteria_dataset.col.values
-    row_coords = criteria_dataset.row.values
+    col_coords = criteria_dataarray.col.values
+    row_coords = criteria_dataarray.row.values
 
     # Condition where the window is outside the image
     condition = (
-        (criteria_dataset.row + criteria_dataset.disp_row < row_coords[0] + offset)
-        | (criteria_dataset.row + criteria_dataset.disp_row > row_coords[-1] - offset)
-        | (criteria_dataset.col + criteria_dataset.disp_col < col_coords[0] + offset)
-        | (criteria_dataset.col + criteria_dataset.disp_col > col_coords[-1] - offset)
+        (criteria_dataarray.row + criteria_dataarray.disp_row < row_coords[0] + offset)
+        | (criteria_dataarray.row + criteria_dataarray.disp_row > row_coords[-1] - offset)
+        | (criteria_dataarray.col + criteria_dataarray.disp_col < col_coords[0] + offset)
+        | (criteria_dataarray.col + criteria_dataarray.disp_col > col_coords[-1] - offset)
     )
 
-    # Swapaxes to have same shape as cost_volumes and criteria_dataset
+    # Swapaxes to have same shape as cost_volumes and criteria_dataarray
     condition_swap = condition.data.swapaxes(1, 3).swapaxes(1, 2)
 
-    # Update criteria dataset
-    criteria_dataset.criteria.data[condition_swap] = (
-        criteria_dataset.criteria.data[condition_swap] | Criteria.PANDORA2D_MSK_PIXEL_RIGHT_DISPARITY_OUTSIDE
+    # Update criteria dataarray
+    criteria_dataarray.data[condition_swap] = (
+        criteria_dataarray.data[condition_swap] | Criteria.PANDORA2D_MSK_PIXEL_RIGHT_DISPARITY_OUTSIDE
     )
+
+
+def mask_left_no_data(left_image: xr.Dataset, window_size: int, criteria_dataaray: xr.DataArray) -> None:
+    """
+    Set Criteria.PANDORA2D_MSK_PIXEL_LEFT_NODATA on pixels where a no_data is present in the window around its
+    position in the mask.
+
+    :param left_image: left image with `msk` data var.
+    :type left_image: xr.Dataset
+    :param window_size: window size
+    :type window_size: int
+    :param criteria_dataaray: criteria dataarray to update
+    :type criteria_dataaray: xr.DataArray
+    """
+    dilated_mask = binary_dilation_msk(left_image, window_size)
+    criteria_dataaray.data[dilated_mask, ...] |= Criteria.PANDORA2D_MSK_PIXEL_LEFT_NODATA
+
+
+def mask_right_no_data(img_right: xr.Dataset, window_size: int, criteria_dataarray: xr.DataArray) -> None:
+    """
+    Set Criteria.PANDORA2D_MSK_PIXEL_RIGHT_NODATA on pixels where a no_data is present in the window around its
+    position in the mask shift by its disparity.
+
+    :param img_right: right image with `msk` data var.
+    :type img_right: xr.Dataset
+    :param window_size: window size
+    :type window_size: int
+    :param criteria_dataarray:
+    :type criteria_dataarray:
+    """
+    right_criteria_mask = np.full_like(img_right["msk"], Criteria.VALID, dtype=Criteria)
+    right_binary_mask = binary_dilation_msk(img_right, window_size)
+    right_criteria_mask[right_binary_mask] |= Criteria.PANDORA2D_MSK_PIXEL_RIGHT_NODATA
+    for row_disp, col_disp in itertools.product(
+        criteria_dataarray.coords["disp_row"], criteria_dataarray.coords["disp_col"]
+    ):
+        row_disp, col_disp = row_disp.data, col_disp.data
+        # We arrange tests to avoid the slice [:0], which doesn’t work, while [0:] is fine.
+        msk_row_slice = np.s_[:row_disp] if row_disp < 0 else np.s_[row_disp:]
+        msk_col_slice = np.s_[:col_disp] if col_disp < 0 else np.s_[col_disp:]
+        criteria_row_slice = np.s_[-row_disp:] if row_disp <= 0 else np.s_[:-row_disp]
+        criteria_col_slice = np.s_[-col_disp:] if col_disp <= 0 else np.s_[:-col_disp]
+        criteria_dataarray.loc[
+            {
+                "row": criteria_dataarray.coords["row"][criteria_row_slice],
+                "col": criteria_dataarray.coords["col"][criteria_col_slice],
+                "disp_col": col_disp,
+                "disp_row": row_disp,
+            }
+        ] |= right_criteria_mask[msk_row_slice, msk_col_slice]