from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import torch
import pandas as pd
import uvicorn
import json
from PIL import Image, UnidentifiedImageError
import io
import os
from datetime import datetime
import os
os.environ['TORCH_HOME'] = '/tmp/torch'   # Set writable cache for PyTorch Hub

# Load YOLOv5s pretrained on COCO
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# Folder to store annotated images
SAVE_DIR = "detections"
os.makedirs(SAVE_DIR, exist_ok=True)

app = FastAPI(title="COCO Object Detection API")

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    try:
        # Read uploaded file into bytes
        image_bytes = await file.read()

        # Open with Pillow safely
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    except UnidentifiedImageError:
        return JSONResponse(
            content={"error": "Unrecognized or invalid image file format."},
            status_code=400
        )

    # Run inference
    results = model(image)

    # Save annotated image locally
    plotted_image = results.render()[0]  # numpy array with boxes
    pil_image = Image.fromarray(plotted_image)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    save_path = os.path.join(SAVE_DIR, f"detection_{timestamp}.jpg")
    pil_image.save(save_path, format="JPEG")
    print (f"Saved annotated image to {save_path}")

    # Convert results to JSON
    df: pd.DataFrame = results.pandas().xyxy[0]
    json_data = json.loads(df.to_json(orient="records"))

    return JSONResponse(content=json_data)

@app.get("/")
def root():
    return {"message": "Send POST /predict with an image file to get detections."}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8080)