Update yolo_predictor.py

yolo_predictor.py

@@ -1,297 +1,100 @@
 # yolo_predictor.py
 import os
-import numpy as np
 import rasterio
 from modified_ultralytics import YOLO
-from ndvi_predictor import normalize_rgb, predict_ndvi
-import tempfile
-from PIL import Image
 import tifffile
 
 def load_yolo_model(model_path):
     """Load YOLO model from .pt file"""
     return YOLO(model_path)
 
-def predict_ndvi_from_rgb(ndvi_model, rgb_array):
+def validate_4channel_tiff(image_path):
     """
-    Predict NDVI channel from RGB array
+    Validate that the input TIFF file has 4 channels and is readable
     
     Args:
-        ndvi_model: Loaded NDVI prediction model
-        rgb_array: RGB image as numpy array (H, W, 3)
+        image_path: Path to input TIFF image
     
     Returns:
-        ndvi_array: Predicted NDVI as numpy array (H, W)
+        bool: True if valid 4-channel TIFF
+        
+    Raises:
+        ValueError: If validation fails
     """
-    # Normalize RGB input
-    norm_rgb = normalize_rgb(rgb_array)
-    
-    # Predict NDVI
-    ndvi_pred = predict_ndvi(ndvi_model, norm_rgb)
+    if not os.path.exists(image_path):
+        raise ValueError(f"Image file does not exist: {image_path}")
     
-    return ndvi_pred
-
-def predict_yolo(yolo_model, image_path, conf=0.001):
-    import tifffile
-
     try:
+        # Primary validation with tifffile
         img_array = tifffile.imread(image_path)
-        print(f"[DEBUG] Loaded image shape: {img_array.shape}, dtype: {img_array.dtype}")
-
+        
+        # Check array shape and channels
         if len(img_array.shape) == 3:
             if img_array.shape[0] == 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"[ERROR] Expected 4 channels, got {img_array.shape[2]}")
-        else:
-            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:
-            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)
+                # Shape is (4, H, W)
+                channels = 4
+                height, width = img_array.shape[1], img_array.shape[2]
+            elif img_array.shape[2] == 4:
+                # Shape is (H, W, 4)
+                channels = 4
+                height, width = img_array.shape[0], img_array.shape[1]
             else:
-                rgb_array = np.clip(rgb_array * 255, 0, 255).astype(np.uint8)
-
-            # Normalize NDVI
-            ndvi_normalized = ((ndvi_array + 1) * 127.5).astype(np.uint8)
-
-            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
-
-            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)
-                image_path = temp_path
-
-        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)
-
-        if 'temp_path' in locals() and os.path.exists(temp_path):
-            os.unlink(temp_path)
-
-        return results[0]
-
-    except Exception as e:
-        raise ValueError(f"Error processing image: {str(e)}")
-
-def create_4channel_tiff(rgb_array, ndvi_array, output_path):
-    """
-    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)
-        ndvi_array: NDVI image as numpy array (H, W)
-        output_path: Path to save the 4-channel TIFF
-    """
-    height, width = rgb_array.shape[:2]
-    
-    # Ensure RGB is in uint8 format
-    if rgb_array.dtype != np.uint8:
-        if rgb_array.max() <= 1.0:
-            rgb_normalized = (rgb_array * 255).astype(np.uint8)
+                channels = min(img_array.shape[0], img_array.shape[2])
+                height, width = img_array.shape[0], img_array.shape[1]
         else:
-            rgb_normalized = np.clip(rgb_array, 0, 255).astype(np.uint8)
-    else:
-        rgb_normalized = rgb_array
-    
-    # Convert NDVI from [-1, 1] to [0, 255] for uint8 storage
-    ndvi_normalized = ((ndvi_array + 1) * 127.5).astype(np.uint8)
-    
-    # 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 explicit 32-bit compatibility and DPI=1
-    tifffile.imwrite(
-        output_path,
-        four_channel,
-        photometric='rgb',
-        compress='lzw',
-        metadata={'axes': 'YXC', 'resolution': (1, 1)},  # DPI=1
-        bitspersample=8  # Explicitly set to 8-bit per channel
-    )
-
-def load_4channel_tiff(image_path):
-    """
-    Load a 4-channel TIFF image
-    
-    Args:
-        image_path: Path to 4-channel TIFF image
-    
-    Returns:
-        rgb_array: RGB channels as numpy array (H, W, 3)
-        ndvi_array: NDVI channel as numpy array (H, W)
-    """
-    try:
-        # Try with tifffile first for better TIFF support
-        img_array = tifffile.imread(image_path)
-        
-        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]
-        
-        # 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
+            raise ValueError(f"Invalid image shape: {img_array.shape}. Expected 3D array with 4 channels.")
+            
+        if channels != 4:
+            raise ValueError(f"YOLO model expects 4-channel images, but got {channels} channels")
+            
+        print(f"Validation successful: {channels} channels, {height}x{width}, dtype: {img_array.dtype}")
+        return True
         
     except Exception as e:
-        # Fallback to rasterio
+        # Fallback validation with 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
+                    raise ValueError(f"YOLO model expects 4-channel images, but got {src.count} channels")
+                    
+                print(f"Validation successful (rasterio): {src.count} channels, {src.width}x{src.height}, dtype: {src.dtypes[0]}")
+                return True
                 
         except Exception as e2:
-            raise ValueError(f"Could not load 4-channel TIFF. Errors: tifffile={e}, rasterio={e2}")
+            raise ValueError(f"Could not validate TIFF file. Errors: tifffile={str(e)}, rasterio={str(e2)}")
 
-def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
+def predict_yolo(yolo_model, image_path, conf=0.001):
     """
-    Full pipeline: Load image -> Extract RGB -> Predict NDVI -> 
-    Create 4-channel TIFF -> Run YOLO prediction
+    Predict using YOLO model on 4-channel TIFF image
     
     Args:
-        ndvi_model: Loaded NDVI prediction model
         yolo_model: Loaded YOLO model
-        image_path: Path to input image (can be RGB or 4-channel TIFF)
-        conf: Confidence threshold for YOLO
+        image_path: Path to 4-channel TIFF image
+        conf: Confidence threshold
     
     Returns:
         results: YOLO results object
     """
-    rgb_array = None
+    # Validate input file
+    validate_4channel_tiff(image_path)
     
-    # Try multiple methods to load the image and extract RGB
-    try:
-        # Method 1: Try with tifffile first (best for complex TIFF files)
-        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 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:
-                    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.PIL(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}")
-    
-    if rgb_array is None:
-        raise ValueError("Failed to extract RGB data from image")
+    # Run YOLO prediction directly on the input file
+    results = yolo_model([image_path], conf=conf)
     
-    # Ensure RGB is in correct format and range
-    if rgb_array.dtype == np.uint8:
-        rgb_float = rgb_array.astype(np.float32) / 255.0
-    else:
-        if rgb_array.max() > 1.0:
-            rgb_float = rgb_array / rgb_array.max()
-        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_float)
+    return results[0]  # Return first result
+
+def predict_pipeline(ndvi_model, yolo_model, image_path, conf=0.001):
+    """
+    Simplified pipeline: Validate input -> Run YOLO prediction
     
-    # Create temporary 4-channel TIFF file
-    with tempfile.NamedTemporaryFile(delete=False, suffix='.tiff') as tmp_file:
-        temp_4ch_path = tmp_file.name
+    Args:
+        ndvi_model: Not used (kept for API compatibility)
+        yolo_model: Loaded YOLO model
+        image_path: Path to input 4-channel TIFF image
+        conf: Confidence threshold for YOLO
     
-    try:
-        # 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_array = tifffile.imread(temp_4ch_path)
-            if len(test_array.shape) == 3 and (test_array.shape[0] == 4 or test_array.shape[2] == 4):
-                channels = 4
-            else:
-                channels = test_array.shape[2] if len(test_array.shape) == 3 else 1
-            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)
-        
-        return results
-        
-    finally:
-        # Clean up temporary file
-        if os.path.exists(temp_4ch_path):
-            os.unlink(temp_4ch_path)
+    Returns:
+        results: YOLO results object
+    """
+    # Simply validate and run prediction on the uploaded file
+    return predict_yolo(yolo_model, image_path, conf=conf)