API/api/endpoints.py

from typing import List
from fastapi import APIRouter, File, UploadFile, HTTPException
from ..models.schemas import PredictionResponse, HealthResponse, BatchPredictionResponse
from ..models.inference import model_manager

router = APIRouter()

@router.get("/health", response_model=HealthResponse)
def health_check():
    """
    Check if the API and models are healthy and loaded.
    """
    if not model_manager.is_loaded:
        return HealthResponse(status="degraded", message="Models are not loaded yet.")
    return HealthResponse(status="ok", message="API is healthy and ready to serve predictions.")

@router.post("/predict", response_model=PredictionResponse)
def predict_image(file: UploadFile = File(..., description="An image file strictly required (e.g., JPEG, PNG) to process.")):
    """
    Perform Biomass inference on an uploaded image.
    Uses standard sync function block as PyTorch relies on synchronous execution 
    which FastAPI safely handles in a threadpool to not block the async event loop.
    """
    if not file.content_type or not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Provided file is not an image.")
        
    try:
        image_bytes = file.file.read()
        predictions = model_manager.predict(image_bytes)
        return PredictionResponse(predictions=predictions, message="Success")
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal prediction error: {str(e)}")

@router.post("/predict-batch", response_model=BatchPredictionResponse)
def predict_batch(files: List[UploadFile] = File(..., description="A list of up to 25 image files to process.")):
    """
    Perform Biomass inference on a batch of up to 25 uploaded images.
    """
    if len(files) > 25:
        raise HTTPException(status_code=400, detail="Maximum 25 images allowed per batch.")
        
    # Validate all files are images before processing any
    for f in files:
        if not f.content_type or not f.content_type.startswith("image/"):
            raise HTTPException(status_code=400, detail=f"File '{f.filename}' is not an image.")

    results = {}
    for f in files:
        try:
            image_bytes = f.file.read()
            predictions = model_manager.predict(image_bytes)
            results[f.filename] = predictions
        except ValueError as e:
            raise HTTPException(status_code=400, detail=f"Error processing '{f.filename}': {str(e)}")
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Internal prediction error for '{f.filename}': {str(e)}")
            
    return BatchPredictionResponse(results=results, message="Batch processing successful")