Update main.py

main.py

@@ -7,6 +7,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from PIL import Image
 import tensorflow as tf
+from huggingface_hub import snapshot_download
 
 # optional gatekeep
 try:
@@ -15,21 +16,19 @@ try:
 except Exception:
     HAS_OPENCV = False
 
-# HF Hub (สำหรับดึง derm-foundation)
-from huggingface_hub import snapshot_download
-
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("skinclassify")
 
 # ---------------------- Config ----------------------
 DERM_MODEL_ID  = os.getenv("DERM_MODEL_ID", "google/derm-foundation")
-DERM_LOCAL_DIR = os.getenv("DERM_LOCAL_DIR", "")  # path ไปยัง SavedModel ถ้ามีออฟไลน์
-HEAD_PATH       = os.getenv("HEAD_PATH", "Models/mlp_head_savedmodel/saved_model.pb")
+DERM_LOCAL_DIR = os.getenv("DERM_LOCAL_DIR", "")  # path to local SavedModel if offline
+
+HEAD_PATH       = os.getenv("HEAD_PATH", "Models/mlp_best.keras")  # <== ใช้ .keras ตรง ๆ
 THRESHOLDS_PATH = os.getenv("THRESHOLDS_PATH", "Models/mlp_thresholds.npy")
 MU_PATH         = os.getenv("MU_PATH", "Models/mu.npy")
 SD_PATH         = os.getenv("SD_PATH", "Models/sd.npy")
 LABELS_PATH     = os.getenv("LABELS_PATH", "Models/class_names.json")
-NPZ_PATH        = os.getenv("NPZ_PATH", "")  # ถ้าอยากโหลด mu/sd/class_names จากไฟล์เดียว
+NPZ_PATH        = os.getenv("NPZ_PATH", "")  # optional fallback for mu/sd/class_names
 
 TOPK = int(os.getenv("TOPK", "5"))
 
@@ -40,15 +39,13 @@ MIN_BRIGHT, MAX_BRIGHT = float(os.getenv("MIN_BRIGHT", "20")), float(os.getenv("
 MIN_SKIN_RATIO         = float(os.getenv("MIN_SKIN_RATIO", "0.15"))
 MIN_SHARPNESS          = float(os.getenv("MIN_SHARPNESS", "30.0"))
 
-# Performance (กัน OOM บน free space)
+# Performance: กัน OOM บน Free Space
 os.environ.setdefault("TF_NUM_INTRAOP_THREADS", "1")
 os.environ.setdefault("TF_NUM_INTEROP_THREADS", "1")
 os.environ.setdefault("OMP_NUM_THREADS", "1")
 os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
 
-# Upload size (กัน DoS)
 MAX_UPLOAD = int(os.getenv("MAX_UPLOAD", str(6 * 1024 * 1024)))  # 6MB
-
 DF_SIZE = (448, 448)
 
 app = FastAPI(title="SkinClassify API (Derm-Foundation)", version="2.0.0")
@@ -79,43 +76,13 @@ else:
     raise RuntimeError("LABELS_PATH not found and NPZ_PATH not provided.")
 C = len(CLASS_NAMES)
 
-# ---------------------- Load head ----------------------
-# logger.info(f"Loading head from {HEAD_PATH}")
-# head = tf.keras.models.load_model(HEAD_PATH, compile=False)
-import inspect
-
-def load_head_any(path: str):
-    # กรณีเป็นโฟลเดอร์ SavedModel
-    if os.path.isdir(path) and os.path.exists(os.path.join(path, "saved_model.pb")):
-        m = tf.saved_model.load(path)
-        sig = m.signatures.get("serving_default") or next(iter(m.signatures.values()))
-        # หาชื่อ input/output อัตโนมัติ
-        in_names  = list(sig.structured_input_signature[1].keys())
-        out_names = list(sig.structured_outputs.keys())
-
-        def _predict(z_np: np.ndarray) -> np.ndarray:
-            z_tf = tf.convert_to_tensor(z_np, dtype=tf.float32)
-            # บางรุ่นรับเป็น args บางรุ่นรับเป็น kwargs
-            if len(in_names) == 1:
-                out_dict = sig(z_tf)
-            else:
-                out_dict = sig(**{in_names[0]: z_tf})
-            y = out_dict[out_names[0]]
-            return y.numpy()
-        return _predict
-
-    # กรณีเป็นไฟล์ .h5 / .keras ที่ tf.keras อ่านได้
-    try:
-        model = tf.keras.models.load_model(path, compile=False)
-        return lambda z_np: model.predict(z_np, verbose=0)
-    except Exception as e:
-        raise RuntimeError(f"Cannot load head from {path}: {e}")
-
-# ตั้ง path ไปยัง SavedModel ที่อัปขึ้น
-HEAD_PATH = os.getenv("HEAD_PATH", "Models/mlp_head_savedmodel/")
-logger.info(f"Loading head from {HEAD_PATH}")
-head_predict = load_head_any(HEAD_PATH)
+# ---------------------- Load head (.keras via Keras3) ----------------------
+def load_head_keras3(path: str):
+    import keras
+    logger.info(f"Loading head (.keras) via Keras3 from {path}")
+    return keras.saving.load_model(path, compile=False)
 
+head = load_head_keras3(HEAD_PATH)
 
 # ---------------------- Load mu/sd ----------------------
 def _load_mu_sd():
@@ -142,14 +109,7 @@ else:
     logger.warning("THRESHOLDS_PATH not found -> default 0.5 for all classes")
     best_th = np.full(C, 0.5, dtype="float32")
 
-# ---------------------- Wrap head with standardization ----------------------
-inp = tf.keras.Input(shape=(mu.shape[-1],), name="embedding")
-x = tf.keras.layers.Lambda(lambda e: (e - mu) / (sd + 1e-6), name="standardize")(inp)
-out = head(x)
-clf = tf.keras.Model(inp, out, name="head_with_norm")
-
 # ---------------------- Load derm-foundation ----------------------
-# ใช้ snapshot_download + tf.saved_model.load (ถูกกับโมเดลของ Google)
 logger.info("Loading Derm Foundation (first time may take a while)...")
 try:
     if DERM_LOCAL_DIR and os.path.isdir(DERM_LOCAL_DIR) and os.path.exists(os.path.join(DERM_LOCAL_DIR, "saved_model.pb")):
@@ -225,19 +185,6 @@ def gatekeep_image(img_bytes: bytes) -> Dict[str, Any]:
     if ratio < MIN_SKIN_RATIO: reasons.append("not_skin_like")
     return {"ok": len(reasons)==0, "reasons": reasons, "metrics": metrics}
 
-# def predict_probs(img_bytes: bytes) -> np.ndarray:
-#     pil = Image.open(io.BytesIO(img_bytes)).convert("RGB").resize(DF_SIZE)
-#     by  = pil_to_png_bytes_448(pil)
-#     ex  = tf.train.Example(features=tf.train.Features(
-#           feature={'image/encoded': tf.train.Feature(bytes_list=tf.train.BytesList(value=[by]))}
-#     )).SerializeToString()
-#     out = infer(inputs=tf.constant([ex]))
-#     # ปกติคีย์จะชื่อ "embedding"
-#     if "embedding" not in out:
-#         raise RuntimeError(f"Unexpected derm-foundation outputs: {list(out.keys())}")
-#     emb = out["embedding"].numpy().astype("float32")  # (1, 6144)
-#     probs = clf.predict(emb, verbose=0)[0]
-#     return probs
 def predict_probs(img_bytes: bytes) -> np.ndarray:
     pil = Image.open(io.BytesIO(img_bytes)).convert("RGB").resize(DF_SIZE)
     by  = pil_to_png_bytes_448(pil)
@@ -247,13 +194,11 @@ def predict_probs(img_bytes: bytes) -> np.ndarray:
     out = infer(inputs=tf.constant([ex]))
     if "embedding" not in out:
         raise RuntimeError(f"Unexpected derm-foundation outputs: {list(out.keys())}")
-    emb = out["embedding"].numpy().astype("float32")  # (1,6144)
+    emb = out["embedding"].numpy().astype("float32")  # (1, 6144)
     z   = (emb - mu) / (sd + 1e-6)
-    probs = head_predict(z)[0]                        # <-- เรียกหัวที่โหลดมา
+    probs = head.predict(z, verbose=0)[0]             # head (.keras) โดยตรง
     return probs
 
-
-
 # ---------------------- Endpoints ----------------------
 @app.get("/health")
 def health():
@@ -266,7 +211,6 @@ def health():
 
 @app.post("/predict")
 async def predict(request: Request, file: UploadFile = File(...)):
-    # limit content-length
     cl = request.headers.get("content-length")
     if cl and int(cl) > MAX_UPLOAD:
         raise HTTPException(413, "File too large")
@@ -297,8 +241,6 @@ async def predict(request: Request, file: UploadFile = File(...)):
         }
     }
 
-# สำหรับรันนอก Docker (เช่นทดสอบ local)
 if __name__ == "__main__":
     import uvicorn
-    port = int(os.getenv("PORT", "7860"))
-    uvicorn.run(app, host="0.0.0.0", port=port, workers=1)
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")), workers=1)