PR #184 - SAM - Normalization and ChipWindow

ops/segment_anything/automatic_segmentation.yaml

@@ -41,7 +41,7 @@ description:
   parameters:
     model_type: SAM's image encoder backbone architecture, among 'vit_h', 'vit_l', or 'vit_b'. Before running the workflow, make sure the desired model has been exported to the cluster by running `scripts/export_sam_models.py`. For more information, refer to the FarmVibes.AI troubleshooting page in the documentation.
     band_names: Name of raster bands that should be selected to compose the 3-channel images expected by SAM. If not provided, will try to use ["R", "G", "B"]. If only a single band name is provided, will replicate it through all three channels.
-    band_scaling: A list of floats to scale each band by to the range of [0.0, 1.0] or [0.0, 255.0]. If not provided, will default to the raster scaling parameter. If a list with a single value is provided, will use it for all three bands.
+    band_scaling: A list of floats to scale each band by to the range of [0.0, 1.0]. If not provided, will default to the raster scaling parameter. If a list with a single value is provided, will use it for all three bands.
     band_offset: A list of floats to offset each band by. If not provided, will default to the raster offset value. If a list with a single value is provided, will use it for all three bands.
     spatial_overlap: Percentage of spatial overlap between chips in the range of [0.0, 1.0).
     points_per_side: The number of points to be sampled along one side of the chip to be prompts. The total number of points is points_per_side**2.

ops/segment_anything/sam_inference.py

@@ -497,7 +497,7 @@ def generate_masks_from_grid(
                 meta = cast(Dict[str, Any], write_info_list[0]["meta"])
                 meta.update({**INT_COMPRESSION_KWARGS})
 
-                write_window = ChipWindow(
+                write_window = (
                     int(read_window.col_off - dataset.offset.width),
                     int(read_window.row_off - dataset.offset.height),
                     int(read_window.width),

ops/segment_anything_combine_masks/combine_sam_masks.py

@@ -12,10 +12,10 @@
 
 def touch_chip_boundaries(bbox: BBox, chip_window: ChipWindow) -> bool:
     return (
-        bbox[0] <= chip_window.col_offset
-        or bbox[1] <= chip_window.row_offset
-        or bbox[2] >= chip_window.col_offset + chip_window.width
-        or bbox[3] >= chip_window.row_offset + chip_window.height
+        bbox[0] <= chip_window[0]  # col_offset
+        or bbox[1] <= chip_window[1]  # row_offset
+        or bbox[2] >= chip_window[0] + chip_window[2]  # col_offset + width
+        or bbox[3] >= chip_window[1] + chip_window[3]  # row_offset + height
     )
 
 

ops/segment_anything_combine_masks/test_combine_sam_masks.py

@@ -8,7 +8,7 @@
 import xarray as xr
 from shapely import geometry as shpg
 
-from vibe_core.data.core_types import ChipWindow, gen_guid
+from vibe_core.data.core_types import gen_guid
 from vibe_core.data.rasters import CategoricalRaster, SamMaskRaster
 from vibe_dev.testing.op_tester import OpTester
 from vibe_lib.raster import save_raster_to_asset
@@ -59,7 +59,7 @@ def create_segmented_raster(
         categories=["background", "foreground"],
         mask_score=[mask_score],
         mask_bbox=[tuple([float(c) for c in mask_bbox])],  # type: ignore
-        chip_window=ChipWindow(0.0, 0.0, float(raster_size), float(raster_size)),
+        chip_window=(0.0, 0.0, float(raster_size), float(raster_size)),
     )
 
 

src/vibe_core/vibe_core/data/core_types.py

@@ -15,7 +15,6 @@
     ClassVar,
     Dict,
     List,
-    NamedTuple,
     Optional,
     Tuple,
     Type,
@@ -51,20 +50,8 @@
 """Type alias for a time range, as a tuple of two `datetime` objects (start, end)."""
 
 
-class ChipWindow(NamedTuple):
-    """Represent a window of a raster chip.
-
-    Attributes:
-        col_offset: The column offset of the window with relation to the raster chip.
-        row_offset: The row offset of the window with relation to the raster chip.
-        width: The width of the window.
-        height: The height of the window.
-    """
-
-    col_offset: float
-    row_offset: float
-    width: float
-    height: float
+ChipWindow = Tuple[float, float, float, float]
+"""Type alias representing a raster chip window, as (col_offset, row_offset, width, height)."""
 
 
 def gen_guid():

src/vibe_lib/vibe_lib/segment_anything.py

@@ -11,12 +11,13 @@
 from geopandas import GeoDataFrame
 from numpy.typing import NDArray
 from rasterio import Affine
+from rasterio.windows import Window
 from shapely.geometry.base import BaseGeometry
 from torchvision.transforms.functional import resize
 
 from vibe_core.data import GeometryCollection, Raster
 from vibe_core.data.core_types import BBox, Point
-from vibe_lib.spaceeye.chip import ChipDataset, Dims, Window
+from vibe_lib.spaceeye.chip import ChipDataset, Dims
 
 LOGGER = logging.getLogger(__name__)
 
@@ -473,7 +474,7 @@ def build_chip_preprocessing_operation(
         elif len(band_scaling) != len(band_names):
             raise ValueError(f"Expected one or three scaling parameters. Got {band_scaling}")
     else:
-        band_scaling = [raster.scale] * 3
+        band_scaling = [float(raster.scale)] * 3
     scale = np.array(band_scaling).reshape(1, 3, 1, 1)
 
     if band_offset:
@@ -483,13 +484,20 @@ def build_chip_preprocessing_operation(
         elif len(band_offset) != len(band_names):
             raise ValueError(f"Expected one or three offset parameters. Got {band_offset}")
     else:
-        band_offset = [raster.offset] * 3
+        band_offset = [float(raster.offset)] * 3
     offset = np.array(band_offset).reshape(1, 3, 1, 1)
 
     def preprocessing_operation(chip: NDArray[Any]) -> NDArray[Any]:
         normalized_chip = chip[:, band_idx, :, :] * scale + offset
-        if np.min(normalized_chip) >= 0 and np.max(normalized_chip) <= 1:
-            normalized_chip = normalized_chip * 255.0
+        if np.min(normalized_chip) < 0 or np.max(normalized_chip) > 1:
+            LOGGER.warning(
+                "Chip values are outside the expected range [0, 1] after scaling and offset. "
+                f"Found max of {np.max(normalized_chip)} and min of {np.min(normalized_chip)}."
+                "Will clip to [0, 1] and normalize to [0, 255]. Please, verify the band_scaling "
+                "and band_offset parameters of the workflow."
+            )
+        normalized_chip = np.clip(normalized_chip, 0, 1)
+        normalized_chip = normalized_chip * 255.0
         return normalized_chip.astype(np.float32)
 
     return preprocessing_operation