Update app.py

app.py

@@ -1,12 +1,19 @@
+# app.py
 import os, csv, random
 from datetime import datetime
 import numpy as np
 from PIL import Image
 import gradio as gr
 
+# ---- PyTorch / Torchvision ----
+import torch, torch.nn as nn
+from torchvision import transforms, models
+
 # ---------- Paths ----------
 IMG_DIR = "durian_images"
 HIS_CSV = "history/history.csv"
+CKPT_PATH = "durian_mnv2_ckpt.pth"  # วางไฟล์น้ำหนักไว้โฟลเดอร์เดียวกับ app.py
+
 os.makedirs(IMG_DIR, exist_ok=True)
 os.makedirs(os.path.dirname(HIS_CSV), exist_ok=True)
 
@@ -118,14 +125,10 @@ RIPENESS_CAPTION_VARIANTS = {
     ],
 }
 
-# ---------- ความมั่นใจ ----------
+# ---------- ความมั่นใจ (ไม่ปั๊มตัวเลข) ----------
 def adjust_confidence(raw_conf: float) -> float:
     p = raw_conf if raw_conf > 1 else raw_conf * 100.0
-    if p < 30.0:
-        return 96.0
-    elif p < 60.0:
-        return 85.0
-    return round(p, 1)
+    return round(max(0.0, min(100.0, p)), 1)
 
 # ---------- เลือกแคปชั่น ----------
 def pick_variant(level: str):
@@ -138,7 +141,7 @@ def generate_caption(level: str, raw_conf_pct: float) -> str:
     body_list = pick_variant(level)
     return " ".join([head] + body_list)
 
-# ---------- Fallback inference ----------
+# ---------- Fallback 4-class by color ----------
 def _classify_4class_by_color(img: Image.Image):
     arr = np.array(img.convert("RGB")).reshape(-1, 3).mean(axis=0)
     R, G, B = arr
@@ -148,10 +151,97 @@ def _classify_4class_by_color(img: Image.Image):
     idx = int(np.argmax(probs))
     return idx, probs
 
+# ---------- Model (PyTorch) ----------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+idx_to_class = {i: name for i, name in enumerate(RIPENESS_LABELS)}
+model = None
+temperature = 1.0
+
+def _build_model(num_classes=4):
+    m = models.mobilenet_v2(weights=None)
+    in_f = m.classifier[1].in_features
+    m.classifier[1] = nn.Linear(in_f, num_classes)
+    return m
+
+def _load_model():
+    """โหลดโมเดล + mapping จาก ckpt ถ้ามี"""
+    global model, idx_to_class, temperature
+    if not os.path.exists(CKPT_PATH):
+        print(f"[WARN] ckpt not found at {CKPT_PATH}. Use color fallback.")
+        return False
+
+    ckpt = torch.load(CKPT_PATH, map_location=device)
+
+    state = None
+    if isinstance(ckpt, dict):
+        for k in ["state_dict", "model_state_dict", "model"]:
+            if k in ckpt and isinstance(ckpt[k], dict):
+                state = ckpt[k]; break
+        if state is None:
+            # บางครั้งเซฟเป็น state_dict ตรง ๆ
+            # ตรวจคร่าว ๆ ว่า value เป็น tensor
+            if all(hasattr(v, "shape") for v in ckpt.values()):
+                state = ckpt
+    else:
+        # ผิดรูปแบบ
+        print("[WARN] Unknown ckpt format, using fallback.")
+        return False
+
+    model_ = _build_model(num_classes=len(RIPENESS_LABELS))
+    if state is not None:
+        state = {k.replace("module.", ""): v for k, v in state.items()}
+        missing, unexpected = model_.load_state_dict(state, strict=False)
+        if missing:
+            print(f"[INFO] missing keys: {missing[:5]}{'...' if len(missing)>5 else ''}")
+        if unexpected:
+            print(f"[INFO] unexpected keys: {unexpected[:5]}{'...' if len(unexpected)>5 else ''}")
+    else:
+        return False
+
+    if "class_to_idx" in ckpt and isinstance(ckpt["class_to_idx"], dict):
+        c2i = ckpt["class_to_idx"]
+        idx_to_class = {i: lbl for lbl, i in c2i.items()}
+
+    temperature = float(ckpt.get("temperature", 1.0))
+    model_.to(device).eval()
+    print(f"[OK] Model loaded. Temperature={temperature}. Classes={list(idx_to_class.values())}")
+    # set global
+    globals()["model"] = model_
+    globals()["idx_to_class"] = idx_to_class
+    globals()["temperature"] = temperature
+    return True
+
+MODEL_READY = _load_model()
+
+# Preprocess (ปรับให้ตรงกับตอนเทรน)
+IM_SIZE = 224
+_base_tf = transforms.Compose([
+    transforms.Resize((IM_SIZE, IM_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485,0.456,0.406],
+                         std=[0.229,0.224,0.225]),
+])
+
+def _predict_proba_with_model(img: Image.Image):
+    """TTA เบา ๆ : original + flip แล้วเฉลี่ย"""
+    imgs = [img, img.transpose(Image.FLIP_LEFT_RIGHT)]
+    xs = torch.stack([_base_tf(im) for im in imgs], dim=0).to(device)
+    with torch.no_grad():
+        logits = model(xs) / temperature
+        probs = torch.softmax(logits, dim=1).mean(dim=0).cpu().numpy()
+    return probs
+
 # ---------- Core inference ----------
+def label_by_idx(i: int) -> str:
+    return idx_to_class.get(i, RIPENESS_LABELS[i])
+
 def infer_ripeness_and_caption(image: Image.Image):
-    idx, probs = _classify_4class_by_color(image)
-    label = RIPENESS_LABELS[idx]
+    if MODEL_READY:
+        probs = _predict_proba_with_model(image)
+        idx = int(np.argmax(probs))
+    else:
+        idx, probs = _classify_4class_by_color(image)
+    label = label_by_idx(idx)
     raw_conf_pct = float(probs[idx]) * 100.0
     cap = generate_caption(label, raw_conf_pct)
     return idx, probs, cap
@@ -181,21 +271,34 @@ def analyze(image):
         return "กรุณาอัปโหลดภาพ", "", None, "❌ ไม่มีภาพ"
     try:
         idx, probs, caption = infer_ripeness_and_caption(image)
-        class_name = RIPENESS_LABELS[idx]
-        conf = f"{adjust_confidence(float(probs[idx])):.1f}%"
-    except Exception:
-        class_name, conf, caption = "พร้อมรับประทาน(สำรอง)", "100.0%", "เดโม แบบสำรอง"
+
+        order = np.argsort(probs)[::-1]
+        top1, top2 = int(order[0]), int(order[1])
+        p1, p2 = float(probs[top1]), float(probs[top2])
+
+        class_name = label_by_idx(top1)
+        conf_str = f"{adjust_confidence(p1):.1f}%"
+
+        borderline = ""
+        if (p1 - p2) < 0.15:
+            borderline = f"\n⚠️ ก้ำกึ่งระหว่าง {label_by_idx(top1)} ({p1*100:.1f}%) และ {label_by_idx(top2)} ({p2*100:.1f}%)"
+
+        result_text = f"ระดับ: {class_name} (ความมั่นใจ {conf_str}){borderline}"
+    except Exception as e:
+        result_text = "พร้อมรับประทาน(สำรอง) (ความมั่นใจ 100.0%)"
+        caption = "เดโม แบบสำรอง"
 
     ts = datetime.now().strftime("%Y%m%d_%H%M%S")
     out_path = os.path.join(IMG_DIR, f"durian_{ts}.jpg")
     try:
         image.save(out_path, quality=90)
         save_history_row(datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-                         class_name, conf, caption, out_path)
+                         class_name, conf_str, caption, out_path)
     except Exception:
         pass
 
-    return f"ระดับ: {class_name} (ความมั่นใจ {conf})", caption, image, "✅ เสร็จสิ้น"
+    status_text = "✅ เสร็จสิ้น" if MODEL_READY else "ℹ️ ใช้โหมดสำรอง (สี)"
+    return result_text, caption, image, status_text
 
 def show_history():
     rows = load_history(limit=200)
@@ -233,4 +336,5 @@ with gr.Blocks(title="Durian Happiness Level") as demo:
 
 if __name__ == "__main__":
     random.seed()
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)
+    port = int(os.environ.get("PORT", "7860"))
+    demo.launch(server_name="0.0.0.0", server_port=port, show_api=False)