New endpoints

app.py

@@ -231,4 +231,108 @@ async def predict_pipeline_api(file: UploadFile = File(...)):
         
     except Exception as e:
         logger.error(f"Error in predict_pipeline_api: {e}")
+        return JSONResponse(status_code=500, content={"error": str(e)})
+    
+# New endpoints to add to your FastAPI app
+from yolo_predictor import predict_yolo_with_image, predict_pipeline_with_image, pil_image_to_bytes
+
+@app.post("/predict_yolo_image/")
+async def predict_yolo_image_api(file: UploadFile = File(...)):
+    """Predict YOLO results from 4-channel TIFF image and return annotated image"""
+    if yolo_model is None:
+        return JSONResponse(status_code=500, content={"error": "YOLO model not loaded"})
+    
+    try:
+        # Save uploaded file temporarily with proper extension
+        file_extension = '.tiff' if file.filename and file.filename.lower().endswith(('.tif', '.tiff')) else '.tiff'
+        
+        with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp_file:
+            contents = await file.read()
+            tmp_file.write(contents)
+            tmp_file.flush()  # Ensure data is written
+            tmp_file_path = tmp_file.name
+        
+        try:
+            # Verify the file was written correctly
+            if not os.path.exists(tmp_file_path) or os.path.getsize(tmp_file_path) == 0:
+                raise ValueError("Failed to create temporary file")
+            
+            logger.info(f"Processing YOLO prediction with image output for file: {file.filename}, temp path: {tmp_file_path}")
+            
+            # Additional validation: check if file has 4 channels
+            try:
+                import tifffile
+                test_array = tifffile.imread(tmp_file_path)
+                if len(test_array.shape) == 3:
+                    if test_array.shape[0] == 4 or test_array.shape[2] == 4:
+                        channels = 4
+                    else:
+                        channels = test_array.shape[0] if test_array.shape[0] <= 4 else test_array.shape[2]
+                else:
+                    channels = 1
+                
+                if channels != 4:
+                    raise ValueError(f"YOLO model expects 4-channel images, but uploaded file has {channels} channels")
+                    
+            except Exception as validation_error:
+                logger.warning(f"Could not validate channels: {validation_error}")
+            
+            # Predict using YOLO model and get annotated image
+            annotated_image = predict_yolo_with_image(yolo_model, tmp_file_path)
+            
+            # Convert PIL Image to bytes for response
+            img_bytes = pil_image_to_bytes(annotated_image, format='PNG')
+            
+            logger.info(f"YOLO prediction with image output completed successfully")
+            
+            return StreamingResponse(
+                img_bytes,
+                media_type="image/png",
+                headers={"Content-Disposition": f"attachment; filename=yolo_annotated_{file.filename}.png"}
+            )
+            
+        finally:
+            # Clean up temporary file
+            if os.path.exists(tmp_file_path):
+                os.unlink(tmp_file_path)
+            
+    except Exception as e:
+        logger.error(f"Error in predict_yolo_image_api: {e}")
+        return JSONResponse(status_code=500, content={"error": str(e)})
+
+@app.post("/predict_pipeline_image/")
+async def predict_pipeline_image_api(file: UploadFile = File(...)):
+    """Full pipeline with image output: RGB -> NDVI -> 32-bit 4-channel TIFF (RGB+NDVI) -> YOLO prediction -> Annotated Image"""
+    if ndvi_model is None or yolo_model is None:
+        return JSONResponse(status_code=500, content={"error": "Models not loaded properly"})
+    
+    try:
+        logger.info(f"Starting full pipeline with image output for file: {file.filename}")
+        
+        # Read uploaded RGB image
+        contents = await file.read()
+        logger.info(f"Read {len(contents)} bytes from uploaded file")
+        
+        # Convert to PIL Image and then to numpy array
+        img = Image.open(BytesIO(contents)).convert("RGB")
+        rgb_array = np.array(img)
+        logger.info(f"Converted to RGB array with shape: {rgb_array.shape}")
+        
+        # Run the full pipeline with image output (includes resizing internally)
+        annotated_image = predict_pipeline_with_image(ndvi_model, yolo_model, rgb_array)
+        logger.info("Pipeline processing with image output completed successfully")
+        
+        # Convert PIL Image to bytes for response
+        img_bytes = pil_image_to_bytes(annotated_image, format='PNG')
+        
+        logger.info(f"Pipeline prediction with image output completed successfully")
+        
+        return StreamingResponse(
+            img_bytes,
+            media_type="image/png",
+            headers={"Content-Disposition": f"attachment; filename=pipeline_annotated_{file.filename}.png"}
+        )
+        
+    except Exception as e:
+        logger.error(f"Error in predict_pipeline_image_api: {e}")
         return JSONResponse(status_code=500, content={"error": str(e)})

yolo_predictor.py

@@ -4,6 +4,9 @@ import logging
 import tempfile
 import numpy as np
 import tifffile
+from io import BytesIO
+import cv2
+from PIL import Image
 from rasterio.transform import from_bounds
 from ultralytics import YOLO
 from ndvi_predictor import normalize_rgb, predict_ndvi
@@ -18,7 +21,7 @@ def load_yolo_model(model_path):
     logger.info(f"Loading YOLO model from: {model_path}")
     return YOLO(model_path)
 
-def predict_yolo(yolo_model, image_path, conf=0.001):
+def predict_yolo(yolo_model, image_path, conf=0.25):
     """
     Predict using YOLO model on 4-channel TIFF image
     
@@ -104,7 +107,7 @@ def create_4channel_tiff(rgb_array, ndvi_array, output_path):
     tifffile.imwrite(output_path, four_channel)
     logger.info(f"Successfully saved 4-channel TIFF (RGB+NDVI format) to: {output_path}")
 
-def predict_pipeline(ndvi_model, yolo_model, rgb_array, conf=0.001):
+def predict_pipeline(ndvi_model, yolo_model, rgb_array, conf=0.25):
     """
     Full pipeline: RGB -> NDVI -> 32-bit 4-channel TIFF (RGB+NDVI) -> YOLO prediction
     
@@ -196,4 +199,162 @@ def validate_4channel_tiff(tiff_path):
         
     except Exception as e:
         logger.error(f"Error validating TIFF file: {e}")
-        return False
+        return False
+
+# Additional functions for yolo_predictor.py
+
+def predict_yolo_with_image(yolo_model, image_path, conf=0.25, save_path=None):
+    """
+    Predict using YOLO model on 4-channel TIFF image and return annotated image
+    
+    Args:
+        yolo_model: Loaded YOLO model
+        image_path: Path to 4-channel TIFF image
+        conf: Confidence threshold
+        save_path: Optional path to save the annotated image
+    
+    Returns:
+        annotated_image: PIL Image object with annotations
+    """
+    logger.info(f"Starting YOLO prediction with image output on: {image_path} with confidence: {conf}")
+    
+    # Verify file exists and has correct format
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+    
+    try:
+        # Quick validation of the TIFF file
+        test_array = tifffile.imread(image_path)
+        logger.info(f"TIFF file shape: {test_array.shape}, dtype: {test_array.dtype}")
+        
+        # Validate channels
+        if len(test_array.shape) == 3:
+            channels = test_array.shape[0] if test_array.shape[0] <= 4 else test_array.shape[2]
+        else:
+            channels = 1
+        
+        if channels != 4:
+            raise ValueError(f"Expected 4-channel image, got {channels} channels")
+            
+    except Exception as e:
+        logger.error(f"Error validating TIFF file: {e}")
+        raise
+    
+    logger.info("Running YOLO model inference with image output...")
+    
+    # Run YOLO prediction directly on the input file
+    results = yolo_model([image_path], conf=conf)
+    result = results[0]
+    
+    # Create temporary file for saving annotated image
+    if save_path is None:
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp_file:
+            save_path = tmp_file.name
+    
+    try:
+        # Save the annotated image using ultralytics built-in method
+        result.save(save_path)
+        logger.info(f"Annotated image saved to: {save_path}")
+        
+        # Load the saved image and convert to PIL Image
+        annotated_image = Image.open(save_path).convert('RGB')
+        logger.info(f"YOLO prediction with image output completed successfully")
+        
+        return annotated_image
+        
+    except Exception as e:
+        logger.error(f"Error saving annotated image: {e}")
+        raise
+    finally:
+        # Clean up temporary file if we created it
+        if save_path.endswith('.png') and os.path.exists(save_path):
+            try:
+                os.unlink(save_path)
+                logger.info(f"Cleaned up temporary annotated image file: {save_path}")
+            except Exception as cleanup_error:
+                logger.warning(f"Failed to clean up temporary file: {cleanup_error}")
+
+def predict_pipeline_with_image(ndvi_model, yolo_model, rgb_array, conf=0.25):
+    """
+    Full pipeline with image output: RGB -> NDVI -> 32-bit 4-channel TIFF (RGB+NDVI) -> YOLO prediction -> Annotated Image
+    
+    Args:
+        ndvi_model: Loaded NDVI prediction model
+        yolo_model: Loaded YOLO model
+        rgb_array: RGB image as numpy array (H, W, 3)
+        conf: Confidence threshold for YOLO
+    
+    Returns:
+        annotated_image: PIL Image object with YOLO annotations
+    """
+    logger.info("Starting full prediction pipeline with image output")
+    logger.info(f"Input RGB array shape: {rgb_array.shape}, dtype: {rgb_array.dtype}")
+    
+    # Step 1: Resize RGB image to target size
+    logger.info("Step 1: Resizing RGB image to target size")
+    target_size = (640, 640)  # (height, width)
+    rgb_resized = resize_image_optimized(rgb_array, target_size)
+    logger.info(f"Resized RGB shape: {rgb_resized.shape}")
+    
+    # Step 2: Normalize RGB image
+    logger.info("Step 2: Normalizing RGB image")
+    normalized_rgb = normalize_rgb(rgb_resized)
+    logger.info(f"Normalized RGB shape: {normalized_rgb.shape}, range: [{normalized_rgb.min():.3f}, {normalized_rgb.max():.3f}]")
+    
+    # Step 3: Predict NDVI
+    logger.info("Step 3: Predicting NDVI from RGB")
+    ndvi_prediction = predict_ndvi(ndvi_model, normalized_rgb)
+    logger.info(f"NDVI prediction shape: {ndvi_prediction.shape}, range: [{ndvi_prediction.min():.3f}, {ndvi_prediction.max():.3f}]")
+    
+    # Step 4: Create 4-channel TIFF file
+    logger.info("Step 4: Creating 4-channel TIFF file (RGB+NDVI)")
+    
+    # Create temporary file for the 4-channel TIFF
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.tiff') as tmp_file:
+        tiff_path = tmp_file.name
+    
+    try:
+        # Create the 4-channel TIFF using resized RGB and predicted NDVI
+        create_4channel_tiff(rgb_resized, ndvi_prediction, tiff_path)
+        
+        # Verify the created file
+        if not os.path.exists(tiff_path):
+            raise FileNotFoundError(f"Failed to create 4-channel TIFF at: {tiff_path}")
+        
+        file_size = os.path.getsize(tiff_path)
+        logger.info(f"Created 4-channel TIFF file size: {file_size} bytes")
+        
+        # Step 5: Run YOLO prediction on the 4-channel TIFF and get annotated image
+        logger.info("Step 5: Running YOLO prediction on 4-channel TIFF with image output")
+        annotated_image = predict_yolo_with_image(yolo_model, tiff_path, conf=conf)
+        
+        logger.info("Full pipeline with image output completed successfully")
+        return annotated_image
+        
+    except Exception as e:
+        logger.error(f"Error in pipeline with image output: {e}")
+        raise
+    finally:
+        # Clean up temporary file
+        if os.path.exists(tiff_path):
+            try:
+                os.unlink(tiff_path)
+                logger.info(f"Cleaned up temporary file: {tiff_path}")
+            except Exception as cleanup_error:
+                logger.warning(f"Failed to clean up temporary file: {cleanup_error}")
+
+def pil_image_to_bytes(image, format='PNG'):
+    """
+    Convert PIL Image to bytes for API response
+    
+    Args:
+        image: PIL Image object
+        format: Image format ('PNG', 'JPEG', etc.)
+    
+    Returns:
+        BytesIO: Image as bytes buffer
+    """
+    img_bytes = BytesIO()
+    image.save(img_bytes, format=format)
+    img_bytes.seek(0)
+    return img_bytes