main.py

"""
FastAPI application for YOLO PyTorch image classification
Accepts image uploads and returns verdict (green/red) based on detected classes
Includes GPU detection and automatic device selection
"""

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict, Optional
import numpy as np
import cv2
from io import BytesIO
from PIL import Image
import logging
import base64

from yolo_inference_pytorch import YOLOInference
from verdict_logic import VerdictEngine

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="YOLO PyTorch Image Classifier - Engineer Selfie",
    description="Upload an image to get classification verdict (green/red) based on detected objects. Supports GPU acceleration.",
    version="2.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class Base64ImageRequest(BaseModel):
    image: str
    
    class Config:
        json_schema_extra = {
            "example": {
                "image": "data:image/jpeg;base64,/9j/4AAQSkZJRg..."
            }
        }

class Detection(BaseModel):
    class_id: int
    class_name: str
    confidence: float
    bbox: List[float]

class VerdictResponse(BaseModel):
    verdict: str
    confidence: float
    detections: List[Detection]
    message: str
    image_size: Dict[str, int]

yolo_model = None
verdict_engine = None

@app.on_event("startup")
async def startup_event():
    """Initialize models on startup"""
    global yolo_model, verdict_engine
    
    try:
        import os
        model_path = "model/best.pt"
        
        if not os.path.exists(model_path):
            logger.warning(f"⚠️  YOLO model not found at {model_path}")
            logger.warning("⚠️  The API will start but /predict endpoint will not work until you add the model")
            logger.warning("⚠️  Please place your best.pt file in the model/ directory")
            
            verdict_engine = VerdictEngine(
                green_classes=[],
                red_classes=[],
                rules_path="config/verdict_rules.json"
            )
            logger.info("Verdict engine initialized successfully")
            return
        
        logger.info("Loading YOLO PyTorch model with GPU detection...")
        yolo_model = YOLOInference(
            model_path=model_path,
            labels_path="model/labels.txt",
            conf_threshold=0.5,
            iou_threshold=0.45,
            device=None
        )
        logger.info("YOLO model loaded successfully")
        logger.info(f"Using confidence threshold: 0.5 (50%)")
        
        logger.info("Initializing verdict engine...")
        verdict_engine = VerdictEngine(
            green_classes=[],
            red_classes=[],
            rules_path="config/verdict_rules.json"
        )
        logger.info("Verdict engine initialized successfully")
        
    except Exception as e:
        logger.error(f"Error during startup: {str(e)}")
        raise

@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "message": "YOLO PyTorch Image Classifier API - Engineer Selfie",
        "version": "2.0.0",
        "description": "Upload images to check for object detection and get approval verdicts. GPU-accelerated when available.",
        "endpoints": {
            "check_image_base64": "/check-image (POST with base64)",
            "check_image_upload": "/check-image/upload (POST with file)",
            "check_images_batch": "/check-images/batch (POST with multiple files)",
            "health": "/health (GET)",
            "model_info": "/model/info (GET)",
            "device_info": "/device/info (GET)"
        },
        "usage": {
            "base64": "POST /check-image with JSON body: {\"image\": \"base64_string\"}",
            "file_upload": "POST /check-image/upload with multipart/form-data",
            "batch": "POST /check-images/batch with multiple files"
        }
    }

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "model_loaded": yolo_model is not None,
        "verdict_engine_loaded": verdict_engine is not None,
        "device": yolo_model.device if yolo_model else "not loaded"
    }

@app.get("/device/info")
async def device_info():
    """Get GPU/device information"""
    if yolo_model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    return yolo_model.get_device_info()

@app.get("/model/info")
async def model_info():
    """Get model information"""
    if yolo_model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    return {
        "model_type": "YOLO PyTorch",
        "model_path": yolo_model.model_path,
        "classes": yolo_model.get_class_names(),
        "num_classes": len(yolo_model.get_class_names()),
        "input_size": yolo_model.get_input_size(),
        "confidence_threshold": yolo_model.conf_threshold,
        "iou_threshold": yolo_model.iou_threshold,
        "device_info": yolo_model.get_device_info()
    }

def process_image_to_numpy(image_data: bytes) -> tuple[np.ndarray, int, int]:
    """Convert image bytes to numpy array for processing"""
    image = Image.open(BytesIO(image_data))
    
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    image_np = np.array(image)
    
    if len(image_np.shape) == 3 and image_np.shape[2] == 3:
        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
    elif len(image_np.shape) == 2:
        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2BGR)
    
    original_height, original_width = image_np.shape[:2]
    return image_np, original_width, original_height


def create_verdict_response(detections: List[Dict], verdict_result: Dict, width: int, height: int) -> VerdictResponse:
    """Create standardized verdict response"""
    formatted_detections = [
        Detection(
            class_id=det["class_id"],
            class_name=det["class_name"],
            confidence=float(det["confidence"]),
            bbox=[float(x) for x in det["bbox"]]
        )
        for det in detections
    ]
    
    return VerdictResponse(
        verdict=verdict_result["verdict"],
        confidence=verdict_result["confidence"],
        detections=formatted_detections,
        message=verdict_result["message"],
        image_size={
            "width": width,
            "height": height
        }
    )


@app.post("/check-image", response_model=VerdictResponse)
async def check_image_base64(request: Base64ImageRequest):
    """
    CHECK IMAGE (Base64 Format)
    
    Send a base64 encoded image and get a verdict (GREEN/RED) based on detected objects.
    """
    if yolo_model is None or verdict_engine is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    try:
        base64_string = request.image
        
        if "base64," in base64_string:
            base64_string = base64_string.split("base64,")[1]
        
        image_bytes = base64.b64decode(base64_string)
        logger.info("Processing base64 encoded image")
        
        image_np, original_width, original_height = process_image_to_numpy(image_bytes)
        logger.info(f"Image size: {original_width}x{original_height}")
        
        detections = yolo_model.predict(image_np)
        logger.info(f"Found {len(detections)} detections")
        
        verdict_result = verdict_engine.get_verdict(detections)
        
        response = create_verdict_response(detections, verdict_result, original_width, original_height)
        
        logger.info(f"Verdict: {verdict_result['verdict']} (confidence: {verdict_result['confidence']:.2f})")
        return response
        
    except base64.binascii.Error:
        raise HTTPException(status_code=400, detail="Invalid base64 image data")
    except Exception as e:
        logger.error(f"Error processing image: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")


@app.post("/check-image/upload", response_model=VerdictResponse)
async def check_image_upload(file: UploadFile = File(...)):
    """
    CHECK IMAGE (File Upload)
    
    Upload an image file and get a verdict (GREEN/RED) based on detected objects.
    """
    if yolo_model is None or verdict_engine is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    if not file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=400,
            detail=f"Invalid file type: {file.content_type}. Please upload an image."
        )
    
    try:
        logger.info(f"Processing image: {file.filename}")
        contents = await file.read()
        
        image_np, original_width, original_height = process_image_to_numpy(contents)
        logger.info(f"Image size: {original_width}x{original_height}")
        
        detections = yolo_model.predict(image_np)
        logger.info(f"Found {len(detections)} detections")
        
        verdict_result = verdict_engine.get_verdict(detections)
        
        response = create_verdict_response(detections, verdict_result, original_width, original_height)
        
        logger.info(f"Verdict: {verdict_result['verdict']} (confidence: {verdict_result['confidence']:.2f})")
        return response
        
    except Exception as e:
        logger.error(f"Error processing image: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")


@app.post("/predict", response_model=VerdictResponse)
async def predict(file: UploadFile = File(...)):
    """
    [LEGACY] Upload an image and get classification verdict
    
    Note: Use /check-image/upload instead for clearer API naming
    """
    return await check_image_upload(file)

@app.post("/check-images/batch")
async def check_images_batch(files: List[UploadFile] = File(...)):
    """
    CHECK MULTIPLE IMAGES (Batch Upload)
    
    Upload multiple image files (max 10) and get verdicts for each.
    """
    if len(files) > 10:
        raise HTTPException(
            status_code=400,
            detail="Maximum 10 images allowed per batch"
        )
    
    results = []
    for file in files:
        try:
            result = await check_image_upload(file)
            results.append({
                "filename": file.filename,
                "result": result
            })
        except Exception as e:
            results.append({
                "filename": file.filename,
                "error": str(e)
            })
    
    return {"results": results}


@app.post("/predict/batch")
async def predict_batch(files: List[UploadFile] = File(...)):
    """
    [LEGACY] Upload multiple images and get classification verdicts
    
    Note: Use /check-images/batch instead for clearer API naming
    """
    return await check_images_batch(files)

if __name__ == "__main__":
    import uvicorn
    import sys
    
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8000
    
    print(f"\n{'='*70}")
    print(f"🚀 Starting Engineer Selfie API on http://0.0.0.0:{port}")
    print(f"📖 API docs available at http://localhost:{port}/docs")
    print(f"🖥️  GPU acceleration will be used if available")
    print(f"{'='*70}\n")
    
    uvicorn.run(app, host="0.0.0.0", port=port)