feature: Use ViT as classifier

app.py

@@ -5,14 +5,14 @@ import cv2
 import PIL
 import torch
 
-from classifier import CustomEfficientNet
+from classifier import CustomEfficientNet, CustomViT
 from model import get_model, predict, prepare_prediction, predict_class
 
 print('Creating the model')
-model = get_model('checkpoint.ckpt')
+model = get_model('efficientDet_icevision.ckpt')
 print('Loading the classifier')
-classifier = CustomEfficientNet(target_size=7, pretrained=False)
-classifier.load_state_dict(torch.load('class_efficientB0_taco_7_class.pth', map_location='cpu'))
+classifier = CustomViT(target_size=7, pretrained=False)
+classifier.load_state_dict(torch.load('class_ViT_taco_7_class.pth', map_location='cpu'))
 
 def plot_img_no_mask(image, boxes, labels):
     colors = {

classifier.py

@@ -1,12 +1,27 @@
 import timm
 import torch.nn as nn
-
+import albumentations as A
 import torch
+import cv2
+
+class CustomNormalization(A.ImageOnlyTransform):
+    def _norm(self, img):
+        return img / 255.
+
+    def apply(self, img, **params):
+        return self._norm(img)
+
+def transform_image(image, size):
+    transforms = [
+        A.Resize(size, size,
+                interpolation=cv2.INTER_NEAREST),
+        CustomNormalization(p=1),
+    ]
 
-def get_efficientnet(model_name):
-    model = timm.create_model(model_name, pretrained=True)
+    augs = A.Compose(transforms)
+    transformed = augs(image=image)
 
-    return model
+    return transformed['image']
 
 class CustomEfficientNet(nn.Module):
     """
@@ -43,3 +58,40 @@ class CustomEfficientNet(nn.Module):
         x = self.model(x)
 
         return x
+
+class CustomViT(nn.Module):
+    """
+    This class defines a custom ViT network.
+
+    Parameters
+    ----------
+    target_size : int
+        Number of units for the output layer.
+    pretrained : bool
+        Determine if pretrained weights are used.
+
+    Attributes
+    ----------
+    model : nn.Module
+        CustomViT model.
+    """
+    def __init__(self, model_name : str = 'vit_base_patch16_224',
+                 target_size : int = 4, pretrained : bool = True):
+        super().__init__()
+        self.model = timm.create_model(model_name,
+                                       pretrained=pretrained,
+                                       num_classes=target_size)
+        
+        in_features = self.model.head.in_features
+        self.model.head = nn.Sequential(
+            #nn.Dropout(0.5),
+            nn.Linear(in_features, 256),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(256, target_size)
+        )
+            
+    def forward(self, x : torch.Tensor) -> torch.Tensor:
+        x = self.model(x)
+
+        return x

model.py

@@ -6,6 +6,8 @@ import torch
 import numpy as np
 import torchvision
 
+from classifier import transform_image
+
 import icevision.models.ross.efficientdet
 
 MODEL_TYPE = icevision.models.ross.efficientdet
@@ -81,10 +83,14 @@ def predict_class(model, image, bboxes):
         img = image.copy()
         bbox = np.array(bbox).astype(int)
         cropped_img = PIL.Image.fromarray(img).crop(bbox)
-        cropped_img = np.array(cropped_img).transpose(2, 0, 1)
-        cropped_img = torch.as_tensor(cropped_img, dtype=torch.float).unsqueeze(0)
-
-        y_preds = model(cropped_img)
+        cropped_img = np.array(cropped_img)
+        #cropped_img = torch.as_tensor(cropped_img, dtype=torch.float).unsqueeze(0)
+
+        tran_image = transform_image(cropped_img, 224)
+        tran_image = tran_image.transpose(2, 0, 1)
+        tran_image = torch.as_tensor(tran_image, dtype=torch.float).unsqueeze(0)
+        print(tran_image.shape)
+        y_preds = model(tran_image)
         preds.append(y_preds.softmax(1).detach().numpy())
 
     preds = np.concatenate(preds).argmax(1)