Update app.py

app.py

@@ -4,149 +4,15 @@ from huggingface_hub import hf_hub_download
 import torch
 import torch.nn as nn
 from torchvision import transforms
+from PIL import Image
 
-
-class SimpleResidualBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, set_stride=False):
-        super().__init__()
-        stride = 2 if in_channels != out_channels and set_stride else 1
-
-        self.conv1 = nn.LazyConv2d(
-            out_channels,
-            kernel_size=3,
-            padding="same" if stride == 1 else 1,
-            stride=stride,
-        )
-        self.conv2 = nn.LazyConv2d(out_channels, kernel_size=3, padding="same")
-
-        self.bn1 = nn.LazyBatchNorm2d()
-        self.bn2 = nn.LazyBatchNorm2d()
-
-        self.relu = nn.ReLU()
-
-        if in_channels != out_channels:
-            self.residual = nn.Sequential(
-                nn.LazyConv2d(out_channels, kernel_size=1, stride=stride),
-                nn.LazyBatchNorm2d(),
-            )
-        else:
-            self.residual = nn.Identity()
-
-    def forward(self, x):
-        out = self.relu(self.bn1(self.conv1(x)))
-        out = self.bn2(self.conv2(out))
-        out += self.residual(x)
-        out = self.relu(out)
-        return out
-
-
-class BottleneckResidualBlock(nn.Module):
-    def __init__(
-        self, in_channels, out_channels, identity_mapping=False, set_stride=False
-    ):
-        super().__init__()
-        stride = 2 if in_channels != out_channels and set_stride else 1
-
-        self.conv1 = nn.LazyConv2d(
-            out_channels,
-            kernel_size=1,
-            padding="same" if stride == 1 else 0,
-            stride=stride,
-        )
-        self.conv2 = nn.LazyConv2d(out_channels, kernel_size=3, padding="same")
-        self.conv3 = nn.LazyConv2d(out_channels * 4, kernel_size=1, padding="same")
-
-        self.bn1 = nn.LazyBatchNorm2d()
-        self.bn2 = nn.LazyBatchNorm2d()
-        self.bn3 = nn.LazyBatchNorm2d()
-
-        self.relu = nn.ReLU()
-
-        if in_channels != out_channels or not identity_mapping:
-            self.residual = nn.Sequential(
-                nn.LazyConv2d(out_channels * 4, kernel_size=1, stride=stride),
-                nn.LazyBatchNorm2d(),
-            )
-        else:
-            self.residual = nn.Identity()
-
-    def forward(self, x):
-        out = self.relu(self.bn1(self.conv1(x)))
-        out = self.relu(self.bn2(self.conv2(out)))
-        out = self.bn3(self.conv3(out))
-        out += self.residual(x)
-        out = self.relu(out)
-        return out
-
-
-RESNET_18 = [2, 2, 2, 2]
-RESNET_34 = [3, 4, 6, 3]
-RESNET_50 = [3, 4, 6, 3]
-RESNET_101 = [3, 4, 23, 3]
-RESNET_152 = [3, 8, 36, 3]
-
-
-class ResNet(nn.Module):
-    def __init__(self, arch=RESNET_18, block="simple", num_classes=256):
-        super().__init__()
-        self.conv1 = nn.Sequential(
-            nn.LazyConv2d(64, kernel_size=7, stride=2, padding=3),
-            nn.LazyBatchNorm2d(),
-            nn.ReLU(),
-        )
-        self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
-        self.conv2 = self._make_layer(64, 64, arch[0], set_stride=False, block=block)
-        self.conv3 = self._make_layer(64, 128, arch[1], block=block)
-        self.conv4 = self._make_layer(128, 256, arch[2], block=block)
-        self.conv5 = self._make_layer(256, 512, arch[3], block=block)
-        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
-        self.flatten = nn.Flatten()
-        self.fc = nn.LazyLinear(num_classes)
-
-    def _make_layer(
-        self, in_channels, out_channels, num_blocks, set_stride=True, block="simple"
-    ):
-        """Block is either 'simple' or 'bottleneck'"""
-        layers = []
-        for i in range(num_blocks):
-            layers.append(
-                SimpleResidualBlock(in_channels, out_channels, set_stride=set_stride)
-                if block == "simple"
-                else BottleneckResidualBlock(
-                    in_channels if i == 0 else out_channels * 4,
-                    out_channels,
-                    set_stride=set_stride,
-                )
-            )
-            set_stride = False
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        out = self.conv1(x)
-        out = self.maxpool(self.conv2(out))
-        out = self.conv3(out)
-        out = self.conv4(out)
-        out = self.conv5(out)
-        out = self.avgpool(out)
-        out = self.flatten(out)
-        out = self.fc(out)
-        return out
-
-    def _init_weights(module):
-        # Initlize weights with glorot uniform
-        if isinstance(module, nn.Conv2d):
-            nn.init.xavier_uniform_(module.weight)
-            nn.init.zeros_(module.bias)
-        elif isinstance(module, nn.Linear):
-            nn.init.xavier_uniform_(module.weight)
-            nn.init.zeros_(module.bias)
-
+# Define the model classes (SimpleResidualBlock, BottleneckResidualBlock, ResNet) here...
 
 class ImageClassifier:
     def __init__(self, checkpoint_path):
         self.checkpoint_path = checkpoint_path
         self.model = self.load_model(checkpoint_path)
-        self.transform = self.get_transform((244, 244))
+        self.transform = self.get_transform((224, 224))  # Typical size for ResNet
         self.labels = [
             "airplane",
             "automobile",
@@ -166,7 +32,7 @@ class ImageClassifier:
             block="simple",
             num_classes=10,
         )
-        classifier.load_state_dict(torch.load(checkpoint_path))
+        classifier.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
         classifier = classifier.cpu()
         classifier.eval()
         return classifier
@@ -190,18 +56,18 @@ class ImageClassifier:
     def classify(self, input_image):
         return self.predict(input_image)
 
-
 def classify(input_image):
     return classifier.classify(input_image)
 
-
 checkpoint_path = hf_hub_download(
     repo_id="SatwikKambham/resnet18-cifar10",
     filename="model.pt",
 )
+
 classifier = ImageClassifier(checkpoint_path)
+
 iface = gr.Interface(
-    classify,
+    fn=classify,
     inputs=[
         gr.Image(label="Input Image", type="pil"),
     ],