Update yolo_predictor.py

yolo_predictor.py

@@ -32,71 +32,65 @@ def predict_ndvi_from_rgb(ndvi_model, rgb_array):
     return ndvi_pred
 
 def predict_yolo(yolo_model, image_path, conf=0.001):
-    """
-    Predict using YOLO model on 4-channel TIFF image
-    
-    Args:
-        yolo_model: Loaded YOLO model
-        image_path: Path to 4-channel TIFF image
-        conf: Confidence threshold
-    
-    Returns:
-        results: YOLO results object
-    """
-    # Verify the image has 4 channels before prediction
+    import tifffile
+
     try:
-        # Use tifffile for 32-bit TIFF support
         img_array = tifffile.imread(image_path)
-        
-        # Handle different array shapes
+        print(f"[DEBUG] Loaded image shape: {img_array.shape}, dtype: {img_array.dtype}")
+
         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))
+                print(f"[DEBUG] Transposed image shape to (H,W,C): {img_array.shape}")
             elif img_array.shape[2] != 4:
-                raise ValueError(f"YOLO model expects 4-channel images, but got {img_array.shape[2]} channels")
+                raise ValueError(f"[ERROR] Expected 4 channels, got {img_array.shape[2]}")
         else:
-            raise ValueError(f"Unexpected image shape: {img_array.shape}")
-        
-        # Convert 32-bit float to uint8 if necessary
+            raise ValueError(f"[ERROR] Unexpected image shape: {img_array.shape}")
+
+        # Confirm channel count
+        if img_array.shape[2] != 4:
+            raise ValueError(f"[ERROR] After transpose, still not 4 channels: got {img_array.shape[2]}")
+
+        print(f"[DEBUG] Image dtype before normalization: {img_array.dtype}")
+
         if img_array.dtype != np.uint8:
-            # Normalize to [0, 255] for RGB channels (first 3)
+            print(f"[DEBUG] Converting image to uint8")
             rgb_array = img_array[:, :, :3]
+            ndvi_array = img_array[:, :, 3]
+
+            # Normalize RGB
             if rgb_array.max() > 1.0:
                 rgb_array = np.clip(rgb_array / rgb_array.max() * 255, 0, 255).astype(np.uint8)
             else:
                 rgb_array = np.clip(rgb_array * 255, 0, 255).astype(np.uint8)
-            
-            # Normalize NDVI (4th channel) from [-1, 1] to [0, 255]
-            ndvi_array = img_array[:, :, 3]
+
+            # Normalize NDVI
             ndvi_normalized = ((ndvi_array + 1) * 127.5).astype(np.uint8)
-            
-            # Recombine into 4-channel uint8 array
+
             img_array = np.zeros((img_array.shape[0], img_array.shape[1], 4), dtype=np.uint8)
             img_array[:, :, :3] = rgb_array
             img_array[:, :, 3] = ndvi_normalized
-            
-            # Save normalized image to temporary file
+
+            print(f"[DEBUG] Image converted to uint8 with shape: {img_array.shape}")
+
+            # Save normalized version to temp file
             with tempfile.NamedTemporaryFile(delete=False, suffix='.tiff') as tmp_file:
                 temp_path = tmp_file.name
-                tifffile.imwrite(
-                    temp_path,
-                    img_array,
-                    photometric='rgb',
-                    compress='lzw',
-                    metadata={'axes': 'YXC', 'resolution': (1, 1)}  # DPI=1
-                )
+                tifffile.imwrite(temp_path, img_array)
                 image_path = temp_path
-        
-        # Run YOLO prediction
+
+        print(f"[DEBUG] Final image ready for YOLO, path: {image_path}")
+
+        # Final safety check
+        assert img_array.shape[2] == 4, "[FATAL] Final image does not have 4 channels."
+
         results = yolo_model([image_path], conf=conf)
-        
-        # Clean up temporary file if created
+
         if 'temp_path' in locals() and os.path.exists(temp_path):
             os.unlink(temp_path)
-        
-        return results[0]  # Return first result
-    
+
+        return results[0]
+
     except Exception as e:
         raise ValueError(f"Error processing image: {str(e)}")