Merge branch '159-add-criteria-file' into 'release'

Resolve "Masques: ajout de constant.py et criteria.py"

See merge request 3d/PandoraBox/pandora2d!139

pandora2d/common.py

@@ -34,14 +34,16 @@
     from xarray import Coordinate as Coordinates
 
 import os
-from typing import Dict
+from typing import Dict, Union
 import xarray as xr
 import numpy as np
+from numpy.typing import NDArray
 
 from rasterio import Affine
 
 from pandora.common import mkdir_p, write_data_array
 from pandora2d.img_tools import remove_roi_margins
+from pandora2d.constants import Criteria
 
 
 def save_dataset(dataset: xr.Dataset, cfg: Dict, output: str) -> None:
@@ -145,3 +147,77 @@ def dataset_disp_maps(
         dataset.attrs = attributes
 
     return dataset
+
+
+def set_out_of_row_disparity_range_to_other_value(
+    data: xr.Dataset,
+    min_disp_grid: NDArray[np.floating],
+    max_disp_grid: NDArray[np.floating],
+    value: Union[int, float, Criteria],
+    data_var_name: str,
+):
+    """
+    Put special value in data (cost_volumes or criteria_dataset) 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 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
+    """
+    # WARNING: if one day we switch disp_row with disp_col index should be -2
+    ndisp_row = data[data_var_name].shape[-1]
+
+    for disp_row in range(ndisp_row):
+        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
+
+
+def set_out_of_col_disparity_range_to_other_value(
+    data: xr.Dataset,
+    min_disp_grid: NDArray[np.floating],
+    max_disp_grid: NDArray[np.floating],
+    value: Union[int, float, Criteria],
+    data_var_name: str,
+):
+    """
+    Put special value in data (cost_volumes or criteria_dataset) 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 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
+    """
+    # WARNING: if one day we switch disp_row with disp_col index should be -1
+    ndisp_col = data[data_var_name].shape[-2]
+
+    for disp_col in range(ndisp_col):
+        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

pandora2d/constants.py

@@ -0,0 +1,49 @@
+# 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.
+#
+"""
+This module contains all the parameters related to the criteria dataset, defining each bit.
+"""
+
+from enum import Flag, auto
+
+
+class Criteria(Flag):
+    """
+    Criteria class
+    """
+
+    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
+    PANDORA2D_MSK_PIXEL_LEFT_NODATA = auto()
+    # The pixel is invalid : nodata in right mask
+    PANDORA2D_MSK_PIXEL_RIGHT_NODATA = auto()
+    # The pixel is invalid : disparity is out the right image
+    PANDORA2D_MSK_PIXEL_RIGHT_DISPARITY_OUTSIDE = auto()
+    # The pixel is invalid : invalidated by validity mask of left image
+    PANDORA2D_MSK_PIXEL_INVALIDITY_MASK_LEFT = auto()
+    # The pixel is invalid : invalidated by validity mask of right 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.
+    PANDORA2D_MSK_PIXEL_PEAK_ON_EDGE = auto()
+    # The disparity is not processed because not included in the disparity range of the current point
+    PANDORA2D_MSK_PIXEL_DISPARITY_UNPROCESSED = auto()

pandora2d/criteria.py

@@ -0,0 +1,113 @@
+# 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.
+#
+"""
+This module contains functions associated to the validity mask and criteria dataset created in the cost volume step.
+"""
+
+from typing import Union
+import xarray as xr
+import numpy as np
+from numpy.typing import NDArray
+
+
+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(
+    cv: xr.Dataset, value: Union[int, float, Criteria] = Criteria.VALID, data_type: Union[np.dtype, None] = None
+) -> xr.Dataset:
+    """
+    This method creates the criteria_dataset 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.
+    :type data_type: Union[np.dtype, None], by default None
+    :return: criteria_dataset: 4D Dataset containing the criteria
+    :rtype: criteria_dataset: xr.Dataset
+    """
+    return xr.Dataset(
+        {
+            "criteria": (["row", "col", "disp_col", "disp_row"], 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},
+    )
+
+
+def set_unprocessed_disp(
+    criteria_dataset: xr.Dataset,
+    min_grid_col: NDArray[np.floating],
+    max_grid_col: NDArray[np.floating],
+    min_grid_row: NDArray[np.floating],
+    max_grid_row: NDArray[np.floating],
+):
+    """
+    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 min_grid_col: grid of min disparity col
+    :type min_grid_col: NDArray[np.floating]
+    :param max_grid_col: grid of max disparity col
+    :type max_grid_col: NDArray[np.floating]
+    :param min_grid_row: grid of min disparity row
+    :type min_grid_row: NDArray[np.floating]
+    :param max_grid_row: grid of max disparity row
+    :type max_grid_row: NDArray[np.floating]
+    """
+    # 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"
+    )
+    # 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"
+    )
+
+
+def mask_border(cost_volumes: xr.Dataset, criteria_dataset: xr.Dataset):
+    """
+    This method raises PANDORA2D_MSK_PIXEL_LEFT_BORDER criteria on the edges of the criteria_dataset
+    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
+    """
+
+    offset = cost_volumes.attrs["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

pandora2d/matching_cost/matching_cost.py

@@ -31,15 +31,14 @@
 import xarray as xr
 import numpy as np
 
-from numpy.typing import NDArray
-
 from pandora import matching_cost
 from pandora.criteria import validity_mask
 from pandora.margins import Margins
 
 
 from pandora2d import img_tools
 import pandora2d.schema as cst_schema
+from pandora2d.common import set_out_of_row_disparity_range_to_other_value
 
 
 class MatchingCost:
@@ -393,34 +392,6 @@ def compute_cost_volumes(
 
         # Delete ROI_margins attributes which we used to calculate the row coordinates in the cost_volumes
         del cost_volumes.attrs["ROI_margins_for_cv"]
-        set_out_of_disparity_range_to_nan(cost_volumes, grid_min_row, grid_max_row)
+        set_out_of_row_disparity_range_to_other_value(cost_volumes, grid_min_row, grid_max_row, np.nan, "cost_volumes")
 
         return cost_volumes
-
-
-def set_out_of_disparity_range_to_nan(
-    cost_volumes: xr.Dataset, min_disp_grid: NDArray[np.floating], max_disp_grid: NDArray[np.floating]
-):
-    """
-    Put NaNs in cost_volumes where the disparity is out of the range defined by disparity grids.
-
-    The operation is done inplace.
-
-    :param cost_volumes: cost_volumes to modify.
-    :type cost_volumes: xr.Dataset
-    :param min_disp_grid: grid of min below which cost_volumes will be set to NaN.
-    :type min_disp_grid: NDArray[np.floating]
-    :param max_disp_grid: grid of max above which cost_volumes will be set to NaN.
-    :type max_disp_grid: NDArray[np.floating]
-    """
-    # WARNING: if one day we switch dis_row with disp_col index should be -2
-    ndisp_row = cost_volumes["cost_volumes"].shape[-1]
-
-    for disp_row in range(ndisp_row):
-        masking = np.nonzero(
-            np.logical_or(
-                cost_volumes.coords["disp_row"].data[disp_row] < min_disp_grid,
-                cost_volumes.coords["disp_row"].data[disp_row] > max_disp_grid,
-            )
-        )
-        cost_volumes["cost_volumes"].data[masking[0], masking[1], :, disp_row] = np.nan