RimsJ/batik-classifier

1"""2Gradio UI application for Batik Classification3Optimized for Hugging Face Spaces deployment4"""5import gradio as gr6import torch7import torch.nn as nn8from torchvision import transforms, models9from PIL import Image10import json11import numpy as np12from typing import Tuple, Dict13from huggingface_hub import hf_hub_download14import os15 16# Global variables17model = None18class_names = []19device = torch.device("cuda" if torch.cuda.is_available() else "cpu")20transform = None21 22 23def load_model():24    global model, class_names, transform25 26    try:27        # Load model configuration28        with open('model_config.json', 'r') as f:29            config = json.load(f)30 31        num_classes = config['num_classes']32        class_names = config['class_names']33        image_size = config.get('image_size', 224)34 35        # Initialize VGG16 model36        model = models.vgg16(weights=None)37        # Modify classifier to match saved model architecture38        model.classifier[3] = nn.Linear(4096, num_classes)39        model.classifier = nn.Sequential(*list(model.classifier.children())[:4])40 41        # Download model from Hugging Face Hub42        print("📥 Downloading model from Hugging Face Hub...")43        model_path = hf_hub_download(44            repo_id="RimsJ/Batik-Classifier",45            filename="vgg16_batik_best.pth"46        )47 48        # Load trained weights49        checkpoint = torch.load(model_path, map_location=device)50 51        # Extract state_dict52        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:53            state_dict = checkpoint['model_state_dict']54        else:55            state_dict = checkpoint56 57        # Remove '_orig_mod.' prefix if present58        new_state_dict = {}59        for key, value in state_dict.items():60            if key.startswith('_orig_mod.'):61                new_key = key.replace('_orig_mod.', '')62                new_state_dict[new_key] = value63            else:64                new_state_dict[key] = value65 66        model.load_state_dict(new_state_dict)67        model = model.to(device)68        model.eval()69 70        # Define image preprocessing71        transform = transforms.Compose([72            transforms.Resize((image_size, image_size)),73            transforms.ToTensor(),74            transforms.Normalize(mean=[0.485, 0.456, 0.406],75                               std=[0.229, 0.224, 0.225])76        ])77 78        print(f"✅ Model loaded successfully on {device}")79        print(f"📊 Number of classes: {num_classes}")80 81    except Exception as e:82        print(f"❌ Error loading model: {str(e)}")83        raise84 85 86def predict_image(image):87    """88    Predict batik class from image89 90    Args:91        image: PIL Image92 93    Returns:94        Tuple of (top_k_dict, formatted_text)95    """96    global model, transform, class_names97 98    try:99        if image is None:100            return None, "❌ Silakan upload gambar batik terlebih dahulu"101 102        if model is None:103            return None, "❌ Model belum dimuat. Silakan refresh halaman."104 105        # Convert to RGB if needed106        if image.mode != 'RGB':107            image = image.convert('RGB')108 109        # Transform and predict110        input_tensor = transform(image).unsqueeze(0).to(device)111 112        with torch.no_grad():113            outputs = model(input_tensor)114            probabilities = torch.nn.functional.softmax(outputs, dim=1)115            top_probs, top_indices = torch.topk(probabilities, min(5, len(class_names)), dim=1)116 117        # Get top prediction118        predicted_class = class_names[top_indices[0][0].item()]119        confidence = top_probs[0][0].item() * 100120 121        # Format top-5 results122        results = {}123        for i in range(min(5, len(class_names))):124            class_name = class_names[top_indices[0][i].item()]125            conf = top_probs[0][i].item()126            results[class_name] = float(conf)127 128        # Format output text129        result_text = f"""130## 🎯 Hasil Prediksi131 132**Motif Batik:** `{predicted_class}`133**Confidence:** `{confidence:.2f}%`134 135---136 137### 📊 Top 5 Prediksi:138"""139 140        for idx, (class_name, conf) in enumerate(list(results.items())[:5], 1):141            bar = "█" * int(conf * 20)142            result_text += f"\n{idx}. **{class_name}** - {conf*100:.2f}%  \n   {bar}"143 144        return results, result_text145 146    except Exception as e:147        import traceback148        traceback.print_exc()149        return None, f"❌ Error: {str(e)}"150 151 152# Load model at startup153print("🔄 Loading model...")154load_model()155print("✅ Model ready!")156 157# Create Gradio interface158with gr.Blocks(159    title="Batik Classification",160    theme=gr.themes.Soft(),161    css=".gradio-container {max-width: 1200px; margin: auto;}"162) as demo:163 164    gr.Markdown("""165    # 🎨 Klasifikasi Motif Batik Indonesia166 167    Upload gambar batik untuk mengetahui motif dan asalnya!168    **Total 111 motif batik** dari berbagai daerah di Indonesia 🇮🇩169    """)170 171    with gr.Row():172        with gr.Column(scale=1):173            input_image = gr.Image(174                type="pil",175                label="📤 Upload Gambar Batik",176                height=400177            )178            predict_btn = gr.Button(179                "🔍 Prediksi Motif Batik",180                variant="primary",181                size="lg"182            )183 184            gr.Markdown("""185            ### 💡 Tips:186            - Gunakan gambar dengan kualitas baik187            - Pastikan motif batik terlihat jelas188            - Format: JPG, PNG, JPEG189            """)190 191        with gr.Column(scale=1):192            output_text = gr.Markdown(label="Hasil Prediksi")193            output_label = gr.Label(194                label="📊 Confidence Score",195                num_top_classes=5196            )197 198    # Event handler199    predict_btn.click(200        fn=predict_image,201        inputs=input_image,202        outputs=[output_label, output_text]203    )204 205    gr.Markdown("""206    ---207    ### 📋 Tentang Model208    - **Dataset:** 111 Motif Batik Indonesia209    - **Kategori:** Batik dari Jawa Tengah, Jawa Timur, Jawa Barat, Bali, Jakarta, Kalimantan, Lampung210 211    ### 🎨 Contoh Motif:212    Parang Kusumo, Megamendung, Kawung, Truntum, Semarangan, dan banyak lagi!213 214    ---215    **Made with ❤️ for Indonesian Batik Heritage**216    """)217 218 219# Launch220if __name__ == "__main__":221    demo.launch(server_name="0.0.0.0", server_port=7860)222