RimsJ/batik-classifier

1"""2Gradio UI application for Batik Classification using VGG16 model3"""4import gradio as gr5import torch6import torch.nn as nn7from torchvision import transforms, models8from PIL import Image9import json10import numpy as np11from typing import Tuple, List12 13# Global variables14model = None15class_names = []16device = torch.device("cuda" if torch.cuda.is_available() else "cpu")17transform = None18 19 20def load_model():21    """Load VGG16 model and configuration"""22    global model, class_names, transform23 24    try:25        # Load model configuration26        with open('model_config.json', 'r') as f:27            config = json.load(f)28 29        num_classes = config['num_classes']30        class_names = config['class_names']31        image_size = config.get('image_size', 224)32 33        # Initialize VGG16 model34        model = models.vgg16(weights=None)35        # Modify classifier to match saved model architecture36        # The saved model has classifier.3 as output layer (111 classes)37        model.classifier[3] = nn.Linear(4096, num_classes)38        # Remove layers after classifier.339        model.classifier = nn.Sequential(*list(model.classifier.children())[:4])40 41        # Load trained weights42        checkpoint = torch.load('models/vgg16_batik_best.pth', map_location=device)43 44        # Check if checkpoint is a dict with 'model_state_dict' key or direct state_dict45        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:46            state_dict = checkpoint['model_state_dict']47        else:48            state_dict = checkpoint49 50        # Remove '_orig_mod.' prefix if present (from torch.compile)51        new_state_dict = {}52        for key, value in state_dict.items():53            if key.startswith('_orig_mod.'):54                new_key = key.replace('_orig_mod.', '')55                new_state_dict[new_key] = value56            else:57                new_state_dict[key] = value58 59        model.load_state_dict(new_state_dict)60        model = model.to(device)61        model.eval()        # Define image preprocessing62        transform = transforms.Compose([63            transforms.Resize((image_size, image_size)),64            transforms.ToTensor(),65            transforms.Normalize(mean=[0.485, 0.456, 0.406],66                               std=[0.229, 0.224, 0.225])67        ])68 69        print(f"✅ Model loaded successfully on {device}")70        print(f"📊 Number of classes: {num_classes}")71 72    except Exception as e:73        print(f"❌ Error loading model: {str(e)}")74        raise75 76 77def predict_single(image: Image.Image) -> Tuple[str, float]:78    """79    Predict single class for an image80 81    Args:82        image: PIL Image83 84    Returns:85        Tuple of (predicted_class, confidence)86    """87    try:88        # Preprocess image89        if image is None:90            return "Error: No image provided", 0.091 92        # Convert to RGB if needed93        if image.mode != 'RGB':94            image = image.convert('RGB')95 96        # Transform and add batch dimension97        input_tensor = transform(image).unsqueeze(0).to(device)98 99        # Make prediction100        with torch.no_grad():101            outputs = model(input_tensor)102            probabilities = torch.nn.functional.softmax(outputs, dim=1)103            confidence, predicted = torch.max(probabilities, 1)104 105        predicted_class = class_names[predicted.item()]106        confidence_score = confidence.item() * 100  # Convert to percentage107 108        return predicted_class, confidence_score109 110    except Exception as e:111        return f"Error: {str(e)}", 0.0112 113 114def predict_top_k(image: Image.Image, k: int = 5) -> dict:115    """116    Predict top-k classes for an image117 118    Args:119        image: PIL Image120        k: Number of top predictions121 122    Returns:123        Dictionary of class names and their confidence scores124    """125    try:126        # Preprocess image127        if image is None:128            return {"Error": 1.0}129 130        # Convert to RGB if needed131        if image.mode != 'RGB':132            image = image.convert('RGB')133 134        # Transform and add batch dimension135        input_tensor = transform(image).unsqueeze(0).to(device)136 137        # Make prediction138        with torch.no_grad():139            outputs = model(input_tensor)140            probabilities = torch.nn.functional.softmax(outputs, dim=1)141            top_probs, top_indices = torch.topk(probabilities, min(k, len(class_names)), dim=1)142 143        # Format results as dictionary for Gradio144        results = {}145        for i in range(min(k, len(class_names))):146            class_name = class_names[top_indices[0][i].item()]147            confidence = top_probs[0][i].item()148            results[class_name] = float(confidence)149 150        return results151 152    except Exception as e:153        return {"Error": f"{str(e)}"}154 155 156def format_prediction(image: Image.Image) -> Tuple[str, dict]:157    """158    Format prediction output for Gradio interface159 160    Args:161        image: PIL Image162 163    Returns:164        Tuple of (formatted_text, top_k_dict)165    """166    try:167        if image is None:168            return "❌ Silakan upload gambar batik terlebih dahulu", {}169 170        # Get single prediction171        predicted_class, confidence = predict_single(image)172 173        # Get top-5 predictions174        top_k_results = predict_top_k(image, k=5)175 176        # Format main result177        result_text = f"""178## 🎯 Hasil Prediksi179 180**Motif Batik:** `{predicted_class}`181**Confidence:** `{confidence:.2f}%`182 183---184 185### 📊 Top 5 Prediksi:186"""187 188        for idx, (class_name, conf) in enumerate(list(top_k_results.items())[:5], 1):189            bar = "█" * int(conf * 20)  # Simple bar visualization190            result_text += f"\n{idx}. **{class_name}** - {conf*100:.2f}%  \n   {bar}"191 192        return result_text, top_k_results193 194    except Exception as e:195        return f"❌ Error: {str(e)}", {}196 197 198def get_model_info() -> str:199    """Get model information"""200    info = f"""201### 📋 Informasi Model202 203- **Arsitektur:** VGG16204- **Device:** {device}205- **Jumlah Kelas:** {len(class_names)}206- **Status:** ✅ Model siap digunakan207 208### 🎨 Kategori Batik:209Total {len(class_names)} motif batik dari berbagai daerah di Indonesia210"""211    return info212 213 214# Load model at startup215load_model()216 217# Create Gradio interface218with gr.Blocks(title="Batik Classification - VGG16", theme=gr.themes.Soft()) as demo:219 220    gr.Markdown("""221    # 🎨 Klasifikasi Motif Batik Indonesia222    ### Menggunakan Model VGG16 Deep Learning223 224    Upload gambar batik untuk mengetahui motif dan asalnya!225    """)226 227    with gr.Tabs():228 229        # Tab 1: Single Prediction230        with gr.Tab("🖼️ Prediksi Tunggal"):231            with gr.Row():232                with gr.Column():233                    input_image = gr.Image(234                        type="pil",235                        label="Upload Gambar Batik",236                        height=400237                    )238                    predict_btn = gr.Button("🔍 Prediksi", variant="primary", size="lg")239 240                    gr.Examples(241                        examples=[],  # Add example images if available242                        inputs=input_image,243                        label="Contoh Gambar (jika tersedia)"244                    )245 246                with gr.Column():247                    output_text = gr.Markdown(label="Hasil Prediksi")248                    output_label = gr.Label(249                        label="Top 5 Prediksi",250                        num_top_classes=5251                    )252 253            predict_btn.click(254                fn=format_prediction,255                inputs=input_image,256                outputs=[output_text, output_label]257            )258 259        # Tab 2: Batch Prediction260        with gr.Tab("📁 Prediksi Batch"):261            gr.Markdown("### Upload multiple gambar batik sekaligus")262 263            batch_input = gr.File(264                file_count="multiple",265                file_types=["image"],266                label="Upload Gambar (Multiple)"267            )268            batch_btn = gr.Button("🔍 Prediksi Semua", variant="primary")269            batch_output = gr.Dataframe(270                headers=["Filename", "Predicted Class", "Confidence (%)"],271                label="Hasil Prediksi Batch"272            )273 274            def predict_batch(files):275                """Predict multiple images"""276                if files is None or len(files) == 0:277                    return []278 279                results = []280                for file in files:281                    try:282                        image = Image.open(file.name)283                        pred_class, confidence = predict_single(image)284                        results.append([file.name.split('/')[-1], pred_class, f"{confidence:.2f}"])285                    except Exception as e:286                        results.append([file.name.split('/')[-1], "Error", str(e)])287 288                return results289 290            batch_btn.click(291                fn=predict_batch,292                inputs=batch_input,293                outputs=batch_output294            )295 296        # Tab 3: Model Info297        with gr.Tab("ℹ️ Info Model"):298            gr.Markdown(get_model_info())299 300            with gr.Accordion("📜 Daftar Semua Kelas Batik", open=False):301                class_list = "\n".join([f"{i+1}. {name}" for i, name in enumerate(class_names)])302                gr.Textbox(303                    value=class_list,304                    label=f"Total {len(class_names)} Kelas",305                    lines=20,306                    max_lines=30307                )308 309    gr.Markdown("""310    ---311    ### 📝 Cara Penggunaan:312    1. **Prediksi Tunggal:** Upload satu gambar batik dan klik tombol Prediksi313    2. **Prediksi Batch:** Upload beberapa gambar sekaligus untuk prediksi massal314    3. **Info Model:** Lihat informasi lengkap tentang model dan daftar kelas315 316    ### 💡 Tips:317    - Gunakan gambar dengan kualitas yang baik untuk hasil terbaik318    - Pastikan gambar menunjukkan motif batik dengan jelas319    - Model mendukung format JPG, PNG, dan format gambar umum lainnya320    """)321 322 323# Launch the app324if __name__ == "__main__":325    try:326        demo.launch(327            server_name="127.0.0.1",328            server_port=7860,329            share=False,  # Ubah ke True jika mau public link330            inbrowser=True,331            quiet=False332        )333    except Exception as e:334        print(f"Error launching Gradio: {e}")335        # Fallback: try simpler launch336        demo.launch()337