switched to transfer learning

app.py

@@ -1,27 +1,29 @@
 import gradio as gr
 import torch
-from torchvision import transforms
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms, models
 import numpy as np
 import os
 import huggingface_hub
 
-from net import Net
-
 if torch.cuda.is_available():
     device = torch.device("cuda:0")
 else:
     device = torch.device("cpu")
 
 
-net = Net()
+net = models.resnet18(pretrained=False)
+net.fc = nn.Linear(net.fc.in_features, 6)
 net.to(device)
 
 HF_Token = os.environ['HF_Token']
 
 model = huggingface_hub.cached_download(huggingface_hub.hf_hub_url(
-    'danyalmalik/sceneryclassifier', '1655684183.7481008_Acc0.87_modelweights.pth'), use_auth_token=HF_Token)
+    'danyalmalik/sceneryclassifier', '1655988285.9725637_Acc0.88_modelweights.pth'), use_auth_token=HF_Token)
 
 net.load_state_dict(torch.load(model, map_location=device))
+net.eval()
 
 mean = np.array([0.5, 0.5, 0.5])
 std = np.array([0.25, 0.25, 0.25])
@@ -49,9 +51,10 @@ def predict(img):
     try:
         img = data_transforms(img)
         img = img.to(device)
+        img = img.unsqueeze(0)
 
         with torch.no_grad():
-            output = net(img)
+            output = F.softmax(net(img), dim=1)
 
             pred = [output[0][i].item() for i in range(len(labels))]
 
@@ -62,5 +65,5 @@ def predict(img):
     return weightage
 
 
-gr.Interface(fn=predict, inputs=gr.Image(shape=(150, 150), type='pil'),
+gr.Interface(fn=predict, inputs=gr.Image(type='pil'),
              outputs='label', title=title,  examples=examples()).launch()

net.py

@@ -1,35 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class Net(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.conv1 = nn.Conv2d(3, 32, 5)
-        self.conv2 = nn.Conv2d(32, 64, 5)
-        self.conv3 = nn.Conv2d(64, 128, 5)
-
-        x = torch.randn(3, 150, 150).view(-1, 3, 150, 150)
-        self._to_linear = None
-        self.convs(x)
-
-        self.fc1 = nn.Linear(self._to_linear, 512)
-        self.fc2 = nn.Linear(512, 6)
-
-    def convs(self, x):
-        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
-        x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
-        x = F.max_pool2d(F.relu(self.conv3(x)), (2, 2))
-
-        if self._to_linear is None:
-            self._to_linear = x[0].shape[0]*x[0].shape[1]*x[0].shape[2]
-        return x
-
-    def forward(self, x):
-        x = self.convs(x)
-        x = x.view(-1, self._to_linear)
-        x = F.relu(self.fc1(x))
-        x = self.fc2(x)
-
-        return F.softmax(x, dim=1)