Upload app.py

app.py

@@ -0,0 +1,405 @@
+import streamlit as st
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+from PIL import Image
+import matplotlib.pyplot as plt
+import numpy as np
+import requests
+import io
+from timm import create_model
+
+# Set page config
+st.set_page_config(
+    page_title="Sports Ball Classifier",
+    page_icon="🏀",
+    layout="wide"
+)
+
+# Custom ConvNeXt model definition (in case the saved model uses a different architecture)
+class ConvNeXtBlock(nn.Module):
+    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), 
+                                requires_grad=True) if layer_scale_init_value > 0 else None
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
+        x = input + x
+        return x
+
+class CustomConvNeXt(nn.Module):
+    def __init__(self, num_classes=15):
+        super().__init__()
+        self.stem = nn.Sequential(
+            nn.Conv2d(3, 96, kernel_size=4, stride=4),
+            nn.LayerNorm([96, 56, 56], eps=1e-6)
+        )
+        
+        # Stage 1
+        self.stage1 = nn.Sequential(*[ConvNeXtBlock(96) for _ in range(3)])
+        
+        # Downsample 1
+        self.downsample1 = nn.Sequential(
+            nn.LayerNorm([96, 56, 56], eps=1e-6),
+            nn.Conv2d(96, 192, kernel_size=2, stride=2)
+        )
+        
+        # Stage 2
+        self.stage2 = nn.Sequential(*[ConvNeXtBlock(192) for _ in range(3)])
+        
+        # Downsample 2
+        self.downsample2 = nn.Sequential(
+            nn.LayerNorm([192, 28, 28], eps=1e-6),
+            nn.Conv2d(192, 384, kernel_size=2, stride=2)
+        )
+        
+        # Stage 3
+        self.stage3 = nn.Sequential(*[ConvNeXtBlock(384) for _ in range(9)])
+        
+        # Downsample 3
+        self.downsample3 = nn.Sequential(
+            nn.LayerNorm([384, 14, 14], eps=1e-6),
+            nn.Conv2d(384, 768, kernel_size=2, stride=2)
+        )
+        
+        # Stage 4
+        self.stage4 = nn.Sequential(*[ConvNeXtBlock(768) for _ in range(3)])
+        
+        # Head
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.norm = nn.LayerNorm(768, eps=1e-6)
+        self.head = nn.Linear(768, num_classes)
+
+    def forward(self, x):
+        x = self.stem(x)
+        x = self.stage1(x)
+        x = self.downsample1(x)
+        x = self.stage2(x)
+        x = self.downsample2(x)
+        x = self.stage3(x)
+        x = self.downsample3(x)
+        x = self.stage4(x)
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.norm(x)
+        x = self.head(x)
+        return x
+
+# Cache the model loading to avoid reloading on every interaction
+@st.cache_resource
+def load_model():
+    """Load the pre-trained ViT model for sports ball classification"""
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    try:
+        # Download model weights from Hugging Face
+        model_url = "https://huggingface.co/Alamgirapi/sports-ball-convnext-classifier/resolve/main/model.pth"
+        response = requests.get(model_url)
+        if response.status_code != 200:
+            raise Exception(f"Failed to download model: HTTP {response.status_code}")
+            
+        model_state = torch.load(io.BytesIO(response.content), map_location=device)
+        
+        # Inspect the state dict to understand the model structure
+        sample_keys = list(model_state.keys())[:10]
+        
+        # Try Vision Transformer models (this is likely what was used)
+        vit_models_to_try = [
+            ("vit_base_patch16_224", lambda: create_model('vit_base_patch16_224', pretrained=False, num_classes=15)),
+            ("vit_small_patch16_224", lambda: create_model('vit_small_patch16_224', pretrained=False, num_classes=15)),
+            ("vit_tiny_patch16_224", lambda: create_model('vit_tiny_patch16_224', pretrained=False, num_classes=15)),
+            ("vit_large_patch16_224", lambda: create_model('vit_large_patch16_224', pretrained=False, num_classes=15)),
+            ("vit_base_patch32_224", lambda: create_model('vit_base_patch32_224', pretrained=False, num_classes=15)),
+        ]
+        
+        st.info("Trying Vision Transformer (ViT) models...")
+        for model_name, model_func in vit_models_to_try:
+            try:
+                model = model_func()
+                model.load_state_dict(model_state)
+                model.eval()
+                model.to(device)
+                st.success(f"✅ Successfully loaded model using: {model_name}")
+                return model, device
+            except Exception as e:
+                st.warning(f"❌ Failed to load with {model_name}: {str(e)[:100]}...")
+                continue
+        
+        # Try ConvNeXt models as fallback
+        convnext_models_to_try = [
+            ("convnext_tiny", lambda: create_model('convnext_tiny', pretrained=False, num_classes=15)),
+            ("convnext_small", lambda: create_model('convnext_small', pretrained=False, num_classes=15)),
+            ("convnext_base", lambda: create_model('convnext_base', pretrained=False, num_classes=15)),
+        ]
+        
+        st.info("Trying ConvNeXt models as fallback...")
+        for model_name, model_func in convnext_models_to_try:
+            try:
+                model = model_func()
+                model.load_state_dict(model_state)
+                model.eval()
+                model.to(device)
+                st.success(f"✅ Successfully loaded model using: {model_name}")
+                return model, device
+            except Exception as e:
+                st.warning(f"❌ Failed to load with {model_name}: {str(e)[:100]}...")
+                continue
+        
+        # Try other common models
+        other_models_to_try = [
+            ("resnet50", lambda: create_model('resnet50', pretrained=False, num_classes=15)),
+            ("efficientnet_b0", lambda: create_model('efficientnet_b0', pretrained=False, num_classes=15)),
+            ("mobilenetv3_large_100", lambda: create_model('mobilenetv3_large_100', pretrained=False, num_classes=15)),
+        ]
+        
+        st.info("Trying other model architectures...")
+        for model_name, model_func in other_models_to_try:
+            try:
+                model = model_func()
+                model.load_state_dict(model_state)
+                model.eval()
+                model.to(device)
+                st.success(f"✅ Successfully loaded model using: {model_name}")
+                return model, device
+            except Exception as e:
+                st.warning(f"❌ Failed to load with {model_name}: {str(e)[:100]}...")
+                continue
+        
+        # If all fail, try loading with strict=False and show detailed info
+        st.info("Attempting to load with strict=False...")
+        try:
+            # Try with the most common ViT model first
+            model = create_model('vit_base_patch16_224', pretrained=False, num_classes=15)
+            missing_keys, unexpected_keys = model.load_state_dict(model_state, strict=False)
+            
+            if missing_keys:
+                st.warning(f"⚠️ Missing keys ({len(missing_keys)}): {missing_keys[:3]}...")
+            if unexpected_keys:
+                st.warning(f"⚠️ Unexpected keys ({len(unexpected_keys)}): {unexpected_keys[:3]}...")
+            
+            model.eval()
+            model.to(device)
+            
+            if len(missing_keys) > 0 or len(unexpected_keys) > 0:
+                st.error("⚠️ Model loaded with mismatched weights - predictions will likely be unreliable!")
+                st.info("💡 The saved model might have been trained with a different architecture.")
+            else:
+                st.success("✅ Model loaded successfully with strict=False")
+            
+            return model, device
+            
+        except Exception as e:
+            st.error(f"❌ Failed to load model with all methods. Error: {str(e)}")
+            st.info("💡 Try checking the model file or re-training with a compatible architecture.")
+            return None, device
+            
+    except Exception as e:
+        st.error(f"❌ Error downloading or loading model: {str(e)}")
+        return None, device
+
+def get_transform():
+    """Define image preprocessing transforms"""
+    return 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 predict_image(image, model, device, transform, label_names, topk=5):
+    """Make predictions on uploaded image"""
+    # Transform image
+    img_tensor = transform(image).unsqueeze(0).to(device)
+    
+    # Predict
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probs = F.softmax(outputs, dim=1)
+        top_probs, top_idxs = torch.topk(probs, k=topk)
+    
+    # Convert to CPU for display
+    top_probs = top_probs[0].cpu().numpy()
+    top_idxs = top_idxs[0].cpu().numpy()
+    
+    return top_probs, top_idxs
+
+def main():
+    st.title("🏀 Sports Ball Classifier")
+    st.markdown("Upload an image of a sports ball and get AI-powered predictions!")
+    
+    # Define label names
+    label_names = [
+        'american_football', 'baseball', 'basketball', 'billiard_ball', 
+        'bowling_ball', 'cricket_ball', 'football', 'golf_ball', 
+        'hockey_ball', 'hockey_puck', 'rugby_ball', 'shuttlecock', 
+        'table_tennis_ball', 'tennis_ball', 'volleyball'
+    ]
+    
+    # Load model
+    with st.spinner("Loading model..."):
+        model, device = load_model()
+    
+    if model is None:
+        st.error("Failed to load model. Please try again later.")
+        return
+    
+    st.success(f"Model loaded successfully! Using device: {device}")
+    
+    # Get image transform
+    transform = get_transform()
+    
+    # Create two columns
+    col1, col2 = st.columns([1, 1])
+    
+    with col1:
+        st.subheader("Upload Image")
+        uploaded_file = st.file_uploader(
+            "Choose an image...", 
+            type=['png', 'jpg', 'jpeg'],
+            help="Upload a clear image of a sports ball for best results"
+        )
+        
+        # Number of top predictions to show
+        topk = st.slider("Number of predictions to show:", 1, 10, 5)
+    
+    with col2:
+        st.subheader("Predictions")
+        
+        if uploaded_file is not None:
+            # Display uploaded image
+            image = Image.open(uploaded_file).convert("RGB")
+            st.image(image, caption="Uploaded Image", use_container_width=True)
+            
+            # Make predictions
+            with st.spinner("Analyzing image..."):
+                try:
+                    top_probs, top_idxs = predict_image(
+                        image, model, device, transform, label_names, topk
+                    )
+                    
+                    # Show original top prediction prominently
+                    top_confidence = float(top_probs[0] * 100)
+                    top_label = label_names[top_idxs[0]].replace('_', ' ').title()
+                    
+                    if top_confidence > 70:
+                        color = "🟢"
+                    elif top_confidence > 40:
+                        color = "🟡"
+                    else:
+                        color = "🔴"
+                    
+                    st.success(f"{color} **Primary Prediction: {top_label}** ({top_confidence:.2f}%)")
+                    st.progress(float(top_confidence / 100))
+                    
+                    # Show top 3 high confidence predictions
+                    st.subheader("Top 3 Predictions:")
+                    
+                    for i in range(min(3, len(top_probs))):
+                        confidence = float(top_probs[i] * 100)
+                        label = label_names[top_idxs[i]].replace('_', ' ').title()
+                        
+                        # Color coding based on confidence
+                        if confidence > 70:
+                            color = "🟢"
+                        elif confidence > 40:
+                            color = "🟡"
+                        else:
+                            color = "🔴"
+                        
+                        st.write(f"{i+1}. {color} **{label}**: {confidence:.2f}%")
+                        
+                        # Progress bar for confidence (convert to Python float)
+                        st.progress(float(confidence / 100))
+                    
+                    # Show all predictions if user wants more
+                    if topk > 3:
+                        with st.expander(f"See all {topk} predictions"):
+                            for i in range(3, len(top_probs)):
+                                confidence = float(top_probs[i] * 100)
+                                label = label_names[top_idxs[i]].replace('_', ' ').title()
+                                
+                                if confidence > 70:
+                                    color = "🟢"
+                                elif confidence > 40:
+                                    color = "🟡"
+                                else:
+                                    color = "🔴"
+                                
+                                st.write(f"{i+1}. {color} **{label}**: {confidence:.2f}%")
+                                st.progress(float(confidence / 100))
+                    
+                    # Show detailed results in expandable section
+                    with st.expander("Detailed Results"):
+                        fig, ax = plt.subplots(figsize=(10, 6))
+                        
+                        labels = [label_names[idx].replace('_', ' ').title() for idx in top_idxs]
+                        probabilities = [float(prob * 100) for prob in top_probs]  # Convert to Python float
+                        
+                        bars = ax.barh(labels[::-1], probabilities[::-1])
+                        ax.set_xlabel('Confidence (%)')
+                        ax.set_title(f'Top {topk} Predictions')
+                        ax.set_xlim(0, 100)
+                        
+                        # Color bars based on confidence
+                        for bar, prob in zip(bars, probabilities[::-1]):
+                            if prob > 70:
+                                bar.set_color('#4CAF50')  # Green
+                            elif prob > 40:
+                                bar.set_color('#FF9800')  # Orange
+                            else:
+                                bar.set_color('#F44336')  # Red
+                        
+                        # Add percentage labels on bars
+                        for i, (bar, prob) in enumerate(zip(bars, probabilities[::-1])):
+                            ax.text(float(prob) + 1, bar.get_y() + bar.get_height()/2, 
+                                   f'{float(prob):.1f}%', va='center')
+                        
+                        plt.tight_layout()
+                        st.pyplot(fig)
+                
+                except Exception as e:
+                    st.error(f"Error during prediction: {str(e)}")
+        
+        else:
+            st.info("👆 Please upload an image to get started!")
+    
+    # Additional information
+    st.markdown("---")
+    st.subheader("Supported Sports Balls")
+    
+    # Display supported categories in a nice grid
+    cols = st.columns(5)
+    for i, label in enumerate(label_names):
+        with cols[i % 5]:
+            st.write(f"• {label.replace('_', ' ').title()}")
+    
+    st.markdown("---")
+    st.markdown("""
+    **About this model:**
+    - Built using ConvNeXt architecture
+    - Trained to classify 15 different types of sports balls
+    - Model weights from: [Alamgirapi/sports-ball-convnext-classifier](https://huggingface.co/Alamgirapi/sports-ball-convnext-classifier)
+    
+    **Tips for best results:**
+    - Use clear, well-lit images
+    - Ensure the ball is the main subject
+    - Avoid cluttered backgrounds when possible
+    """)
+
+if __name__ == "__main__":
+    main()