Update app.py

app.py

@@ -1,24 +1,23 @@
-import os
-import cv2
 import gradio as gr
 import torch
-import numpy as np
-import mediapipe as mp
-from PIL import Image, ImageOps
 from transformers import SwinForImageClassification, AutoFeatureExtractor
+import mediapipe as mp
+import cv2
+from PIL import Image
+import numpy as np
+import os
 
-# =========================
-# Konfigurasi & Metadata
-# =========================
-FACE_SHAPE_DESCRIPTIONS = {
-    "Heart":  "dengan dahi lebar dan dagu yang runcing.",
+# --- Face shape descriptions
+face_shape_descriptions = {
+    "Heart": "dengan dahi lebar dan dagu yang runcing.",
     "Oblong": "yang lebih panjang dari lebar dengan garis pipi lurus.",
-    "Oval":   "dengan proporsi seimbang dan dagu sedikit melengkung.",
-    "Round":  "dengan garis rahang melengkung dan pipi penuh.",
+    "Oval": "dengan proporsi seimbang dan dagu sedikit melengkung.",
+    "Round": "dengan garis rahang melengkung dan pipi penuh.",
     "Square": "dengan rahang tegas dan dahi lebar."
 }
 
-GLASSES_IMAGES = {
+# --- Glasses frame images
+glasses_images = {
     "Oval": "glasses/oval.jpg",
     "Round": "glasses/round.jpg",
     "Square": "glasses/square.jpg",
@@ -27,209 +26,141 @@ GLASSES_IMAGES = {
     "Pilot (Aviator)": "glasses/aviator.jpg"
 }
 
-# Pastikan folder 'glasses' ada & isi dummy kalau file hilang
-os.makedirs("glasses", exist_ok=True)
-for _, p in GLASSES_IMAGES.items():
-    if not os.path.exists(p):
-        Image.new("RGB", (300, 160), color="gray").save(p)
+# Ensure folder exists
+if not os.path.exists("glasses"):
+    os.makedirs("glasses")
+    for _, path in glasses_images.items():
+        if not os.path.exists(path):
+            Image.new('RGB', (200, 100), color='gray').save(path)
 
-ID2LABEL = {0: 'Heart', 1: 'Oblong', 2: 'Oval', 3: 'Round', 4: 'Square'}
-LABEL2ID = {v: k for k, v in ID2LABEL.items()}
+id2label = {0: 'Heart', 1: 'Oblong', 2: 'Oval', 3: 'Round', 4: 'Square'}
+label2id = {v: k for k, v in id2label.items()}
 
-# =========================
-# Model & Device
-# =========================
+# --- Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model_ckpt = "microsoft/swin-tiny-patch4-window7-224"
+model_checkpoint = "microsoft/swin-tiny-patch4-window7-224"
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_checkpoint)
 
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
 model = SwinForImageClassification.from_pretrained(
-    model_ckpt,
-    label2id=LABEL2ID,
-    id2label=ID2LABEL,
+    model_checkpoint,
+    label2id=label2id,
+    id2label=id2label,
     ignore_mismatched_sizes=True
 )
 
-# Muat bobot terlatih jika ada
-if os.path.exists("LR-0001-adamW-32-64swin.pth"):
-    state_dict = torch.load("LR-0001-adamW-32-64swin.pth", map_location=device)
+# Load trained weights if available
+if os.path.exists('LR-0001-adamW-32-64swin.pth'):
+    state_dict = torch.load('LR-0001-adamW-32-64swin.pth', map_location=device)
     model.load_state_dict(state_dict, strict=False)
+model.to(device).eval()
 
-model.to(device)
-model.eval()
+# --- Mediapipe
+mp_face_detection = mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5)
 
-# =========================
-# Utils
-# =========================
-def recommend_glasses_tree(face_shape: str):
-    s = face_shape.strip().lower()
-    if s == "square":
+# --- Decision tree rules
+def recommend_glasses_tree(face_shape):
+    face_shape = face_shape.lower()
+    if face_shape == "square":
         return ["Oval", "Round"]
-    if s == "round":
+    elif face_shape == "round":
         return ["Square", "Octagon", "Cat Eye"]
-    if s == "oval":
+    elif face_shape == "oval":
         return ["Oval", "Pilot (Aviator)", "Cat Eye", "Round"]
-    if s == "heart":
+    elif face_shape == "heart":
         return ["Oval", "Round", "Cat Eye", "Pilot (Aviator)"]
-    if s == "oblong":
+    elif face_shape == "oblong":
         return ["Square", "Pilot (Aviator)", "Cat Eye"]
-    return []
-
-def _pil_to_bgr_ndarray(img: Image.Image) -> np.ndarray:
-    """Pastikan image RGB, buang alpha/EXIF, lalu ke BGR (OpenCV)."""
-    if img.mode not in ("RGB", "L"):
-        # handle RGBA/CMYK/LA dll
-        img = img.convert("RGB")
-    elif img.mode == "L":
-        img = ImageOps.colorize(img, black="black", white="white").convert("RGB")
-
-    # strip EXIF untuk safety
-    img_no_exif = Image.new(img.mode, img.size)
-    img_no_exif.putdata(list(img.getdata()))
-
-    arr = np.array(img_no_exif, dtype=np.uint8)  # RGB
-    return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
-
-def _safe_crop(img: np.ndarray, x1, y1, x2, y2):
-    h, w = img.shape[:2]
-    x1 = max(0, min(w-1, x1))
-    y1 = max(0, min(h-1, y1))
-    x2 = max(x1+1, min(w, x2))
-    y2 = max(y1+1, min(h, y2))
-    return img[y1:y2, x1:x2]
-
-def _center_crop_square(img: np.ndarray) -> np.ndarray:
-    h, w = img.shape[:2]
-    side = min(h, w)
-    x1 = (w - side) // 2
-    y1 = (h - side) // 2
-    return img[y1:y1+side, x1:x1+side]
-
-def preprocess_image(pil_image: Image.Image) -> torch.Tensor:
-    """Deteksi wajah (mediapipe). Jika gagal → fallback center crop. Resize 224, ke pixel_values tensor."""
-    bgr = _pil_to_bgr_ndarray(pil_image)
-
-    # mediapipe face detection (buat objek per-call biar thread-safe)
-    with mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5) as fd:
-        results = fd.process(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
-
-    if results.detections:
-        det = results.detections[0]
-        bbox = det.location_data.relative_bounding_box
-        h, w = bgr.shape[:2]
-        x1 = int(bbox.xmin * w)
-        y1 = int(bbox.ymin * h)
-        x2 = int((bbox.xmin + bbox.width) * w)
-        y2 = int((bbox.ymin + bbox.height) * h)
-
-        face = _safe_crop(bgr, x1, y1, x2, y2)
-        if face.size == 0 or face.shape[0] < 32 or face.shape[1] < 32:
-            # kalau box aneh → fallback
-            face = _center_crop_square(bgr)
     else:
-        # fallback kalau wajah tidak terdeteksi
-        face = _center_crop_square(bgr)
-
-    face = cv2.resize(face, (224, 224), interpolation=cv2.INTER_AREA)
-    face_rgb = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
-    inputs = feature_extractor(images=face_rgb, return_tensors="pt")
-    return inputs["pixel_values"].squeeze(0)
-
-def format_gallery_items(frames):
-    items = []
-    for frame in frames:
-        path = GLASSES_IMAGES.get(frame)
-        if path and os.path.exists(path):
-            try:
-                img = Image.open(path).convert("RGB")
-                items.append((img, frame))
-            except Exception as e:
-                print(f"[WARN] gagal buka gambar frame {frame}: {e}")
-    return items
-
-# =========================
-# Prediksi
-# =========================
-@torch.inference_mode()
-def predict(image: Image.Image):
-    """
-    Return:
-      - bentuk_wajah_terdeteksi: str
-      - penjelasan: str
-      - rekomendasi_kacamata: list[(PIL.Image, caption)]
-    """
+        return []
+
+# --- Preprocess image
+def preprocess_image(image):
+    img = np.array(image, dtype=np.uint8)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+
+    results = mp_face_detection.process(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+
+    if not results.detections:
+        return None  # no face detected
+
+    detection = results.detections[0]
+    bbox = detection.location_data.relative_bounding_box
+    h, w, _ = img.shape
+    x1 = max(int(bbox.xmin * w), 0)
+    y1 = max(int(bbox.ymin * h), 0)
+    x2 = min(int((bbox.xmin + bbox.width) * w), w)
+    y2 = min(int((bbox.ymin + bbox.height) * h), h)
+
+    if x2 <= x1 or y2 <= y1:
+        return None
+
+    face_img = img[y1:y2, x1:x2]
+    face_img = cv2.resize(face_img, (224, 224))
+    face_img = cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB)
+
+    inputs = feature_extractor(images=face_img, return_tensors="pt")
+    return inputs['pixel_values'].squeeze(0)
+
+# --- Prediction
+def predict(image):
     try:
-        if image is None:
-            return "Unknown", "Tidak ada gambar yang diunggah.", []
-
-        pixel_values = preprocess_image(image).unsqueeze(0).to(device)
-
-        # AMP untuk GPU
-        if device.type == "cuda":
-            with torch.cuda.amp.autocast():
-                outputs = model(pixel_values)
-        else:
-            outputs = model(pixel_values)
-
-        probs = torch.softmax(outputs.logits, dim=1)[0]
-        pred_idx = int(torch.argmax(probs).item())
-        pred_label = ID2LABEL[pred_idx]
-        conf = float(probs[pred_idx].item()) * 100.0
-
-        recs = recommend_glasses_tree(pred_label)
-        gallery = format_gallery_items(recs)
-
-        desc = FACE_SHAPE_DESCRIPTIONS.get(pred_label, "tidak dikenali")
-        if recs:
-            rec_text = ", ".join(recs)
-            explanation = (
-                f"Bentuk wajah kamu adalah {pred_label} ({conf:.2f}%). "
-                f"Kamu memiliki bentuk wajah {desc} "
-                f"Rekomendasi bentuk kacamata yang sesuai: {rec_text}."
-            )
-        else:
-            explanation = (
-                f"Bentuk wajah kamu adalah {pred_label} ({conf:.2f}%). "
-                f"Belum ada rekomendasi frame untuk bentuk wajah ini."
-            )
-
-        return pred_label, explanation, gallery
+        inputs = preprocess_image(image)
+        if inputs is None:
+            return "Unknown", "⚠️ Wajah tidak terdeteksi. Silakan upload foto dengan wajah yang jelas.", []
+
+        inputs = inputs.unsqueeze(0).to(device)
+        with torch.no_grad():
+            outputs = model(inputs)
+            probs = torch.nn.functional.softmax(outputs.logits, dim=1)
+            pred_idx = torch.argmax(probs, dim=1).item()
+            pred_label = id2label[pred_idx]
+            pred_prob = probs[0][pred_idx].item() * 100
+
+        # --- threshold confidence
+        if pred_prob < 70:
+            return "Unknown", f"⚠️ Prediksi tidak yakin (Confidence {pred_prob:.2f}%). Silakan gunakan foto wajah yang lebih jelas.", []
+
+        # --- glasses recommendation
+        frame_recommendations = recommend_glasses_tree(pred_label)
+        gallery_items = []
+        for frame in frame_recommendations:
+            frame_image_path = glasses_images.get(frame)
+            if frame_image_path and os.path.exists(frame_image_path):
+                gallery_items.append((Image.open(frame_image_path), frame))
+
+        description = face_shape_descriptions.get(pred_label, "tidak dikenali")
+        recommended_frames_text = ', '.join(frame_recommendations)
+
+        explanation = (
+            f"Bentuk wajah kamu adalah {pred_label} ({pred_prob:.2f}%). "
+            f"Kamu memiliki bentuk wajah {description} "
+            f"Rekomendasi bentuk kacamata: {recommended_frames_text}."
+        )
+
+        return pred_label, explanation, gallery_items
 
     except Exception as e:
-        # Jangan pernah crash: selalu balikin tuple sesuai schema
-        return "Unknown", f"Terjadi kesalahan saat memproses gambar: {str(e)}", []
+        return "Error", f"Terjadi kesalahan: {str(e)}", []
 
-# =========================
-# Gradio UI
-# =========================
+# --- Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as iface:
-    gr.Markdown("# Program Rekomendasi Kacamata Berdasarkan Bentuk Wajah")
+    gr.Markdown("# 👓 Program Rekomendasi Kacamata Berdasarkan Bentuk Wajah")
     gr.Markdown("Upload foto wajahmu untuk mendapatkan rekomendasi bentuk kacamata yang sesuai.")
 
     with gr.Row():
         with gr.Column():
-            # penting: type='pil' supaya API melihat argumen bernama 'image'
-            image_input = gr.Image(type="pil", label="Gambar Wajah")
-            confirm_button = gr.Button("Konfirmasi", variant="primary")
+            image_input = gr.Image(type="pil")
+            confirm_button = gr.Button("Konfirmasi")
             restart_button = gr.Button("Restart")
         with gr.Column():
-            detected_shape = gr.Textbox(label="Bentuk Wajah Terdeteksi", interactive=False)
-            explanation_output = gr.Textbox(label="Penjelasan", lines=4, interactive=False)
-            recommendation_gallery = gr.Gallery(label="Rekomendasi Kacamata", columns=3, show_label=True)
-
-    confirm_button.click(
-        predict,
-        inputs=[image_input],
-        outputs=[detected_shape, explanation_output, recommendation_gallery]
-    )
-    restart_button.click(
-        lambda: (None, "", []),
-        inputs=None,
-        outputs=[image_input, detected_shape, explanation_output, recommendation_gallery]
-    )
-
-    gr.Markdown("**Sumber gambar kacamata**: Katalog dari glassdirect.co.uk")
+            detected_shape = gr.Textbox(label="Bentuk Wajah Terdeteksi")
+            explanation_output = gr.Textbox(label="Penjelasan")
+            recommendation_gallery = gr.Gallery(label="Rekomendasi Kacamata", columns=3, show_label=False)
+
+    confirm_button.click(predict, inputs=image_input, outputs=[detected_shape, explanation_output, recommendation_gallery])
+    restart_button.click(lambda: (None, "", "", []), inputs=None, outputs=[image_input, detected_shape, explanation_output, recommendation_gallery])
+
+    gr.Markdown("**Sumber gambar kacamata**: Katalog dari [glassdirect.co.uk](https://www.glassdirect.co.uk)")
 
 if __name__ == "__main__":
-    # show_error=True biar pesan error server-side tampil jelas saat debug
-    iface.queue().launch(show_error=True, server_name="0.0.0.0", server_port=7860)
+    iface.launch()