Create app.py

app.py

@@ -0,0 +1,178 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+from PIL import Image
+import numpy as np
+
+# Define the ResNet architecture (same as in training)
+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.bn1 = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        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 != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, 
+                         stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+
+    def forward(self, x):
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = self.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.relu = nn.ReLU(inplace=True)
+        
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        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.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+    def _make_layer(self, block, out_channels, num_blocks, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        layers = []
+        for stride in strides:
+            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 = self.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = self.avgpool(out)
+        out = torch.flatten(out, 1)
+        out = self.fc(out)
+        return out
+
+def ResNet34():
+    return ResNet(BasicBlock, [3, 4, 6, 3], num_classes=100)
+
+# CIFAR-100 class names
+CIFAR100_CLASSES = [
+    'apple', 'aquarium_fish', 'baby', 'bear', 'beaver', 'bed', 'bee', 'beetle', 
+    'bicycle', 'bottle', 'bowl', 'boy', 'bridge', 'bus', 'butterfly', 'camel', 
+    'can', 'castle', 'caterpillar', 'cattle', 'chair', 'chimpanzee', 'clock', 
+    'cloud', 'cockroach', 'couch', 'crab', 'crocodile', 'cup', 'dinosaur', 
+    'dolphin', 'elephant', 'flatfish', 'forest', 'fox', 'girl', 'hamster', 
+    'house', 'kangaroo', 'keyboard', 'lamp', 'lawn_mower', 'leopard', 'lion',
+    'lizard', 'lobster', 'man', 'maple_tree', 'motorcycle', 'mountain', 'mouse',
+    'mushroom', 'oak_tree', 'orange', 'orchid', 'otter', 'palm_tree', 'pear',
+    'pickup_truck', 'pine_tree', 'plain', 'plate', 'poppy', 'porcupine',
+    'possum', 'rabbit', 'raccoon', 'ray', 'road', 'rocket', 'rose',
+    'sea', 'seal', 'shark', 'shrew', 'skunk', 'skyscraper', 'snail', 'snake',
+    'spider', 'squirrel', 'streetcar', 'sunflower', 'sweet_pepper', 'table',
+    'tank', 'telephone', 'television', 'tiger', 'tractor', 'train', 'trout',
+    'tulip', 'turtle', 'wardrobe', 'whale', 'willow_tree', 'wolf', 'woman', 'worm'
+]
+
+# Load model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = ResNet34()
+
+# Load checkpoint
+checkpoint = torch.load('cifar100_resnet34_final.pth', map_location=device)
+model.load_state_dict(checkpoint['model_state_dict'])
+model.to(device)
+model.eval()
+
+# Get normalization parameters from checkpoint
+mean = checkpoint.get('normalization_mean', [0.5071, 0.4867, 0.4408])
+std = checkpoint.get('normalization_std', [0.2675, 0.2565, 0.2761])
+
+# Define transforms
+transform = transforms.Compose([
+    transforms.Resize((32, 32)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean, std),
+])
+
+def predict(image):
+    """
+    Predict the class of an image using the trained ResNet model.
+    """
+    if image is None:
+        return None
+    
+    # Convert to PIL Image if necessary
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image.astype('uint8'), 'RGB')
+    
+    # Apply transforms
+    img_tensor = transform(image).unsqueeze(0).to(device)
+    
+    # Make prediction
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probabilities = torch.nn.functional.softmax(outputs, dim=1)
+        
+        # Get top 5 predictions
+        top5_prob, top5_idx = torch.topk(probabilities[0], 5)
+        
+        # Create dictionary of top 5 predictions
+        predictions = {}
+        for i in range(5):
+            class_name = CIFAR100_CLASSES[top5_idx[i].item()]
+            probability = top5_prob[i].item()
+            predictions[class_name] = float(probability)
+    
+    return predictions
+
+# Create Gradio interface
+title = "CIFAR-100 Image Classifier (ResNet34)"
+description = """
+This is a ResNet34 model trained on CIFAR-100 dataset with 80%+ accuracy.
+Upload an image to classify it into one of 100 categories.
+
+The model works best with:
+- Natural images (animals, objects, vehicles, etc.)
+- Images with clear subjects
+- Square aspect ratio images
+
+Note: The model was trained on 32x32 images, so very high resolution details might not be fully utilized.
+"""
+
+examples = [
+    # You can add example image paths here if you have them
+]
+
+# Create the interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Label(num_top_classes=5, label="Top 5 Predictions"),
+    title=title,
+    description=description,
+    examples=examples if examples else None,
+    theme="default",
+    allow_flagging="never"
+)
+
+# Launch the app
+if __name__ == "__main__":
+    iface.launch()