Update app.py

app.py

@@ -4,15 +4,149 @@ from huggingface_hub import hf_hub_download
 import torch
 import torch.nn as nn
 from torchvision import transforms
-from PIL import Image
 
-# Define the model classes (SimpleResidualBlock, BottleneckResidualBlock, ResNet) here...
+
+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)
+
 
 class ImageClassifier:
     def __init__(self, checkpoint_path):
         self.checkpoint_path = checkpoint_path
         self.model = self.load_model(checkpoint_path)
-        self.transform = self.get_transform((224, 224))  # Typical size for ResNet
+        self.transform = self.get_transform((244, 244))
         self.labels = [
             "airplane",
             "automobile",
@@ -32,7 +166,7 @@ class ImageClassifier:
             block="simple",
             num_classes=10,
         )
-        classifier.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
+        classifier.load_state_dict(torch.load(checkpoint_path))
         classifier = classifier.cpu()
         classifier.eval()
         return classifier
@@ -56,18 +190,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(
-    fn=classify,
+    classify,
     inputs=[
         gr.Image(label="Input Image", type="pil"),
     ],