Update app.py

app.py

@@ -1,29 +1,15 @@
-#!/usr/bin/env python3
-"""
-Gradio web interface for artifact classification
-"""
-
-import os
-# Fix SSL issue on Windows
-os.environ['SSL_CERT_FILE'] = ''
-
 import gradio as gr
 import torch
-import torch.nn as nn
-from torchvision import transforms
 from PIL import Image
-import numpy as np
-import os
-import json
-from pathlib import Path
-
-# Define the model architecture directly (standalone)
-import torch
+import torchvision.transforms as transforms
 import torch.nn as nn
 from torchvision import models
+from typing import Dict, Tuple
+import os
+
 
 class MultiOutputModel(nn.Module):
-    """Multi-output model for artifact classification"""
+    """Multi-output model for artifact classification (matches UI)"""
 
     def __init__(self, num_object_classes, num_material_classes, hidden_size=512):
         super(MultiOutputModel, self).__init__()
@@ -34,21 +20,17 @@ class MultiOutputModel(nn.Module):
         self.backbone = nn.Sequential(*list(self.backbone.children())[:-1])
 
         # Freeze early layers for transfer learning
-        for param in list(self.backbone.parameters())[:-2]:
+        for param in list(self.backbone.parameters())[:-4]:  # Unfreeze more layers for better fine-tuning
             param.requires_grad = False
 
         # Classification heads for each attribute
         self.object_classifier = nn.Linear(2048, num_object_classes)
         self.material_classifier = nn.Linear(2048, num_material_classes)
 
-        # Dropout for regularization
-        self.dropout = nn.Dropout(0.3)
-
     def forward(self, x):
         # Extract features using backbone
         features = self.backbone(x)
         features = features.view(features.size(0), -1)
-        features = self.dropout(features)
 
         # Get predictions for each attribute
         object_pred = self.object_classifier(features)
@@ -59,270 +41,211 @@ class MultiOutputModel(nn.Module):
             'material': material_pred,
         }
 
-print("MultiOutputModel class defined directly in app (standalone)")
 
-class ArtifactClassifier:
-    def __init__(self, model_path="train/outputs/best_model.pth"):
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        print(f"Using device: {self.device}")
 
-        # Try to load from local file first, then from HuggingFace
-        self.model = self.load_model(model_path)
-        self.model.to(self.device)
-        self.model.eval()
 
-        # Set up transforms (same as training)
-        self.transform = transforms.Compose([
-            transforms.Resize((224, 224)),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-        ])
+def load_model(model_path: str) -> Tuple[torch.nn.Module, Dict[str, Dict[int, str]]]:
+    """Load the model from checkpoint and return model and label mappings."""
+    print(f"Loading model from {model_path}...")
+    checkpoint = torch.load(model_path, map_location="cpu")
 
-        # Load label mappings if available
-        self.label_mappings = self.load_label_mappings()
-        print("Model loaded successfully!")
+    # Get label mappings to determine number of classes
+    label_mappings = checkpoint.get('label_mappings', {})
+    num_object_classes = len(label_mappings.get('object_name', {}))
+    num_material_classes = len(label_mappings.get('material', {}))
 
-    def load_model(self, model_path):
-        """Load the trained model from local file or HuggingFace Hub"""
-        # First try to load from local file
-        if os.path.exists(model_path):
-            print(f"Loading model from local file: {model_path}")
-            return self._load_model_from_path(model_path)
-
-        # If local file doesn't exist, try to download from HuggingFace
-        print(f"Local model not found, downloading from HuggingFace...")
-        try:
-            return self._load_model_from_hub()
-        except Exception as e:
-            print(f"Failed to download from HuggingFace: {e}")
-            print("Falling back to local model creation...")
-            return self._create_model_with_defaults()
-
-    def _load_model_from_path(self, model_path):
-        """Load model from local file"""
-        checkpoint = torch.load(model_path, map_location=self.device)
-
-        # Get label mappings to determine number of classes
-        label_mappings = checkpoint.get('label_mappings', {})
-        num_object_classes = len(label_mappings.get('object_name', {}))
-        num_material_classes = len(label_mappings.get('material', {}))
-
-        if num_object_classes == 0:
-            print("Warning: No label mappings found, using fallback class counts")
-            num_object_classes, num_material_classes = 1018, 192
-
-        # Create model
-        model = MultiOutputModel(num_object_classes, num_material_classes)
+    if num_object_classes == 0:
+        print("Warning: No label mappings found, using fallback class counts")
+        num_object_classes, num_material_classes = 1018, 192
+
+    # Check model type based on state_dict keys to determine which architecture to use
+    model_state_dict = checkpoint.get('model_state_dict', {})
+    state_dict_keys = set(model_state_dict.keys())
+
+    # Only support v1 model (MultiOutputModel) with ResNet backbone
+    print(f"Loading v1 model (MultiOutputModel) with ResNet backbone")
+    model = MultiOutputModel(num_object_classes, num_material_classes)
+
+    # Load state dict
+    if 'model_state_dict' in checkpoint:
         model.load_state_dict(checkpoint['model_state_dict'])
+    else:
+        print("Warning: No model_state_dict found in checkpoint")
 
-        return model
+    # Create reverse mappings (id2label)
+    reverse_mappings = {}
+    for attr, mapping in label_mappings.items():
+        reverse_mappings[attr] = {int(v): str(k) for k, v in mapping.items()}
+        print(f"Loaded {attr} mappings: {len(reverse_mappings[attr])} classes")
 
-    def _load_model_from_hub(self):
-        """Download and load model from HuggingFace Hub"""
-        try:
-            from huggingface_hub import hf_hub_download
+    return model, reverse_mappings
 
-            print("Downloading model from HuggingFace Hub...")
-            model_file = hf_hub_download(
-                repo_id="SpyC0der77/artifact-classification-model",
-                filename="best_model.pth"
-            )
 
-            print(f"Model downloaded to: {model_file}")
-            return self._load_model_from_path(model_file)
+def run_inference(model: torch.nn.Module, pixel_values: torch.Tensor, device: str) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Run inference on pixel_values and return predictions and confidences for both object_name and material."""
+    model.eval()
+    model.to(device)
+    pixel_values = pixel_values.to(device)
 
-        except Exception as e:
-            print(f"Error downloading from HuggingFace: {e}")
-            raise
+    with torch.no_grad():
+        outputs = model(pixel_values)
 
-    def _create_model_with_defaults(self):
-        """Create model with default parameters when no model is available"""
-        print("Creating model with default parameters...")
-        print("Note: This model won't have the trained weights!")
+        # Handle different output formats
+        if isinstance(outputs, dict):
+            # Multi-output model format
+            if 'object_name' in outputs and 'material' in outputs:
+                logits_obj = outputs['object_name']
+                logits_mat = outputs['material']
+            else:
+                raise ValueError("Expected 'object_name' and 'material' in model outputs")
+        else:
+            raise ValueError("Expected dict output with 'object_name' and 'material' keys")
 
-        # Use default class counts
-        num_object_classes, num_material_classes = 1018, 192
+        preds_obj = torch.argmax(logits_obj, dim=-1)
+        probs_obj = torch.softmax(logits_obj, dim=-1)
+        max_probs_obj = torch.max(probs_obj, dim=-1)[0]
 
-        # Create model
-        model = MultiOutputModel(num_object_classes, num_material_classes)
-
-        return model
-
-    def load_label_mappings(self):
-        """Load label mappings for decoding predictions"""
-        # First try local model
-        model_path = "train/outputs/best_model.pth"
-        if os.path.exists(model_path):
-            try:
-                checkpoint = torch.load(model_path, map_location='cpu')
-                mappings = checkpoint.get('label_mappings', {})
-
-                # Create reverse mappings
-                reverse_mappings = {}
-                for attr, mapping in mappings.items():
-                    reverse_mappings[attr] = {v: k for k, v in mapping.items()}
-
-                return reverse_mappings
-            except Exception as e:
-                print(f"Could not load local label mappings: {e}")
-
-        # Try to download from HuggingFace
-        try:
-            print("Downloading label mappings from HuggingFace...")
-            from huggingface_hub import hf_hub_download
-
-            mappings_file = hf_hub_download(
-                repo_id="SpyC0der77/artifact-classification-model",
-                filename="best_model.pth"  # Contains the mappings
-            )
-
-            checkpoint = torch.load(mappings_file, map_location='cpu')
-            mappings = checkpoint.get('label_mappings', {})
-
-            # Create reverse mappings
-            reverse_mappings = {}
-            for attr, mapping in mappings.items():
-                reverse_mappings[attr] = {v: k for k, v in mapping.items()}
-
-            print(f"Loaded {len(reverse_mappings)} label mappings from HuggingFace")
-            return reverse_mappings
-
-        except Exception as e:
-            print(f"Could not load label mappings from HuggingFace: {e}")
-
-        return {}
-
-    def predict(self, image):
-        """Make prediction on uploaded image"""
-        try:
-            # Convert to PIL Image if needed
-            if isinstance(image, np.ndarray):
-                image = Image.fromarray(image).convert('RGB')
-            elif not isinstance(image, Image.Image):
-                image = Image.open(image).convert('RGB')
-
-            # Apply transforms
-            image_tensor = self.transform(image).unsqueeze(0).to(self.device)
-
-            # Make prediction
-            with torch.no_grad():
-                outputs = self.model(image_tensor)
-
-            # Process results
-            results = {}
-            for attr in ['object_name', 'material']:
-                if attr in outputs:
-                    # Get probabilities and prediction
-                    probs = torch.softmax(outputs[attr], dim=1)
-                    confidence, predicted_idx = torch.max(probs, dim=1)
-
-                    pred_class = predicted_idx.item()
-                    conf = confidence.item()
-
-                    # Convert to label name
-                    if attr in self.label_mappings and pred_class in self.label_mappings[attr]:
-                        pred_label = self.label_mappings[attr][pred_class]
-                    else:
-                        pred_label = f"Class_{pred_class}"
-
-                    results[attr] = {
-                        'prediction': pred_label,
-                        'confidence': conf,
-                        'class_id': pred_class
-                    }
-
-            return results
-
-        except Exception as e:
-            return {"error": str(e)}
-
-# Global classifier instance
-classifier = None
-
-def classify_image(image):
-    """Gradio interface function"""
-    global classifier
-
-    if classifier is None:
-        return "Error: Model not loaded. Please restart the app."
+        preds_mat = torch.argmax(logits_mat, dim=-1)
+        probs_mat = torch.softmax(logits_mat, dim=-1)
+        max_probs_mat = torch.max(probs_mat, dim=-1)[0]
 
-    try:
-        results = classifier.predict(image)
+    return preds_obj.cpu(), max_probs_obj.cpu(), preds_mat.cpu(), max_probs_mat.cpu()
 
-        if "error" in results:
-            return f"Prediction failed: {results['error']}"
 
-        # Format results
-        output = "PREDICTION RESULTS\n\n"
+# Global variables for model and label mappings
+model = None
+label_mappings = None
+device = None
+
+def preprocess_image(image: Image.Image) -> torch.Tensor:
+    """Preprocess image for model inference."""
+    # Define transforms
+    transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    # Apply transforms
+    image = image.convert('RGB')
+    tensor = transform(image).unsqueeze(0)  # Add batch dimension
+
+    return tensor
+
+def predict_artifact(image: Image.Image) -> tuple[str, float, str, float]:
+    """Predict object and material from image."""
+    global model, label_mappings, device
 
-        for attr, result in results.items():
-            status = "OK" if result['confidence'] > 0.5 else "LOW"
-            output += f"{status} {attr.upper()}: {result['prediction']}\n"
-            output += f"   Confidence: {result['confidence']:.3f}\n"
-            output += f"   Class ID: {result['class_id']}\n\n"
+    if model is None:
+        raise ValueError("Model not loaded. Please restart the application.")
 
-        # Overall confidence
-        confidences = [r['confidence'] for r in results.values()]
-        avg_confidence = sum(confidences) / len(confidences)
-        output += f"Average Confidence: {avg_confidence:.3f}"
+    # Preprocess image
+    pixel_values = preprocess_image(image)
 
-        return output
+    # Run inference
+    preds_obj, confs_obj, preds_mat, confs_mat = run_inference(model, pixel_values, device)
+
+    # Get predictions
+    object_pred_id = preds_obj[0].item()
+    material_pred_id = preds_mat[0].item()
+    object_conf = confs_obj[0].item()
+    material_conf = confs_mat[0].item()
+
+    # Convert IDs to labels
+    object_name = label_mappings['object_name'].get(object_pred_id, f"class_{object_pred_id}")
+    material_name = label_mappings['material'].get(material_pred_id, f"class_{material_pred_id}")
+
+    return object_name, object_conf, material_name, material_conf
+
+def gradio_predict(image):
+    """Gradio interface function."""
+    if image is None:
+        return "Please upload an image", "", "", ""
+
+    try:
+        object_name, object_conf, material_name, material_conf = predict_artifact(image)
+
+        # Format results
+        object_result = f"**{object_name}** ({object_conf:.1%} confidence)"
+        material_result = f"**{material_name}** ({material_conf:.1%} confidence)"
+
+        return object_result, material_result, f"{object_conf:.3f}", f"{material_conf:.3f}"
 
     except Exception as e:
-        return f"Error during prediction: {str(e)}"
+        return f"Error: {str(e)}", "", "", ""
 
-def create_interface():
-    """Create and launch the Gradio interface"""
-    global classifier
+def load_model_on_startup():
+    """Load model when the application starts."""
+    global model, label_mappings, device
+
+    # Set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Load model
+    model_path = "model/v1/best_model.pth"
+    if not os.path.exists(model_path):
+        print(f"Warning: Model file not found at {model_path}")
+        print("Please ensure the model file exists before running the application.")
+        return
 
-    # Initialize classifier
     try:
-        print("Loading model...")
-        classifier = ArtifactClassifier()
+        model, label_mappings = load_model(model_path)
         print("Model loaded successfully!")
+        print(f"Object classes: {len(label_mappings.get('object_name', {}))}")
+        print(f"Material classes: {len(label_mappings.get('material', {}))}")
     except Exception as e:
-        print(f"Failed to load model: {e}")
-        return
+        print(f"Error loading model: {e}")
+
+# Load model on startup
+load_model_on_startup()
+
+# Create Gradio interface
+with gr.Blocks(title="Artifact Classification v1", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🏺 Artifact Classification Model v1")
+    gr.Markdown("Upload an image of an artifact to classify its **object type** and **material composition**.")
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Upload Artifact Image", type="pil")
+            submit_btn = gr.Button("🔍 Classify Artifact", variant="primary")
 
-    # Create interface
-    interface = gr.Interface(
-        fn=classify_image,
-        inputs=gr.Image(type="pil", label="Upload Artifact Image"),
-        outputs=gr.Textbox(label="Classification Results", lines=10),
-        title="Artifact Classification",
-        description="""
-        Upload an image of an archaeological artifact to get AI-powered classification!
-
-        Features:
-        - Object type identification (coin, vase, statue, etc.)
-        - Material classification (gold, silver, pottery, etc.)
-        - Confidence scores for each prediction
-        - GPU-accelerated processing (if available)
-        - Auto-downloads model from HuggingFace Hub
-        - Completely standalone - no training code needed
-
-        Supported formats: JPG, PNG, JPEG
-        """,
-        article="""
-        How to use:
-        1. Upload an artifact image using the file picker
-        2. Click "Submit" to run classification
-        3. View results with confidence scores and predictions
-
-        Model trained on: British Museum artifact dataset
-        Accuracy: ~71% for objects, ~62% for materials
-        """,
-        examples=[]
+        with gr.Column():
+            gr.Markdown("### 📊 Classification Results")
+
+            object_output = gr.Markdown(label="**Object Type**")
+            material_output = gr.Markdown(label="**Material**")
+
+            with gr.Accordion("📈 Confidence Scores", open=False):
+                object_conf = gr.Textbox(label="Object Confidence", interactive=False)
+                material_conf = gr.Textbox(label="Material Confidence", interactive=False)
+
+    # Connect the interface
+    submit_btn.click(
+        fn=gradio_predict,
+        inputs=image_input,
+        outputs=[object_output, material_output, object_conf, material_conf]
     )
 
-    # Launch
-    print("Starting Gradio interface...")
-    interface.launch(
-        server_name="0.0.0.0",  # Allow external connections
-        server_port=7860,
-        share=False,  # Set to True for public link
-        debug=False
+    # Example images
+    gr.Examples(
+        examples=[
+            # You can add example image paths here if available
+        ],
+        inputs=image_input,
+        outputs=[object_output, material_output, object_conf, material_conf],
+        fn=gradio_predict,
+        cache_examples=False
     )
 
+    gr.Markdown("""
+    ### ℹ️ About
+    This model uses a ResNet-50 backbone to classify museum artifacts into object types (vase, statue, pottery, etc.)
+    and material compositions (ceramic, bronze, stone, etc.).
+
+    **Model**: MultiOutputModel with ResNet-50 backbone
+    **Training Data**: Oriental Museum artifacts dataset
+    """)
+
 if __name__ == "__main__":
-    create_interface()
+    demo.launch()