update

sorghum_pipeline/data/preprocessor.py

@@ -27,12 +27,21 @@ class ImagePreprocessor:
         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(box), dtype=float) for box in boxes], axis=-1)
-        green, red, red_edge, nir = np.split(stack, 4, axis=-1)
         
-        composite = np.concatenate([green, red_edge, red], axis=-1)
+        # Split into 4 bands and squeeze the extra dimension
+        green, red, red_edge, nir = [band.squeeze(-1) for band in np.split(stack, 4, axis=-1)]
+        
+        # Stack into (H, W, 3) composite: green, red_edge, red
+        composite = np.stack([green, red_edge, red], axis=-1)
         composite_uint8 = self.convert_to_uint8(composite)
         
-        spectral_bands = {"green": green, "red": red, "red_edge": red_edge, "nir": nir}
+        # Keep spectral bands with single channel dimension for consistency
+        spectral_bands = {
+            "green": green[..., np.newaxis], 
+            "red": red[..., np.newaxis], 
+            "red_edge": red_edge[..., np.newaxis], 
+            "nir": nir[..., np.newaxis]
+        }
         return composite_uint8, spectral_bands
     
     def create_composites(self, plants: Dict[str, Dict[str, Any]]) -> Dict[str, Dict[str, Any]]: