Update yolo_predictor.py

yolo_predictor.py

@@ -44,6 +44,36 @@ def predict_yolo(yolo_model, image_path, conf=0.001):
     Returns:
         results: YOLO results object
     """
+    # Verify the image has 4 channels before prediction
+    try:
+        # Check image format and channels
+        with Image.open(image_path) as img:
+            if hasattr(img, 'n_frames'):
+                # Multi-frame TIFF
+                channels = img.n_frames
+            else:
+                # Regular image
+                channels = len(img.getbands()) if hasattr(img, 'getbands') else 3
+        
+        # If not 4 channels, try with tifffile
+        if channels != 4:
+            img_array = tifffile.imread(image_path)
+            if len(img_array.shape) == 3:
+                if img_array.shape[0] == 4:
+                    channels = 4
+                elif img_array.shape[2] == 4:
+                    channels = 4
+                else:
+                    channels = img_array.shape[0] if img_array.shape[0] <= 4 else img_array.shape[2]
+            else:
+                channels = 1
+        
+        if channels != 4:
+            raise ValueError(f"YOLO model expects 4-channel images, but got {channels} channels")
+            
+    except Exception as e:
+        raise ValueError(f"Error reading image channels: {str(e)}")
+    
     # Run YOLO prediction
     results = yolo_model([image_path], conf=conf)
     
@@ -51,7 +81,7 @@ def predict_yolo(yolo_model, image_path, conf=0.001):
 
 def create_4channel_tiff(rgb_array, ndvi_array, output_path):
     """
-    Create a 4-channel TIFF file from RGB and NDVI arrays
+    Create a 4-channel TIFF file from RGB and NDVI arrays compatible with PIL and YOLO
     
     Args:
         rgb_array: RGB image as numpy array (H, W, 3)
@@ -60,24 +90,33 @@ 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((4, height, width), dtype=np.float32)
-    
-    # Convert RGB to proper format and range
-    if rgb_array.dtype == np.uint8:
-        rgb_normalized = rgb_array.astype(np.float32) / 255.0
+    # Ensure RGB is in uint8 format for better compatibility
+    if rgb_array.dtype != np.uint8:
+        if rgb_array.max() <= 1.0:
+            rgb_normalized = (rgb_array * 255).astype(np.uint8)
+        else:
+            rgb_normalized = np.clip(rgb_array, 0, 255).astype(np.uint8)
     else:
-        rgb_normalized = rgb_array.astype(np.float32)
+        rgb_normalized = rgb_array
     
-    # 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
+    # Convert NDVI from [-1, 1] to [0, 255] for uint8 storage
+    ndvi_normalized = ((ndvi_array + 1) * 127.5).astype(np.uint8)
     
-    # Use tifffile for better compatibility with YOLO
-    import tifffile
-    tifffile.imwrite(output_path, four_channel, photometric='rgb')
+    # Create 4-channel array in (H, W, 4) format
+    four_channel = np.zeros((height, width, 4), dtype=np.uint8)
+    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_normalized          # NDVI
+    
+    # Save using tifffile with proper format for YOLO compatibility
+    tifffile.imwrite(
+        output_path, 
+        four_channel, 
+        photometric='rgb',
+        compress='lzw',
+        metadata={'axes': 'YXC'}
+    )
 
 def load_4channel_tiff(image_path):
     """
@@ -91,45 +130,52 @@ def load_4channel_tiff(image_path):
         ndvi_array: NDVI channel as numpy array (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
+        # Try with tifffile first for better TIFF support
         img_array = tifffile.imread(image_path)
         
-        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 len(img_array.shape) == 3:
+            if img_array.shape[0] == 4:
+                # Shape is (4, H, W) - transpose to (H, W, 4)
+                img_array = np.transpose(img_array, (1, 2, 0))
+            elif img_array.shape[2] != 4:
+                raise ValueError(f"Expected 4 channels, got {img_array.shape}")
+        
+        # Extract RGB and NDVI from (H, W, 4) format
+        rgb_array = img_array[:, :, :3]
+        ndvi_array = img_array[:, :, 3]
         
-        # Normalize NDVI if needed
+        # Convert NDVI back from [0, 255] to [-1, 1] if it was stored as uint8
         if img_array.dtype == np.uint8:
             ndvi_array = (ndvi_array.astype(np.float32) / 127.5) - 1
         
         return rgb_array, ndvi_array
+        
+    except Exception as e:
+        # Fallback to rasterio
+        try:
+            with rasterio.open(image_path) as src:
+                if src.count != 4:
+                    raise ValueError(f"Expected 4 channels, got {src.count}")
+                
+                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)
+                
+                # Convert NDVI if needed
+                if channels.dtype == np.uint8:
+                    ndvi_array = (ndvi_array.astype(np.float32) / 127.5) - 1
+                
+                return rgb_array, ndvi_array
+                
+        except Exception as e2:
+            raise ValueError(f"Could not load 4-channel TIFF. Errors: tifffile={e}, rasterio={e2}")
 
 def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
     """
-    Full pipeline: Load 4-channel image -> Extract RGB -> Predict NDVI -> 
-    Create new 4-channel with predicted NDVI -> Run YOLO prediction
+    Full pipeline: Load image -> Extract RGB -> Predict NDVI -> 
+    Create 4-channel TIFF -> Run YOLO prediction
     
     Args:
         ndvi_model: Loaded NDVI prediction model
@@ -142,40 +188,48 @@ def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
     """
     rgb_array = None
     
-    # Try multiple methods to load the image
+    # Try multiple methods to load the image and extract RGB
     try:
         # 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[0] == 3:
+                # Shape is (3, H, W) - transpose to RGB
+                rgb_array = np.transpose(img_array, (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))
+        elif len(img_array.shape) == 2:
+            # Grayscale - convert to RGB
+            rgb_array = np.stack([img_array] * 3, axis=-1)
+            
     except Exception as e1:
         try:
             # Method 2: Try with rasterio
             with rasterio.open(image_path) as src:
+                channels = src.read()
                 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))
+                    rgb_array = np.transpose(channels[:3], (1, 2, 0))
+                elif src.count == 1:
+                    # Single channel - convert to RGB
+                    single_channel = channels[0]
+                    rgb_array = np.stack([single_channel] * 3, axis=-1)
+                    
         except Exception as e2:
             try:
                 # Method 3: Fall back to PIL for standard formats
-                img = Image.open(image_path).convert("RGB")
+                img = Image.open(image_path)
+                if img.mode != 'RGB':
+                    img = img.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}")
     
@@ -183,11 +237,19 @@ def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
         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
+    if rgb_array.dtype == np.uint8:
+        # Keep as uint8 but also create float version for NDVI prediction
+        rgb_float = rgb_array.astype(np.float32) / 255.0
+    else:
+        # Already float, ensure range is [0, 1]
+        if rgb_array.max() > 1.0:
+            rgb_float = rgb_array / 255.0
+        else:
+            rgb_float = rgb_array
+        rgb_array = (rgb_float * 255).astype(np.uint8)
     
     # Predict NDVI from RGB
-    ndvi_pred = predict_ndvi_from_rgb(ndvi_model, rgb_array)
+    ndvi_pred = predict_ndvi_from_rgb(ndvi_model, rgb_float)
     
     # Create temporary 4-channel TIFF file
     with tempfile.NamedTemporaryFile(delete=False, suffix='.tiff') as tmp_file:
@@ -197,6 +259,21 @@ def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
         # Create 4-channel TIFF with predicted NDVI
         create_4channel_tiff(rgb_array, ndvi_pred, temp_4ch_path)
         
+        # Verify the created file can be read
+        try:
+            test_img = Image.open(temp_4ch_path)
+            if hasattr(test_img, 'n_frames'):
+                channels = test_img.n_frames
+            else:
+                channels = len(test_img.getbands())
+            test_img.close()
+            
+            if channels != 4:
+                raise ValueError(f"Created TIFF has {channels} channels instead of 4")
+                
+        except Exception as e:
+            raise ValueError(f"Created TIFF file is not readable: {str(e)}")
+        
         # Run YOLO prediction on 4-channel image
         results = predict_yolo(yolo_model, temp_4ch_path, conf=conf)