Create yolo_predictor.py

yolo_predictor.py

@@ -0,0 +1,168 @@
+# yolo_predictor.py
+import os
+import numpy as np
+import rasterio
+from ultralytics import YOLO
+from ndvi_predictor import normalize_rgb, predict_ndvi
+import tempfile
+from rasterio.transform import from_bounds
+from PIL import Image
+
+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):
+    """
+    Predict NDVI channel from RGB array
+    
+    Args:
+        ndvi_model: Loaded NDVI prediction model
+        rgb_array: RGB image as numpy array (H, W, 3)
+    
+    Returns:
+        ndvi_array: Predicted NDVI as numpy array (H, W)
+    """
+    # Normalize RGB input
+    norm_rgb = normalize_rgb(rgb_array)
+    
+    # Predict NDVI
+    ndvi_pred = predict_ndvi(ndvi_model, norm_rgb)
+    
+    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
+    """
+    # Run YOLO prediction
+    results = yolo_model([image_path], conf=conf)
+    
+    return results[0]  # Return first result
+
+def create_4channel_tiff(rgb_array, ndvi_array, output_path):
+    """
+    Create a 4-channel TIFF file from RGB and NDVI arrays
+    
+    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]
+    
+    # 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
+    
+    # Normalize NDVI to match RGB data type 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)
+    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)
+
+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)
+    """
+    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
+
+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
+    
+    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
+    
+    Returns:
+        results: YOLO results object
+    """
+    # Try to load as 4-channel TIFF first, fall back to RGB
+    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)
+    
+    # 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:
+        temp_4ch_path = tmp_file.name
+    
+    try:
+        # Create 4-channel TIFF with predicted NDVI
+        create_4channel_tiff(rgb_array, ndvi_pred, temp_4ch_path)
+        
+        # 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)