Update app.py

app.py

@@ -5,69 +5,87 @@ from PIL import Image
 import torch
 import io
 
-# ★ 토마토 잎 병해 전용 모델
+# 토마토 잎 병해 전용 모델
 MODEL_ID = "wellCh4n/tomato-leaf-disease-classification-resnet50"
 
 app = FastAPI(title="SmartFarm Tomato Disease API")
 
-# --- 모델 & 전처리기 로드 (서버 시작 시 1번만) ---
+# --- 모델 & 전처리기 로딩 (서버 시작 시 1번만) ---
 processor = AutoImageProcessor.from_pretrained(MODEL_ID)
 model = ResNetForImageClassification.from_pretrained(MODEL_ID)
 model.eval()  # 추론 모드
 
-@torch.no_grad()
-def infer_image(img_bytes: bytes):
-    """
-    이미지 바이트를 받아서
-    - 원본
-    - 좌우 반전본
-    두 장을 모델에 넣은 뒤,
-    평균 logits로 softmax → top5 {label, score} 리스트 반환
-    """
-    base_image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+# --- 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",
+}
 
-    # 1) TTA용 이미지들
-    tta_images = [
-        base_image,
-        base_image.transpose(Image.FLIP_LEFT_RIGHT),
-    ]
-
-    logits_sum = None
+# 모델 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()}
 
-    for img in tta_images:
-        inputs = processor(images=img, return_tensors="pt")
-        outputs = model(**inputs)
-        if logits_sum is None:
-            logits_sum = outputs.logits
-        else:
-            logits_sum = logits_sum + outputs.logits
 
-    # 2) 평균 logits에 softmax
-    logits_mean = logits_sum / len(tta_images)
-    probs = torch.nn.functional.softmax(logits_mean, dim=-1)[0]
+@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]
 
-    # 3) top5 뽑기
-    values, indices = probs.topk(5)
+    values, indices = probs.topk(topk)
     values = values.tolist()
     indices = indices.tolist()
 
     id2label = model.config.id2label
-    result = []
+    results = []
     for score, idx in zip(values, indices):
-        label = id2label.get(str(idx), f"Unknown_{idx}")
-        result.append({
+        label = id2label.get(int(idx), f"Unknown_{idx}")
+        results.append({
             "label": label,
-            "score": float(score),  # 0~1 사이 값
+            "score": float(score),
         })
-    return result
+    return results
 
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     """
     PHP에서 보내는 이미지 파일 하나를 받아서
-    top5 결과를 그대로 JSON으로 반환
+    HF Inference API와 비슷한 형식으로 결과 반환:
+    [
+      {"label": "...", "score": 0.87},
+      {"label": "...", "score": 0.05},
+      ...
+    ]
     """
-    img_bytes = await file.read()
-    result = infer_image(img_bytes)
-    return JSONResponse(content=result)
+    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,
+        )