app.py

from fastapi import FastAPI, UploadFile, File, HTTPException
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch
import io

app = FastAPI()

MODEL_NAME = "dima806/skin_types_image_detection"
try:
    processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
    model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
    print(f"Successfully loaded model: {MODEL_NAME}")
except Exception as e:
    print(f"Error loading model: {e}")
    processor = None
    model = None

@app.get("/")
def read_root():
    return {"status": "AuraSkin Skin Type Analyzer is running"}

@app.post("/analyze")
async def analyze_image(image_file: UploadFile = File(...)):
    if not model or not processor:
        raise HTTPException(status_code=500, detail="Model is not available.")
        
    if not image_file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Invalid file type. Please upload an image.")

    try:
        image_bytes = await image_file.read()
        image = Image.open(io.BytesIO(image_bytes))
        
        inputs = processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits

        # --- THIS IS A SINGLE-LABEL MODEL ---
        # We use softmax to get probabilities that add up to 1.
        probabilities = torch.nn.functional.softmax(logits, dim=-1)
        
        # Create a list of all possible labels and their scores
        results = []
        for i, score in enumerate(probabilities[0]):
            label = model.config.id2label[i]
            results.append({"label": label, "score": score.item()})

        # Return the full list, sorted by score from highest to lowest.
        # Your app can then pick the top one as the most likely skin type.
        return sorted(results, key=lambda x: x['score'], reverse=True)

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