Update yolo_predictor.py

yolo_predictor.py

@@ -7,6 +7,7 @@ from ndvi_predictor import normalize_rgb, predict_ndvi
 import tempfile
 from rasterio.transform import from_bounds
 from PIL import Image
+import tifffile
 
 def load_yolo_model(model_path):
     """Load YOLO model from .pt file"""
@@ -60,35 +61,23 @@ def create_4channel_tiff(rgb_array, ndvi_array, output_path):
     height, width = rgb_array.shape[:2]
     
     # Stack RGB and NDVI to create 4-channel image
-    four_channel = np.zeros((height, width, 4), dtype=rgb_array.dtype)
-    four_channel[:, :, :3] = rgb_array  # RGB channels
+    four_channel = np.zeros((4, height, width), dtype=np.float32)
     
-    # Normalize NDVI to match RGB data type range
+    # Convert RGB to proper format and range
     if rgb_array.dtype == np.uint8:
-        # Scale NDVI from [-1, 1] to [0, 255]
-        ndvi_scaled = ((ndvi_array + 1) * 127.5).astype(np.uint8)
+        rgb_normalized = rgb_array.astype(np.float32) / 255.0
     else:
-        # Keep NDVI in original range for float types
-        ndvi_scaled = ndvi_array.astype(rgb_array.dtype)
-    
-    four_channel[:, :, 3] = ndvi_scaled  # NDVI channel
-    
-    # Create transform (assuming no specific georeferencing needed)
-    transform = from_bounds(0, 0, width, height, width, height)
-    
-    # Write 4-channel TIFF
-    with rasterio.open(
-        output_path,
-        'w',
-        driver='GTiff',
-        height=height,
-        width=width,
-        count=4,
-        dtype=four_channel.dtype,
-        transform=transform
-    ) as dst:
-        for i in range(4):
-            dst.write(four_channel[:, :, i], i + 1)
+        rgb_normalized = rgb_array.astype(np.float32)
+    
+    # Assign channels in (C, H, W) format for rasterio
+    four_channel[0] = rgb_normalized[:, :, 0]  # Red
+    four_channel[1] = rgb_normalized[:, :, 1]  # Green  
+    four_channel[2] = rgb_normalized[:, :, 2]  # Blue
+    four_channel[3] = ndvi_array.astype(np.float32)  # NDVI
+    
+    # Use tifffile for better compatibility with YOLO
+    import tifffile
+    tifffile.imwrite(output_path, four_channel, photometric='rgb')
 
 def load_4channel_tiff(image_path):
     """
@@ -101,16 +90,38 @@ def load_4channel_tiff(image_path):
         rgb_array: RGB channels as numpy array (H, W, 3)
         ndvi_array: NDVI channel as numpy array (H, W)
     """
-    with rasterio.open(image_path) as src:
-        # Read all 4 channels
-        channels = src.read()  # Shape: (4, H, W)
+    try:
+        with rasterio.open(image_path) as src:
+            # Read all 4 channels
+            channels = src.read()  # Shape: (4, H, W)
+            
+            # Extract RGB and NDVI
+            rgb_array = np.transpose(channels[:3], (1, 2, 0))  # (H, W, 3)
+            ndvi_array = channels[3]  # (H, W)
+            
+            # If NDVI was scaled to uint8, convert back to [-1, 1] range
+            if channels.dtype == np.uint8:
+                ndvi_array = (ndvi_array.astype(np.float32) / 127.5) - 1
+            
+            return rgb_array, ndvi_array
+    except Exception as e:
+        # Try with tifffile as fallback
+        import tifffile
+        img_array = tifffile.imread(image_path)
         
-        # Extract RGB and NDVI
-        rgb_array = np.transpose(channels[:3], (1, 2, 0))  # (H, W, 3)
-        ndvi_array = channels[3]  # (H, W)
+        if len(img_array.shape) == 3 and img_array.shape[0] == 4:
+            # Shape is (4, H, W)
+            rgb_array = np.transpose(img_array[:3], (1, 2, 0))  # (H, W, 3)
+            ndvi_array = img_array[3]  # (H, W)
+        elif len(img_array.shape) == 3 and img_array.shape[2] == 4:
+            # Shape is (H, W, 4)
+            rgb_array = img_array[:, :, :3]  # (H, W, 3)
+            ndvi_array = img_array[:, :, 3]  # (H, W)
+        else:
+            raise ValueError(f"Unexpected image shape: {img_array.shape}")
         
-        # If NDVI was scaled to uint8, convert back to [-1, 1] range
-        if channels.dtype == np.uint8:
+        # Normalize NDVI if needed
+        if img_array.dtype == np.uint8:
             ndvi_array = (ndvi_array.astype(np.float32) / 127.5) - 1
         
         return rgb_array, ndvi_array
@@ -129,28 +140,57 @@ def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
     Returns:
         results: YOLO results object
     """
-    # Try to load as 4-channel TIFF first, fall back to RGB
+    rgb_array = None
+    
+    # Try multiple methods to load the image
     try:
-        with rasterio.open(image_path) as src:
-            if src.count == 4:
-                # Load 4-channel TIFF
-                rgb_array, _ = load_4channel_tiff(image_path)
-            elif src.count == 3:
-                # Load as RGB TIFF
-                channels = src.read()
-                rgb_array = np.transpose(channels, (1, 2, 0))
-            else:
-                raise ValueError(f"Unsupported number of channels: {src.count}")
-    except:
-        # Fall back to PIL for standard image formats
-        img = Image.open(image_path).convert("RGB")
-        rgb_array = np.array(img)
+        # Method 1: Try with tifffile first (best for complex TIFF files)
+        import tifffile
+        img_array = tifffile.imread(image_path)
+        
+        if len(img_array.shape) == 3:
+            if img_array.shape[0] == 4:
+                # Shape is (4, H, W) - extract RGB
+                rgb_array = np.transpose(img_array[:3], (1, 2, 0))
+            elif img_array.shape[2] == 4:
+                # Shape is (H, W, 4) - extract RGB
+                rgb_array = img_array[:, :, :3]
+            elif img_array.shape[2] == 3:
+                # Shape is (H, W, 3) - already RGB
+                rgb_array = img_array
+            elif img_array.shape[0] == 3:
+                # Shape is (3, H, W) - transpose to RGB
+                rgb_array = np.transpose(img_array, (1, 2, 0))
+    except Exception as e1:
+        try:
+            # Method 2: Try with rasterio
+            with rasterio.open(image_path) as src:
+                if src.count >= 3:
+                    channels = src.read()
+                    if src.count == 4:
+                        rgb_array = np.transpose(channels[:3], (1, 2, 0))
+                    else:
+                        rgb_array = np.transpose(channels, (1, 2, 0))
+        except Exception as e2:
+            try:
+                # Method 3: Fall back to PIL for standard formats
+                img = Image.open(image_path).convert("RGB")
+                rgb_array = np.array(img)
+            except Exception as e3:
+                raise ValueError(f"Could not load image with any method. Errors: tifffile={e1}, rasterio={e2}, PIL={e3}")
+    
+    if rgb_array is None:
+        raise ValueError("Failed to extract RGB data from image")
+    
+    # Ensure RGB is in correct format and range
+    if rgb_array.max() > 1:
+        rgb_array = rgb_array.astype(np.float32) / 255.0
     
     # Predict NDVI from RGB
     ndvi_pred = predict_ndvi_from_rgb(ndvi_model, rgb_array)
     
     # Create temporary 4-channel TIFF file
-    with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as tmp_file:
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.tiff') as tmp_file:
         temp_4ch_path = tmp_file.name
     
     try: