Minimal bare-bones version

app.py

@@ -1,8 +1,3 @@
-"""
-CIFAR-100 Image Classification - Hugging Face Space (Fixed)
-Uses model from: https://huggingface.co/ivantv/cifar100-resnet18
-"""
-
 import torch
 import torch.nn as nn
 from PIL import Image
@@ -10,7 +5,6 @@ import torchvision.transforms as transforms
 import gradio as gr
 from huggingface_hub import hf_hub_download
 
-# CIFAR-100 class names
 CIFAR100_CLASSES = [
     'apple', 'aquarium_fish', 'baby', 'bear', 'beaver', 'bed', 'bee', 'beetle',
     'bicycle', 'bottle', 'bowl', 'boy', 'bridge', 'bus', 'butterfly', 'camel',
@@ -28,31 +22,20 @@ CIFAR100_CLASSES = [
     'turtle', 'wardrobe', 'whale', 'willow_tree', 'wolf', 'woman', 'worm'
 ]
 
-print("=" * 60)
-print("Loading CIFAR-100 Classification Model")
-print("=" * 60)
-
-# Define model architecture
 class BasicBlock(nn.Module):
     expansion = 1
-    
     def __init__(self, in_channels, out_channels, stride=1):
         super(BasicBlock, self).__init__()
-        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, 
-                               stride=stride, padding=1, bias=False)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
         self.bn1 = nn.BatchNorm2d(out_channels)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3,
-                               stride=1, padding=1, bias=False)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
         self.bn2 = nn.BatchNorm2d(out_channels)
-        
         self.shortcut = nn.Sequential()
         if stride != 1 or in_channels != self.expansion * out_channels:
             self.shortcut = nn.Sequential(
-                nn.Conv2d(in_channels, self.expansion * out_channels,
-                         kernel_size=1, stride=stride, bias=False),
+                nn.Conv2d(in_channels, self.expansion * out_channels, kernel_size=1, stride=stride, bias=False),
                 nn.BatchNorm2d(self.expansion * out_channels)
             )
-    
     def forward(self, x):
         out = torch.relu(self.bn1(self.conv1(x)))
         out = self.bn2(self.conv2(out))
@@ -60,12 +43,10 @@ class BasicBlock(nn.Module):
         out = torch.relu(out)
         return out
 
-
 class ResNet(nn.Module):
     def __init__(self, block, num_blocks, num_classes=100):
         super(ResNet, self).__init__()
         self.in_channels = 64
-        
         self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
         self.bn1 = nn.BatchNorm2d(64)
         self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
@@ -73,7 +54,6 @@ class ResNet(nn.Module):
         self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
         self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
         self.linear = nn.Linear(512 * block.expansion, num_classes)
-    
     def _make_layer(self, block, out_channels, num_blocks, stride):
         strides = [stride] + [1] * (num_blocks - 1)
         layers = []
@@ -81,7 +61,6 @@ class ResNet(nn.Module):
             layers.append(block(self.in_channels, out_channels, stride))
             self.in_channels = out_channels * block.expansion
         return nn.Sequential(*layers)
-    
     def forward(self, x):
         out = torch.relu(self.bn1(self.conv1(x)))
         out = self.layer1(out)
@@ -93,108 +72,46 @@ class ResNet(nn.Module):
         out = self.linear(out)
         return out
 
-
 def ResNet18():
     return ResNet(BasicBlock, [2, 2, 2, 2])
 
-
-# Load model
-print("📥 Downloading model from Hugging Face Hub...")
-model_path = hf_hub_download(
-    repo_id="ivantv/cifar100-resnet18",
-    filename="best_model.pth"
-)
-
+print("Loading model...")
+model_path = hf_hub_download(repo_id="ivantv/cifar100-resnet18", filename="best_model.pth")
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-print(f"💻 Using device: {device}")
-
 model = ResNet18().to(device)
 model.load_state_dict(torch.load(model_path, map_location=device))
 model.eval()
-print("✅ Model loaded successfully!")
-print("=" * 60)
+print("Model loaded!")
 
-# Define image transforms
 transform = transforms.Compose([
     transforms.Resize((32, 32)),
     transforms.ToTensor(),
     transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761))
 ])
 
-
-def predict_image(image):
-    """Make prediction on uploaded image"""
+def classify_image(image):
     if image is None:
-        return None
-    
-    try:
-        # Convert to PIL Image if needed
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image)
-        
-        # Convert to RGB if needed
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        
-        # Preprocess image
-        image_tensor = transform(image).unsqueeze(0).to(device)
-        
-        # Make prediction
-        with torch.no_grad():
-            output = model(image_tensor)
-            probabilities = torch.nn.functional.softmax(output, dim=1)
-            
-            # Get top-5 predictions
-            top_probs, top_indices = torch.topk(probabilities[0], 5)
-        
-        # Prepare results
-        top_probs = top_probs.cpu().numpy()
-        top_indices = top_indices.cpu().numpy()
-        
-        # Create predictions dictionary for Gradio
-        predictions = {}
-        for prob, idx in zip(top_probs, top_indices):
-            class_name = CIFAR100_CLASSES[idx]
-            predictions[class_name] = float(prob)
-        
-        return predictions
-    
-    except Exception as e:
-        print(f"Error during prediction: {e}")
-        return {"error": f"Prediction failed: {str(e)}"}
-
-
-# Create simple Gradio interface (avoiding the gr.Blocks bug)
-demo = gr.Interface(
-    fn=predict_image,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Label(num_top_classes=5, label="Top-5 Predictions"),
-    title="🖼️ CIFAR-100 Image Classification",
-    description="""
-    Upload an image to classify it into one of **100 CIFAR-100 categories**!
-    
-    **Model:** [ivantv/cifar100-resnet18](https://huggingface.co/ivantv/cifar100-resnet18)  
-    **Architecture:** ResNet-18 (11.2M parameters)  
-    **Test Accuracy:** 75.84%  
-    
-    The model can recognize 100 categories including animals, vehicles, household objects, plants, and more.
-    """,
-    article="""
-    ### About the Model
-    
-    This ResNet-18 model was trained on the CIFAR-100 dataset with 50 epochs of training, achieving 75.84% test accuracy.
+        return {}
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
     
-    **Categories include:**
-    - 🐾 Animals: bear, tiger, elephant, dolphin, fox, lion, wolf, etc.
-    - 🚗 Vehicles: bicycle, bus, train, motorcycle, pickup_truck, etc.
-    - 🏠 Objects: chair, table, lamp, clock, keyboard, etc.
-    - 🌳 Plants: maple_tree, oak_tree, rose, tulip, orchid, etc.
+    img_tensor = transform(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        output = model(img_tensor)
+        probs = torch.nn.functional.softmax(output, dim=1)[0]
+        top5_prob, top5_idx = torch.topk(probs, 5)
     
-    **Tips:** Use clear, well-lit images with simple backgrounds for best results.
-    """,
-    theme="soft"
+    return {CIFAR100_CLASSES[idx]: prob.item() for prob, idx in zip(top5_prob, top5_idx)}
+
+iface = gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=5),
+    title="CIFAR-100 Classifier",
+    description="ResNet-18 model trained on CIFAR-100 (75.84% accuracy)"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+iface.launch()