app.py

from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from transformers import AutoImageProcessor, ResNetForImageClassification
from PIL import Image
import torch
import io

# 토마토 잎 병해 전용 모델
MODEL_ID = "wellCh4n/tomato-leaf-disease-classification-resnet50"

app = FastAPI(title="SmartFarm Tomato Disease API")

# --- 모델 & 전처리기 로딩 (서버 시작 시 1번만) ---
processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = ResNetForImageClassification.from_pretrained(MODEL_ID)
model.eval()  # 추론 모드

# --- id2label 강제 오버라이드 (Unknown_* 방지용) ---
# 모델 config 안의 id2label이 이상하면 우리가 직접 지정한다.
custom_id2label = {
    0: "Tomato_healthy",
    1: "Tomato_Bacterial_spot",
    2: "Tomato_Early_blight",
    3: "Tomato_Late_blight",
    4: "Tomato_Leaf_Mold",
    5: "Tomato_Septoria_leaf_spot",
    6: "Tomato_Spider_mites_Two_spotted_spider_mite",
    7: "Tomato_Target_Spot",
    8: "Tomato_Tomato_Yellow_Leaf_Curl_Virus",
    9: "Tomato_Tomato_mosaic_virus",
}

# 모델 config 에도 반영 (혹시 내부에서 참조할 수도 있으니까)
model.config.id2label = {int(k): v for k, v in custom_id2label.items()}
model.config.label2id = {v: int(k) for k, v in custom_id2label.items()}


@torch.no_grad()
def infer_image(img_bytes: bytes, topk: int = 5):
    """
    이미지 바이트 -> topk [{label, score}, ...] 반환
    score는 0.0~1.0 사이 float
    """
    image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]

    values, indices = probs.topk(topk)
    values = values.tolist()
    indices = indices.tolist()

    id2label = model.config.id2label
    results = []
    for score, idx in zip(values, indices):
        label = id2label.get(int(idx), f"Unknown_{idx}")
        results.append({
            "label": label,
            "score": float(score),
        })
    return results


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    """
    PHP에서 보내는 이미지 파일 하나를 받아서
    HF Inference API와 비슷한 형식으로 결과 반환:
    [
      {"label": "...", "score": 0.87},
      {"label": "...", "score": 0.05},
      ...
    ]
    """
    try:
        img_bytes = await file.read()
        if not img_bytes:
            return JSONResponse(
                {"error": True, "message": "Empty file"},
                status_code=400,
            )

        raw = infer_image(img_bytes, topk=5)
        return JSONResponse(raw, status_code=200)

    except Exception as e:
        # 에러 나면 PHP에서 메시지 확인하기 쉽도록 문자열로 내려줌
        return JSONResponse(
            {"error": True, "message": str(e)},
            status_code=500,
        )