index.py

from fastapi import FastAPI, File, UploadFile, HTTPException
from ultralytics import YOLO
from PIL import Image
import io
from typing import List
from pydantic import BaseModel
import uvicorn

# Define the response model
class Prediction(BaseModel):
    x1: float
    y1: float
    x2: float
    y2: float
    confidence: float
    class_name: str

class PredictionResponse(BaseModel):
    predictions: List[Prediction]

# Initialize FastAPI app and model
app = FastAPI(title="Skin Condition Detection API")
model = YOLO('best.pt')
class_names = ['Acne', 'Dark circles', 'blackheads', 'eczema', 'rosacea', 'whiteheads', 'wrinkles']

@app.post("/classify", response_model=PredictionResponse)
async def classify_image(file: UploadFile = File(...)):
    """
    Endpoint to classify skin conditions in an uploaded image
    """
    # Validate file
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image")
    
    try:
        # Read and process image
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        resized_image = image.copy()
        resized_image.thumbnail((640, 640))
        
        # Get predictions
        results = model(resized_image)[0]
        predictions = []

        if results.boxes is not None:
            boxes = results.boxes.xyxy
            confidences = results.boxes.conf
            classes = results.boxes.cls

            for i in range(len(boxes)):
                box = boxes[i]
                confidence = confidences[i].item()
                class_id = int(classes[i].item())
                
                prediction = Prediction(
                    x1=float(box[0].item()),
                    y1=float(box[1].item()),
                    x2=float(box[2].item()),
                    y2=float(box[3].item()),
                    confidence=confidence,
                    class_name=class_names[class_id]
                )
                predictions.append(prediction)

        return PredictionResponse(predictions=predictions)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# For local development
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)