Create app.py

app.py

@@ -0,0 +1,591 @@
+#!/usr/bin/env python3
+"""
+🔥 Triple Hybrid API + Gradio UI
+FastAPI backend + Gradio frontend (1 app, 2 interface)
+"""
+
+import os
+import io
+import re
+import gc
+import base64
+import json
+import math
+import logging
+import warnings
+from pathlib import Path
+from typing import List, Union, Optional, Tuple
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.cuda.amp import autocast
+
+import numpy as np
+from PIL import Image
+import cv2
+import torchvision.transforms as T
+
+from fastapi import FastAPI, Request, HTTPException, File, UploadFile
+from fastapi.concurrency import run_in_threadpool
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse
+from pydantic import BaseModel
+from contextlib import asynccontextmanager
+
+# Gradio
+import gradio as gr
+
+from ultralytics import YOLO
+import open_clip
+from peft import PeftModel
+
+warnings.filterwarnings("ignore")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+# ═══════════════════════════════════════════════════════════════
+# CONFIG
+# ═══════════════════════════════════════════════════════════════
+@dataclass
+class Config:
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    yolo_model_path: str = "./yolo/best.pt"
+    yolo_conf: float = 0.25
+    yolo_img_size: int = 640
+    yolo_min_crop_size: int = 32
+    yolo_min_conf: float = 0.3
+    lora_model_dir: str = "./icon_model_hybrid"
+    model_name: str = "ViT-L-16-SigLIP-256"
+    img_size: int = 256
+
+CONFIG = Config()
+
+
+# ═══════════════════════════════════════════════════════════════
+# PROJECTION HEAD (Sesuai Training)
+# ═══════════════════════════════════════════════════════════════
+class ProjectionHead(nn.Module):
+    def __init__(self, in_dim=1024, hidden_dim=512, out_dim=256, dropout=0.1):
+        super().__init__()
+        self.layer1 = nn.Linear(in_dim, hidden_dim)
+        self.bn1 = nn.BatchNorm1d(hidden_dim)
+        self.act = nn.GELU()
+        self.dropout = nn.Dropout(dropout)
+        self.layer2 = nn.Linear(hidden_dim, out_dim)
+        self.bn2 = nn.BatchNorm1d(out_dim)
+        self.skip = nn.Linear(in_dim, out_dim)
+
+    def forward(self, x):
+        if x.shape[0] <= 1:
+            self.train(False)
+        residual = self.skip(x)
+        x = self.layer1(x)
+        x = self.bn1(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        x = self.layer2(x)
+        x = self.bn2(x)
+        return F.normalize(x + residual, p=2, dim=1)
+
+
+class IconClassifier(nn.Module):
+    def __init__(self, backbone, head_state, num_classes=21):
+        super().__init__()
+        self.backbone = backbone
+        self.projection = ProjectionHead(1024, 512, 256)
+        self.weight = nn.Parameter(head_state['arc_margin']['weight'])
+        self.projection.load_state_dict(head_state['projection'])
+        self.s = 30.0
+
+    def forward(self, x):
+        features = self.backbone(x)
+        embedding = self.projection(features)
+        logits = F.linear(F.normalize(embedding), F.normalize(self.weight))
+        return logits * self.s
+
+    def get_embedding(self, x):
+        features = self.backbone(x)
+        embedding = self.projection(features)
+        return F.normalize(embedding, p=2, dim=1)
+
+
+# ═══════════════════════════════════════════════════════════════
+# PREPROCESSING (ImageNet — sesuai training)
+# ═══════════════════════════════════════════════════════════════
+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
+# ═══════════════════════════════════════════════════════════════
+def decode_image_advanced(data: Union[str, bytes], filename: str = "image") -> Image.Image:
+    if isinstance(data, str):
+        clean = data.strip().replace("\n", "").replace("\r", "").replace(" ", "")
+        clean = re.sub(r"^data:image/[a-zA-Z0-9.+-]+;base64,", "", clean)
+        missing = len(clean) % 4
+        if missing:
+            clean += '=' * (4 - missing)
+        try:
+            raw_bytes = base64.b64decode(clean, validate=False)
+        except Exception as e:
+            raise ValueError(f"Base64 decode failed: {e}")
+    else:
+        raw_bytes = data
+
+    try:
+        img = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+        return img
+    except Exception:
+        pass
+
+    common_sizes = [50, 60, 64, 70, 80, 96, 100, 128]
+    for size in common_sizes:
+        expected_bytes = size * size * 4
+        if len(raw_bytes) >= expected_bytes:
+            try:
+                arr = np.frombuffer(raw_bytes[:expected_bytes], dtype=np.uint8)
+                arr = arr.reshape((size, size, 4))
+                rgb_arr = arr[:, :, :3]
+                return Image.fromarray(rgb_arr)
+            except Exception:
+                continue
+
+    raise ValueError(f"Cannot decode {filename}")
+
+
+# ═══════════════════════════════════════════════════════════════
+# MODEL MANAGER
+# ═══════════════════════════════════════════════════════════════
+class ModelManager:
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance.initialized = False
+        return cls._instance
+
+    def init(self):
+        if self.initialized:
+            return
+
+        logger.info("=" * 60)
+        logger.info("🔥 INITIALIZING TRIPLE HYBRID API + Gradio")
+        logger.info("=" * 60)
+
+        logger.info(f"📦 Loading YOLO: {CONFIG.yolo_model_path}")
+        self.yolo = YOLO(CONFIG.yolo_model_path)
+        self.yolo.to(CONFIG.device)
+        logger.info("✅ YOLO loaded")
+
+        logger.info(f"📦 Loading Hybrid LoRA from: {CONFIG.lora_model_dir}")
+
+        config_path = Path(CONFIG.lora_model_dir) / "config.json"
+        with open(config_path) as f:
+            self.model_config = json.load(f)
+
+        num_classes = self.model_config.get('num_classes')
+        if num_classes is None:
+            class_to_idx = self.model_config.get('class_to_idx', {})
+            num_classes = len(class_to_idx)
+
+        model_name = self.model_config.get('model_name', CONFIG.model_name)
+        logger.info(f"   Config: {num_classes} classes, {model_name}")
+
+        backbone, _, _ = open_clip.create_model_and_transforms(model_name, pretrained='webli')
+        backbone = backbone.visual.to(CONFIG.device)
+        model_with_lora = PeftModel.from_pretrained(backbone, CONFIG.lora_model_dir)
+
+        head_path = Path(CONFIG.lora_model_dir) / "head_weights.pt"
+        head_state = torch.load(head_path, map_location=CONFIG.device)
+
+        self.inference_model = IconClassifier(model_with_lora, head_state, num_classes).to(CONFIG.device).eval()
+        logger.info("✅ Hybrid LoRA loaded")
+
+        self.idx_to_class = {int(k): v for k, v in self.model_config.get('idx_to_class', {}).items()}
+        self.class_names = [self.idx_to_class.get(i, f"class_{i}") for i in range(num_classes)]
+
+        self.initialized = True
+        logger.info("=" * 60)
+        logger.info("🚀 ALL MODELS READY!")
+        logger.info("=" * 60)
+
+    def yolo_detect_and_crop(self, image: Image.Image) -> Image.Image:
+        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)
+
+        best_box = None
+        best_conf = 0
+        for result in results:
+            if result.boxes is None:
+                continue
+            for box in result.boxes:
+                conf = float(box.conf[0])
+                if conf > best_conf:
+                    best_conf = conf
+                    best_box = box.xyxy[0].cpu().numpy().astype(int)
+
+        if best_box is not None:
+            x1, y1, x2, y2 = best_box
+            h, w = img_np.shape[:2]
+            x1, y1 = max(0, x1), max(0, y1)
+            x2, y2 = min(w, x2), min(h, y2)
+            crop_w = x2 - x1
+            crop_h = y2 - y1
+
+            if best_conf < CONFIG.yolo_min_conf:
+                logger.warning(f"YOLO conf too low ({best_conf:.3f}), using full image")
+                return image
+            if crop_w < CONFIG.yolo_min_crop_size or crop_h < CONFIG.yolo_min_crop_size:
+                logger.warning(f"YOLO crop too small ({crop_w}x{crop_h}), using full image")
+                return image
+
+            logger.info(f"YOLO crop: ({x1},{y1})→({x2},{y2}) [{crop_w}x{crop_h}], conf={best_conf:.3f}")
+            return Image.fromarray(img_np[y1:y2, x1:x2])
+        else:
+            logger.warning("YOLO: No detection, using full image")
+            return image
+
+    def extract_embedding(self, image: Image.Image) -> torch.Tensor:
+        img_tensor = custom_preprocess(image).unsqueeze(0).to(CONFIG.device)
+        with torch.no_grad():
+            with autocast():
+                emb = self.inference_model.get_embedding(img_tensor)
+        return emb
+
+    def match_options(self, challenge_crop: Image.Image, options: List[Image.Image]) -> dict:
+        challenge_emb = self.extract_embedding(challenge_crop)
+        option_embs = []
+        for opt_img in options:
+            emb = self.extract_embedding(opt_img)
+            option_embs.append(emb)
+
+        option_stack = torch.cat(option_embs, dim=0)
+        similarities = (challenge_emb @ option_stack.T).squeeze(0).cpu().numpy()
+
+        best_idx = int(np.argmax(similarities))
+        best_score = float(similarities[best_idx])
+
+        exp_sim = np.exp(similarities * 10)
+        probs = exp_sim / exp_sim.sum()
+
+        return {
+            "predicted_index": best_idx,
+            "confidence": float(probs[best_idx]),
+            "similarity_score": best_score,
+            "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")
+        }
+
+
+# ═══════════════════════════════════════════════════════════════
+# GRADIO FUNCTION
+# ═══════════════════════════════════════════════════════════════
+def solve_gradio(challenge_img, opt1, opt2, opt3, opt4, opt5, opt6, opt7, opt8):
+    """
+    Gradio interface function.
+    Input: 9 images (1 challenge + 8 options)
+    Output: Result text + highlighted best option
+    """
+    manager = ModelManager()
+    if not manager.initialized:
+        return "Error: Model not initialized", None
+
+    try:
+        options = [opt1, opt2, opt3, opt4, opt5, opt6, opt7, opt8]
+
+        # Validate
+        if challenge_img is None:
+            return "Error: Challenge image required", None
+        if any(o is None for o in options):
+            return "Error: All 8 options required", None
+
+        # Convert numpy/PIL
+        if isinstance(challenge_img, np.ndarray):
+            challenge_img = Image.fromarray(challenge_img)
+
+        options_pil = []
+        for i, opt in enumerate(options):
+            if isinstance(opt, np.ndarray):
+                opt = Image.fromarray(opt)
+            options_pil.append(opt)
+
+        # YOLO crop
+        challenge_crop = manager.yolo_detect_and_crop(challenge_img)
+
+        # Match
+        result = manager.match_options(challenge_crop, options_pil)
+
+        # Format output
+        best_idx = result["predicted_index"]
+        confidence = result["confidence"]
+        predicted_class = result["predicted_class"]
+
+        output_text = f"""
+🎯 **RESULT**
+
+**Best Option:** {best_idx + 1} (index {best_idx})
+**Confidence:** {confidence:.4f} ({confidence*100:.2f}%)
+**Predicted Class:** {predicted_class}
+**Similarity Score:** {result['similarity_score']:.4f}
+
+📊 **All Scores:**
+"""
+        for i in range(8):
+            score = result["all_scores"][f"option_{i}"]
+            marker = "✅" if i == best_idx else "  "
+            output_text += f"{marker} Option {i+1}: {score['probability']:.4f} ({score['similarity']:.4f})
+"
+
+        # Return best option image
+        best_image = options_pil[best_idx]
+
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        return output_text, best_image
+
+    except Exception as e:
+        logger.error(f"Gradio error: {e}", exc_info=True)
+        return f"Error: {str(e)}", None
+
+
+# ═══════════════════════════════════════════════════════════════
+# BUILD GRADIO INTERFACE
+# ═══════════════════════════════════════════════════════════════
+def build_gradio_interface() -> gr.Blocks:
+    with gr.Blocks(title="🔥 Icon Challenge Solver", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 🔥 Icon Challenge Solver
+        ### YOLO + SigLIP LoRA Hybrid
+
+        Upload 1 challenge image + 8 option images. The AI will:
+        1. **YOLO**: Detect & crop the icon from challenge
+        2. **SigLIP LoRA**: Match the crop with 8 options
+        3. Return the best matching option
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### 🎯 Challenge Image")
+                challenge_input = gr.Image(
+                    label="Challenge (1 image)",
+                    type="numpy",
+                    height=200
+                )
+
+            with gr.Column(scale=2):
+                gr.Markdown("### 🎲 Option Images (8 total)")
+                with gr.Row():
+                    opt1 = gr.Image(label="Option 1", type="numpy", height=150)
+                    opt2 = gr.Image(label="Option 2", type="numpy", height=150)
+                    opt3 = gr.Image(label="Option 3", type="numpy", height=150)
+                    opt4 = gr.Image(label="Option 4", type="numpy", height=150)
+                with gr.Row():
+                    opt5 = gr.Image(label="Option 5", type="numpy", height=150)
+                    opt6 = gr.Image(label="Option 6", type="numpy", height=150)
+                    opt7 = gr.Image(label="Option 7", type="numpy", height=150)
+                    opt8 = gr.Image(label="Option 8", type="numpy", height=150)
+
+        with gr.Row():
+            solve_btn = gr.Button("🔍 Solve Challenge", variant="primary", size="lg")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### 📊 Result")
+                result_text = gr.Textbox(
+                    label="Analysis",
+                    lines=15,
+                    max_lines=20,
+                    interactive=False
+                )
+            with gr.Column(scale=1):
+                gr.Markdown("### ✅ Best Match")
+                best_image = gr.Image(label="Best Option", height=200)
+
+        # Bind button
+        solve_btn.click(
+            fn=solve_gradio,
+            inputs=[challenge_input, opt1, opt2, opt3, opt4, opt5, opt6, opt7, opt8],
+            outputs=[result_text, best_image]
+        )
+
+        gr.Markdown("""
+        ---
+        **API Endpoint:** `POST /solve` (JSON or multipart)
+        **Health Check:** `GET /health`
+        """)
+
+    return demo
+
+
+# ═══════════════════════════════════════════════════════════════
+# FASTAPI APP
+# ═══════════════════════════════════════════════════════════════
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    manager = ModelManager()
+    manager.init()
+    yield
+    logger.info("Shutting down...")
+    gc.collect()
+    torch.cuda.empty_cache()
+
+app = FastAPI(
+    title="Triple Hybrid API + Gradio: YOLO + SigLIP LoRA",
+    version="4.0.0",
+    lifespan=lifespan
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+class SolveRequest(BaseModel):
+    challenge: str
+    options: List[str]
+
+
+@app.get("/")
+async def root():
+    return {
+        "message": "Triple Hybrid API + Gradio: YOLO + SigLIP LoRA",
+        "endpoints": {
+            "solve": "POST /solve",
+            "health": "GET /health",
+            "gradio": "GET /gradio"
+        },
+        "preprocessing": "mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)"
+    }
+
+@app.get("/health")
+async def health_check():
+    manager = ModelManager()
+    return {
+        "status": "healthy",
+        "yolo_loaded": hasattr(manager, 'yolo') and manager.yolo is not None,
+        "hybrid_loaded": hasattr(manager, 'inference_model') and manager.inference_model is not None,
+        "device": CONFIG.device,
+        "num_classes": len(manager.idx_to_class) if hasattr(manager, 'idx_to_class') else 0
+    }
+
+@app.post("/solve")
+async def solve(request: Request):
+    manager = ModelManager()
+    content_type = request.headers.get("content-type", "").lower()
+    challenge_img: Optional[Image.Image] = None
+    option_imgs: List[Image.Image] = []
+
+    try:
+        if "application/json" in content_type:
+            body = await request.json()
+            challenge_input = body.get("challenge")
+            options_input = body.get("options", [])
+
+            if not challenge_input:
+                raise HTTPException(400, "'challenge' wajib diisi.")
+            if not isinstance(options_input, list) or len(options_input) != 8:
+                raise HTTPException(400, "'options' harus list berisi tepat 8 base64 images.")
+
+            challenge_img = decode_image_advanced(challenge_input, "challenge")
+            option_imgs = [decode_image_advanced(o, f"option_{i}") for i, o in enumerate(options_input)]
+
+        elif "multipart/form-data" in content_type:
+            form = await request.form()
+
+            challenge_item = form.get("challenge")
+            if challenge_item is None:
+                raise HTTPException(400, "Field 'challenge' wajib diisi.")
+
+            if hasattr(challenge_item, "read"):
+                content = await challenge_item.read()
+                challenge_img = decode_image_advanced(content, "challenge")
+            else:
+                challenge_img = decode_image_advanced(challenge_item, "challenge")
+
+            options_items = form.getlist("options") or []
+            if len(options_items) != 8:
+                options_items = []
+                for i in range(1, 9):
+                    item = form.get(f"option_{i}")
+                    if item is None:
+                        break
+                    options_items.append(item)
+
+            if len(options_items) != 8:
+                raise HTTPException(400, f"Harus ada tepat 8 options, ditemukan {len(options_items)}.")
+
+            for i, item in enumerate(options_items):
+                try:
+                    if hasattr(item, "read"):
+                        content = await item.read()
+                        option_imgs.append(decode_image_advanced(content, f"option_{i}"))
+                    else:
+                        option_imgs.append(decode_image_advanced(item, f"option_{i}"))
+                except ValueError as e:
+                    raise HTTPException(400, f"Error option_{i+1}: {str(e)}")
+        else:
+            raise HTTPException(415, "Content-Type tidak didukung.")
+
+        logger.info(f"📥 Received: challenge({challenge_img.size}) + {len(option_imgs)} options")
+
+        challenge_crop = await run_in_threadpool(manager.yolo_detect_and_crop, challenge_img)
+        result = await run_in_threadpool(manager.match_options, challenge_crop, option_imgs)
+
+        logger.info(f"✅ Result: Option {result['predicted_index']} (conf: {result['confidence']:.4f})")
+
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        return {
+            "success": True,
+            "predicted_index": result["predicted_index"],
+            "confidence": result["confidence"],
+            "similarity_score": result["similarity_score"],
+            "predicted_class": result["predicted_class"],
+            "all_scores": result["all_scores"],
+            "debug": {
+                "challenge_crop_size": challenge_crop.size,
+                "device": CONFIG.device,
+                "preprocessing": "mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)"
+            }
+        }
+
+    except HTTPException:
+        raise
+    except ValueError as e:
+        raise HTTPException(400, str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}", exc_info=True)
+        raise HTTPException(500, f"Internal error: {str(e)}")
+
+
+# ═══════════════════════════════════════════════════════════════
+# MOUNT GRADIO TO FASTAPI
+# ═══════════════════════════════════════════════════════════════
+gradio_app = build_gradio_interface()
+app = gr.mount_gradio_app(app, gradio_app, path="/gradio")
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)