Trained SqueezeNet with torch weights

services/backend/model/model.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:aa39db0e7695da5c388e69de4f56b3b5e5a6f33df3a9bdb197f3e59d8343a16f
-size 45041378
+oid sha256:5541ad1fb46690d0d0de99beafa579f013e8cfcda6ffa1e7b616acc246e69eb3
+size 3215922

services/backend/model/transform.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c26dd204ca8676a5affd07a2cbbd162b79e0e01ebba451fdf599aa6abb242a46
+oid sha256:471b2bb67abb7f23f5e283cd84643e3148776dfc6a7fb0597775b397bdb8a710
 size 4363

services/backend/src/architectures/squeeze_net.py

@@ -6,6 +6,7 @@ import torch
 from torch import nn
 from collections import OrderedDict
 from src.architectures.deep_cnn import CNNBlock
+from src.architectures.base import Backbone
 
 
 class FireBlock(nn.Module):
@@ -38,7 +39,7 @@ class FireBlock(nn.Module):
         return out
 
 
-class SqueezeNet(nn.Module):
+class SqueezeNet1_0(nn.Module):
     def __init__(
         self,
         in_channels: int = 3,
@@ -101,3 +102,28 @@ class SqueezeNet(nn.Module):
     @property
     def name(self):
         return "SqueezeNet"
+
+
+from typing import Literal
+
+
+class SqueezeNet:
+    def __new__(
+        cls,
+        in_channels: int,
+        version: Literal["squeezenet1_0", "squeezenet1_1"],
+        load_from_torch: bool = False,
+        pretrained: bool = False,
+        freeze_extractor: bool = False,
+    ):
+        if load_from_torch:
+            _net = torch.hub.load("pytorch/vision:v0.10.0", version, pretrained=pretrained)
+            _net.classifier[0] = torch.nn.Identity()
+            _net.classifier[1] = torch.nn.Identity()
+            _net.classifier[2] = torch.nn.Identity()
+            net = Backbone(_net, out_channels=512, name=version)
+            if freeze_extractor:
+                net.freeze()
+        else:
+            net = SqueezeNet1_0(in_channels=in_channels)
+        return net

services/backend/src/data.py

@@ -41,7 +41,7 @@ class FlowersDataset(torchvision.datasets.Flowers102):
             self.idx2label = {i: label for i, label in enumerate(labels)}
 
         self._named_labels = np.array([self.idx2label[target] for target in self._labels])
-        self.classes = list(label2idx.keys())
+        self.classes = list(self.idx2label.values())
 
     def get_raw_img(self, idx: int):
         image_file = self._image_files[idx]

services/backend/src/main.py

@@ -21,7 +21,7 @@ app.add_middleware(
 
 model_path = "model/model.pt"
 transform_path = "model/transform.pt"
-mapping_path = "model/mapping.json"
+mapping_path = "model/mapping.txt"
 
 model = torch.jit.load(model_path)
 model.eval()
@@ -29,9 +29,10 @@ model.to("cpu")
 transform = torch.jit.load(transform_path)
 
 
-with open(mapping_path, "rb") as f:
-    label2idx = json.load(f)
-LABELS = list(label2idx.keys())
+with open(mapping_path) as f:
+    labels = f.readlines()
+    labels = [label.strip() for label in labels]
+    idx2label = {i: label for i, label in enumerate(labels)}
 
 
 def load_image_into_numpy_array(data):
@@ -45,8 +46,8 @@ async def predict(file: UploadFile = File(...)):
     img = transform(img).unsqueeze(0)
     log_probs = model(img)[0]
     probs = torch.exp(log_probs)
-    label_probs = {LABELS[i]: float(probs[i]) for i in range(len(LABELS))}
-    return label_probs
+    confidences = {idx2label[i]: float(probs[i]) for i in range(len(idx2label))}
+    return confidences
 
 
 @app.get("/")

services/backend/src/model.py

@@ -136,14 +136,20 @@ def load_net(num_classes: int) -> nn.Module:
     #     pool_kernels=[2, 1, 2, 1, 2],
     # )
 
-    # backbone = SqueezeNet()
-    backbone = ResNet(
+    backbone = SqueezeNet(
         in_channels=3,
-        version="resnet18",
+        version="squeezenet1_0",
         load_from_torch=True,
         pretrained=True,
         freeze_extractor=True,
     )
+    # backbone = ResNet(
+    #     in_channels=3,
+    #     version="resnet18",
+    #     load_from_torch=True,
+    #     pretrained=True,
+    #     freeze_extractor=True,
+    # )
 
     return nn.Sequential(
         backbone,