Update app.py

app.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-🔥 Triple Hybrid API + Gradio UI - HF Spaces Optimized
+🔥 Icon Challenge Solver - Optimized for HF Spaces
 """
 
 import os
@@ -12,7 +12,7 @@ import json
 import logging
 import warnings
 from pathlib import Path
-from typing import List, Union, Optional
+from typing import List, Union
 
 import torch
 import torch.nn as nn
@@ -27,10 +27,6 @@ import torchvision.transforms as T
 import gradio as gr
 import spaces
 
-from fastapi import FastAPI, Request, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from contextlib import asynccontextmanager
-
 from ultralytics import YOLO
 import open_clip
 from peft import PeftModel
@@ -52,7 +48,7 @@ class Config:
 
 CONFIG = Config()
 
-# ===================== MODEL DEFINITION =====================
+# ===================== MODEL CLASSES =====================
 class ProjectionHead(nn.Module):
     def __init__(self, in_dim=1024, hidden_dim=512, out_dim=256, dropout=0.1):
         super().__init__()
@@ -96,16 +92,14 @@ class IconClassifier(nn.Module):
         embedding = self.projection(features)
         return F.normalize(embedding, p=2, dim=1)
 
-# Preprocess
 custom_preprocess = T.Compose([
     T.Resize((256, 256), interpolation=T.InterpolationMode.NEAREST),
     T.ToTensor(),
     T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
 ])
 
-# ===================== IMAGE DECODE =====================
+# ===================== UTILS =====================
 def decode_image_advanced(data: Union[str, bytes], filename: str = "image") -> Image.Image:
-    # ... (sama seperti sebelumnya, saya skip tulis ulang biar tidak terlalu panjang)
     if isinstance(data, str):
         clean = data.strip().replace("\n", "").replace("\r", "").replace(" ", "")
         clean = re.sub(r"^data:image/[a-zA-Z0-9.+-]+;base64,", "", clean)
@@ -121,7 +115,6 @@ def decode_image_advanced(data: Union[str, bytes], filename: str = "image") -> I
     except:
         pass
 
-    # fallback raw bytes
     common_sizes = [50, 60, 64, 70, 80, 96, 100, 128]
     for size in common_sizes:
         expected = size * size * 4
@@ -136,7 +129,6 @@ def decode_image_advanced(data: Union[str, bytes], filename: str = "image") -> I
 # ===================== MODEL MANAGER =====================
 class ModelManager:
     _instance = None
-
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super().__new__(cls)
@@ -144,12 +136,10 @@ class ModelManager:
         return cls._instance
 
     def init(self):
-        if self.initialized:
-            return
-        logger.info("🔥 Initializing models on GPU...")
+        if self.initialized: return
+        logger.info("Loading models on GPU...")
         self.yolo = YOLO(CONFIG.yolo_model_path).to(CONFIG.device)
-        
-        # Load SigLIP + LoRA
+
         config_path = Path(CONFIG.lora_model_dir) / "config.json"
         with open(config_path) as f:
             model_config = json.load(f)
@@ -160,20 +150,18 @@ class ModelManager:
         model_with_lora = PeftModel.from_pretrained(backbone, CONFIG.lora_model_dir)
 
         head_state = torch.load(Path(CONFIG.lora_model_dir) / "head_weights.pt", map_location=CONFIG.device)
-        
         self.inference_model = IconClassifier(model_with_lora, head_state, num_classes).to(CONFIG.device).eval()
-        
+
         self.idx_to_class = {int(k): v for k, v in model_config.get('idx_to_class', {}).items()}
         self.initialized = True
-        logger.info("✅ All models loaded successfully!")
+        logger.info("✅ Models loaded!")
 
+    # (Copy fungsi yolo_detect_and_crop, extract_embedding, match_options dari kode sebelumnya)
     def yolo_detect_and_crop(self, image: Image.Image) -> Image.Image:
-        # ... (sama seperti kode lama kamu)
         img_np = np.array(image)
         img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
         results = self.yolo(img_bgr, conf=CONFIG.yolo_conf, imgsz=CONFIG.yolo_img_size, verbose=False)
         
-        # (saya asumsikan bagian ini tetap sama, copy dari kode lama kamu)
         best_box = None
         best_conf = 0
         for result in results:
@@ -185,9 +173,9 @@ class ModelManager:
                     best_box = box.xyxy[0].cpu().numpy().astype(int)
 
         if best_box is not None and best_conf >= CONFIG.yolo_min_conf:
-            x1,y1,x2,y2 = best_box
-            h,w = img_np.shape[:2]
-            x1,y1,x2,y2 = max(0,x1),max(0,y1),min(w,x2),min(h,y2)
+            x1, y1, x2, y2 = best_box
+            h, w = img_np.shape[:2]
+            x1, y1, x2, y2 = max(0,x1), max(0,y1), min(w,x2), min(h,y2)
             if (x2-x1) >= CONFIG.yolo_min_crop_size and (y2-y1) >= CONFIG.yolo_min_crop_size:
                 return Image.fromarray(img_np[y1:y2, x1:x2])
         return image
@@ -210,23 +198,24 @@ class ModelManager:
         return {
             "predicted_index": best_idx,
             "confidence": float(probs[best_idx]),
-            "similarity_score": float(similarities[best_idx]),
-            "all_scores": {f"option_{i}": {"similarity": float(similarities[i]), "probability": float(probs[i])} for i in range(len(options))},
+            "all_scores": {f"option_{i}": {
+                "similarity": float(similarities[i]),
+                "probability": float(probs[i])
+            } for i in range(len(options))},
             "predicted_class": self.idx_to_class.get(best_idx, "unknown")
         }
 
 manager = ModelManager()
 
-# ===================== GRADIO FUNCTION (GPU) =====================
-@spaces.GPU(request_timeout=120, enable=True)
-def solve_gradio(challenge_img, opt1, opt2, opt3, opt4, opt5, opt6, opt7, opt8):
+# ===================== MAIN FUNCTION =====================
+@spaces.GPU(request_timeout=180, enable=True)
+def solve_gradio(challenge_img, *options):
     if not manager.initialized:
         manager.init()
 
     try:
-        options = [opt1, opt2, opt3, opt4, opt5, opt6, opt7, opt8]
         if challenge_img is None or any(o is None for o in options):
-            return "Error: Semua gambar harus diupload!", None
+            return "❌ Semua gambar harus diisi!", None
 
         # Convert ke PIL
         if isinstance(challenge_img, np.ndarray):
@@ -237,17 +226,18 @@ def solve_gradio(challenge_img, opt1, opt2, opt3, opt4, opt5, opt6, opt7, opt8):
         result = manager.match_options(challenge_crop, options_pil)
 
         best_idx = result["predicted_index"]
-        lines = ["🎯 RESULT\n",
-                 f"Best Option: {best_idx + 1}",
+        lines = ["🎯 HASIL\n",
+                 f"✅ Best Option: {best_idx + 1}",
                  f"Confidence: {result['confidence']:.4f} ({result['confidence']*100:.2f}%)",
-                 f"Predicted Class: {result['predicted_class']}\n",
-                 "📊 All Scores:"]
+                 f"Class: {result['predicted_class']}\n",
+                 "📊 Skor Semua Opsi:"]
 
         for i in range(8):
             s = result["all_scores"][f"option_{i}"]
-            marker = "✅" if i == best_idx else "  "
+            marker = "✅" if i == best_idx else "   "
             lines.append(f"{marker} Option {i+1}: {s['probability']:.4f} ({s['similarity']:.4f})")
 
+        # Cleanup
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -255,38 +245,35 @@ def solve_gradio(challenge_img, opt1, opt2, opt3, opt4, opt5, opt6, opt7, opt8):
         return "\n".join(lines), options_pil[best_idx]
 
     except Exception as e:
-        logger.error(f"Error: {e}", exc_info=True)
-        return f"Runtime Error: {str(e)}", None
+        logger.error(str(e))
+        return f"❌ Error: {str(e)}", None
+
 
+# ===================== GRADIO INTERFACE =====================
+with gr.Blocks(title="🔥 Icon Challenge Solver", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🔥 Icon Challenge Solver\n**YOLO + SigLIP LoRA Hybrid**")
 
-# ===================== GRADIO UI =====================
-with gr.Blocks(title="Icon Challenge Solver", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🔥 Icon Challenge Solver\n**YOLO + SigLIP LoRA**")
-    
     with gr.Row():
         with gr.Column(scale=1):
-            challenge_input = gr.Image(label="Challenge Image", type="numpy", height=220)
+            challenge_input = gr.Image(label="🖼️ Challenge Image", type="numpy", height=240)
         with gr.Column(scale=2):
-            gr.Markdown("### 8 Option Images")
+            gr.Markdown("### 🎲 8 Option Images")
             with gr.Row():
-                opts = [gr.Image(label=f"Option {i+1}", type="numpy", height=140) for i in range(8)]
-    
+                opt_inputs = [gr.Image(label=f"Option {i+1}", type="numpy", height=130) for i in range(8)]
+
     solve_btn = gr.Button("🔍 Solve Challenge", variant="primary", size="large")
-    
+
     with gr.Row():
-        result_text = gr.Textbox(label="Result", lines=12, interactive=False)
-        best_img = gr.Image(label="Best Match")
+        result_text = gr.Textbox(label="📋 Result", lines=14, interactive=False)
+        best_image = gr.Image(label="🏆 Best Match")
 
     solve_btn.click(
         fn=solve_gradio,
-        inputs=[challenge_input] + opts,
-        outputs=[result_text, best_img]
+        inputs=[challenge_input] + opt_inputs,
+        outputs=[result_text, best_image]
     )
 
-# ===================== FASTAPI (opsional) =====================
-app = FastAPI(title="Icon Solver")
-app = gr.mount_gradio_app(app, demo, path="/")
+    gr.Markdown("---\nMade with ❤️ for Icon Challenge")
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+# Launch (HF Spaces akan handle ini otomatis)
+demo.launch()