Update main.py

main.py

@@ -82,28 +82,40 @@ C = len(CLASS_NAMES)
 # ---------------------- Load head ----------------------
 # logger.info(f"Loading head from {HEAD_PATH}")
 # head = tf.keras.models.load_model(HEAD_PATH, compile=False)
-def load_head_model(path: str):
-    # 1) ลอง tf.keras (.h5 / SavedModel)
+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:
-        logger.info(f"Loading head via tf.keras from {path}")
-        return tf.keras.models.load_model(path, compile=False)
-    except Exception as e_tf:
-        logger.warning(f"tf.keras load failed: {e_tf}")
+        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}")
 
-    # 2) ถ้าเป็น .keras (Keras 3)
-    try:
-        import keras
-        logger.info(f"Falling back to keras.saving.load_model for {path}")
-        return keras.saving.load_model(path, compile=False)
-    except Exception as e_k:
-        raise RuntimeError(
-            f"Cannot load head model: {path}\n"
-            f"tf.keras error: {e_tf}\nkeras3 error: {e_k}\n"
-            "If your file is .keras (Keras 3), either add 'keras>=3.3' to requirements "
-            "or re-export to .h5 / SavedModel."
-        )
+# ตั้ง path ไปยัง SavedModel ที่อัปขึ้น
+HEAD_PATH = os.getenv("HEAD_PATH", "Models/mlp_head_savedmodel")
 logger.info(f"Loading head from {HEAD_PATH}")
-head = tf.keras.models.load_model(HEAD_PATH, compile=False)
+head_predict = load_head_any(HEAD_PATH)
+
 
 # ---------------------- Load mu/sd ----------------------
 def _load_mu_sd():
@@ -232,17 +244,16 @@ def predict_probs(img_bytes: bytes) -> np.ndarray:
     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]))
     if "embedding" not in out:
         raise RuntimeError(f"Unexpected derm-foundation outputs: {list(out.keys())}")
-
-    emb = out["embedding"].numpy().astype("float32")  # (1, 6144)
-    z   = (emb - mu) / (sd + 1e-6)                    # ทำ normalization ด้วย NumPy
-    probs = head.predict(z, verbose=0)[0]             # เรียกหัวโมเดลโดยตรง
+    emb = out["embedding"].numpy().astype("float32")  # (1,6144)
+    z   = (emb - mu) / (sd + 1e-6)
+    probs = head_predict(z)[0]                        # <-- เรียกหัวที่โหลดมา
     return probs
 
 
+
 # ---------------------- Endpoints ----------------------
 @app.get("/health")
 def health():