update

sorghum_pipeline/data/preprocessor.py

@@ -1,9 +1,4 @@
-"""Minimal image preprocessing following the requested composite/spectral logic.
-
-Self-contained: builds composite/spectral directly (no external imports).
-Assumes a 2x2 mosaic: TL, TR, BL, BR → green, red, red_edge, nir.
-Composite BGR order: (green, red_edge, red) so displayed RGB = (red, red_edge, green).
-"""
+"""Minimal image preprocessing following the requested composite/spectral logic."""
 
 import numpy as np
 from PIL import Image
@@ -27,34 +22,20 @@ class ImagePreprocessor:
         return norm.astype(np.uint8)
     
     def process_raw_image(self, pil_img: Image.Image) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
-        """Split 4-band mosaic, return 8-bit BGR composite and spectral dict."""
-        # Determine quadrant size
-        d_w = pil_img.size[0] // 2
-        d_h = pil_img.size[1] // 2
-        boxes = [
-            (0, 0, d_w, d_h),                    # TL → green
-            (d_w, 0, d_w * 2, d_h),              # TR → red
-            (0, d_h, d_w, d_h * 2),              # BL → red_edge
-            (d_w, d_h, d_w * 2, d_h * 2),        # BR → nir
-        ]
-        # Extract each quadrant as grayscale float
-        quads = [
-            np.array(pil_img.crop(b).convert('L'), dtype=float)
-            for b in boxes
-        ]
-        green, red, red_edge, nir = quads
-
-        # Build composite in BGR order: (green, red_edge, red)
-        comp = np.dstack([green, red_edge, red])
-        comp_uint8 = self.convert_to_uint8(comp)
-
-        spectral_bands = {
-            "green": green[..., np.newaxis],
-            "red": red[..., np.newaxis],
-            "red_edge": red_edge[..., np.newaxis],
-            "nir": nir[..., np.newaxis],
-        }
-        return comp_uint8, spectral_bands
+        """Split the 4-band RAW into tiles and build composite + spectral stack."""
+        d = pil_img.size[0] // 2
+        boxes = [(j, i, j + d, i + d)
+                 for i, j in product(range(0, pil_img.height, d),
+                                    range(0, pil_img.width,  d))]
+        stack = np.stack([np.array(pil_img.crop(b), float) for b in boxes], axis=-1)
+        # bands come in order: [green, red, red_edge, nir]
+        green, red, red_edge, nir = np.split(stack, 4, axis=-1)
+        # build composite using bands: Red, Red-Edge, Green
+        # Note: OpenCV expects BGR ordering when saving/displaying.
+        # To achieve displayed RGB = (red, red_edge, green), we arrange BGR = (green, red_edge, red).
+        comp = np.concatenate([green, red_edge, red], axis=-1)
+        comp8 = self.convert_to_uint8(comp)
+        return comp8, {"green": green, "red": red, "red_edge": red_edge, "nir": nir}
     
     def create_composites(self, plants: Dict[str, Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
         """