add guard

inference.py

@@ -2,39 +2,56 @@
 import os, json, time
 from typing import List, Dict, Iterable, Any, Optional
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
 import tensorflow as tf
 
 # ---------- Label utils ----------
-_DEFAULT_LABELS = ["Heart","Oblong","Oval","Round","Square"]
+_DEFAULT_LABELS = ["Heart", "Oblong", "Oval", "Round", "Square"]
 
 def _load_labels() -> List[str]:
+    """Prioritas: models/class_indices.json -> models/idx2class.json -> default."""
     p_ci  = os.path.join("models", "class_indices.json")   # {"Label": idx}
-    p_i2c = os.path.join("models", "idx2class.json")       # {"0":"Label"}
+    p_i2c = os.path.join("models", "idx2class.json")       # {"0": "Label"} atau ["Label", ...]
+    # 1) class_indices.json (label->idx)
     try:
         with open(p_ci, "r") as f:
             ci = json.load(f)
-        labels = [k for k,_ in sorted(ci.items(), key=lambda kv: kv[1])]
-        print("[LABEL] from class_indices.json ->", labels)
-        return labels
+        if isinstance(ci, dict):
+            labels = [k for k, _ in sorted(ci.items(), key=lambda kv: kv[1])]
+            print("[LABEL] from class_indices.json ->", labels)
+            return labels
+        if isinstance(ci, list):
+            print("[LABEL] from class_indices.json (list) ->", ci)
+            return list(ci)
     except Exception:
         pass
+    # 2) idx2class.json (idx->label)
     try:
         with open(p_i2c, "r") as f:
             i2c = json.load(f)
-        labels = [i2c[str(i)] for i in range(len(i2c))]
-        print("[LABEL] from idx2class.json ->", labels)
-        return labels
+        if isinstance(i2c, dict):
+            n = len(i2c)
+            labels = [i2c[str(i)] if str(i) in i2c else i2c[i] for i in range(n)]
+            print("[LABEL] from idx2class.json (dict) ->", labels)
+            return labels
+        if isinstance(i2c, list):
+            print("[LABEL] from idx2class.json (list) ->", i2c)
+            return list(i2c)
     except Exception:
         pass
     print("[LABEL] fallback default ->", _DEFAULT_LABELS)
-    return _DEFAULT_LABELS
+    return list(_DEFAULT_LABELS)
 
 def _generate_config_if_missing(model: tf.keras.Model, labels: List[str], path="config.json"):
-    if os.path.exists(path): return
+    """Auto-tulis config.json jika belum ada."""
+    if os.path.exists(path):
+        return
     ishape = model.input_shape
-    h = int(ishape[1])
-    assert h > 0, f"Input shape aneh: {ishape}"
+    try:
+        h = int(ishape[1]); assert h > 0
+    except Exception as e:
+        raise AssertionError(f"Input shape tidak valid untuk config: {ishape}") from e
     cfg = {
         "architectures": ["EfficientNetB4"],
         "image_size": h,
@@ -42,75 +59,105 @@ def _generate_config_if_missing(model: tf.keras.Model, labels: List[str], path="
         "id2label": {str(i): lbl for i, lbl in enumerate(labels)},
         "label2id": {lbl: i for i, lbl in enumerate(labels)},
     }
-    with open(path, "w") as f: json.dump(cfg, f, indent=2)
+    with open(path, "w") as f:
+        json.dump(cfg, f, indent=2)
     print(f"[CFG] wrote {path} (image_size={h}, num_labels={len(labels)})")
 
 # ---------- Model wrapper ----------
 class FaceShapeModel:
     def __init__(self, model_path="models/model.keras"):
-        self.labels = _load_labels()
+        self.labels: List[str] = _load_labels()
+
         full_path = os.path.join(os.getcwd(), model_path)
         print(f"[LOAD] {full_path}")
-        self.model = tf.keras.models.load_model(full_path, compile=False)
+        self.model: tf.keras.Model = tf.keras.models.load_model(full_path, compile=False)
 
+        # input size (H=W)
         ishape = self.model.input_shape
-        self.img_size = int(ishape[1])
+        self.img_size: int = int(ishape[1])
         print(f"[MODEL] input img_size = {self.img_size}")
 
-        names_lower = [l.name.lower() for l in self.model.layers[:10]]
-        has_pp = any(("rescaling" in n) or ("normalization" in n) for n in names_lower)
-        self.external_rescale = not has_pp
-        print(f"[MODEL] internal_preproc={has_pp} -> external_rescale={self.external_rescale}")
+        # apakah model sudah termasuk preprocessing internal?
+        names_lower = [l.name.lower() for l in self.model.layers[:12]]
+        has_internal_pp = any(("rescaling" in n) or ("normalization" in n) for n in names_lower)
+        self.external_rescale: bool = not has_internal_pp
+        print(f"[MODEL] internal_preproc={has_internal_pp} -> external_rescale={self.external_rescale}")
+
+        # sinkronkan jumlah label dengan jumlah output model
+        num_out = int(self.model.output_shape[-1])
+        if len(self.labels) != num_out:
+            print(f"[WARN] labels({len(self.labels)}) != model_out({num_out}) -> menyesuaikan.")
+            if len(self.labels) >= num_out:
+                self.labels = self.labels[:num_out]
+            else:
+                # pad label generik jika kurang
+                self.labels += [f"class_{i}" for i in range(len(self.labels), num_out)]
 
         _generate_config_if_missing(self.model, self.labels)
 
+        # warmup (optional)
         try:
             _ = self.model(tf.zeros((1, self.img_size, self.img_size, 3), dtype=tf.float32))
         except Exception as e:
             print("[WARN] warmup failed:", e)
 
-    def _to_rgb(self, img: Image.Image) -> Image.Image:
+    # ---- preprocessing ----
+    @staticmethod
+    def _to_rgb(img: Image.Image) -> Image.Image:
         return img if img.mode == "RGB" else img.convert("RGB")
 
     def _preprocess(self, img: Image.Image) -> np.ndarray:
         img = self._to_rgb(img).resize((self.img_size, self.img_size))
         x = np.asarray(img, dtype=np.float32)
-        if self.external_rescale: x = x / 255.0
+        if self.external_rescale:
+            x = x / 255.0
         return np.expand_dims(x, 0)  # (1,H,W,3)
 
+    # ---- predict ----
     def predict_dict(self, img: Image.Image) -> Dict[str, float]:
+        """Return dict {label: prob} untuk gr.Label."""
         t0 = time.perf_counter()
-        p = self.model.predict(self._preprocess(img), verbose=0)[0]
+        probs = self.model.predict(self._preprocess(img), verbose=0)[0]  # (C,)
+        # pastikan float murni
+        out = {lbl: float(p) for lbl, p in zip(self.labels, probs)}
         dt = (time.perf_counter() - t0) * 1000.0
         print(f"[INF] {len(self.labels)}-class in {dt:.1f} ms")
-        return {lbl: float(prob) for lbl, prob in zip(self.labels, p)}
+        return out
 
 # singleton
 _MODEL = FaceShapeModel()
 
 # ---------- Public API ----------
 def predict(image: Image.Image) -> Dict[str, float]:
+    # NOTE: gr.Label butuh mapping label->float; jangan kirim string/list.
     if image is None:
-        return {"Error": "No image"}
+        return {"Error": 1.0}
     return _MODEL.predict_dict(image)
 
 def predict_batch(images: Iterable[Any]) -> List[Dict[str, float]]:
+    """Kembalikan list of dict label->prob; cocok untuk gr.JSON di tab Batch."""
     from PIL import Image as _PILImage
     import os as _os
     results: List[Dict[str, float]] = []
 
     def _as_pil(x: Any) -> Optional[_PILImage.Image]:
-        if x is None: return None
-        if isinstance(x, _PILImage.Image): return x
+        if x is None:
+            return None
+        if isinstance(x, _PILImage.Image):
+            return x
         if isinstance(x, (str, bytes, _os.PathLike)):
-            try: return _PILImage.open(x).convert("RGB")
-            except Exception: return None
-        try: return _PILImage.open(x).convert("RGB")
-        except Exception: return None
+            try:
+                return _PILImage.open(x).convert("RGB")
+            except Exception:
+                return None
+        try:
+            return _PILImage.open(x).convert("RGB")
+        except Exception:
+            return None
 
     for x in (images or []):
         im = _as_pil(x)
-        results.append({"Error": "Invalid image"} if im is None else _MODEL.predict_dict(im))
+        results.append({"Error": 1.0} if im is None else _MODEL.predict_dict(im))
     return results
 
 __all__ = ["predict", "predict_batch"]