Multiple refactors

app.py

@@ -3,8 +3,7 @@ import PIL
 import torch
 
 from utils import plot_img_no_mask, get_models
-from classifier import CustomEfficientNet, CustomViT
-from model import get_model, predict, prepare_prediction, predict_class
+from model import predict, prepare_prediction, predict_class
 
 DET_CKPT = 'efficientDet_icevision.ckpt'
 CLASS_CKPT = 'class_ViT_taco_7_class.pth'

model.py

@@ -1,11 +1,10 @@
 from io import BytesIO
-from typing import Union
+from typing import Dict, Tuple, Union
 from icevision import *
 from icevision.models.checkpoint import model_from_checkpoint
 from classifier import transform_image
 from icevision.models import ross
 
-import collections
 import PIL
 import torch
 import numpy as np
@@ -13,44 +12,34 @@ import torchvision
 
 MODEL_TYPE = ross.efficientdet
 
-def get_model(checkpoint_path : str):
-    checkpoint_and_model = model_from_checkpoint(
-                                checkpoint_path,
-                                model_name='ross.efficientdet',
-                                backbone_name='d0',
-                                img_size=512,
-                                classes=['Waste'],
-                                revise_keys=[(r'^model\.', '')])
-
-    model = checkpoint_and_model['model']
-
-    return model
-
-def get_checkpoint(checkpoint_path : str):
-    ckpt = torch.load(checkpoint_path, map_location=torch.device('cpu'))
-
-    fixed_state_dict = collections.OrderedDict()
-
-    for k, v in ckpt['state_dict'].items():
-        new_k = k[6:]
-        fixed_state_dict[new_k] = v
-
-    return fixed_state_dict
-
-def predict(model : object, image : Union[str, BytesIO], detection_threshold : float):
+def predict(det_model : torch.nn.Module, image : Union[str, BytesIO],
+            detection_threshold : float) -> Dict:
+    """
+    Make a prediction with the detection model.
+
+    Args:
+        det_model (torch.nn.Module): Detection model
+        image (Union[str, BytesIO]): Image filepath if the image is one of
+            the example images and BytesIO if the image is a custom image
+            uploaded by the user.
+        detection_threshold (float): Detection threshold
+
+    Returns:
+        Dict: Prediction dictionary.
+    """        
     img = PIL.Image.open(image)
-    #img = PIL.Image.open(BytesIO(image))
-    img = np.array(img)
-    img = PIL.Image.fromarray(img)
+
+    # Class map and transforms
     class_map = ClassMap(classes=['Waste'])
     transforms = tfms.A.Adapter([
                     *tfms.A.resize_and_pad(512),
                     tfms.A.Normalize()
                 ])
-
+    
+    # Single prediction
     pred_dict  = MODEL_TYPE.end2end_detect(img,
                                            transforms, 
-                                           model,
+                                           det_model,
                                            class_map=class_map,
                                            detection_threshold=detection_threshold,
                                            return_as_pil_img=False,
@@ -61,32 +50,67 @@ def predict(model : object, image : Union[str, BytesIO], detection_threshold : f
 
     return pred_dict
 
-def prepare_prediction(pred_dict, threshold):
+def prepare_prediction(pred_dict : Dict,
+                       nms_threshold : str) -> Tuple[torch.Tensor, np.ndarray]:
+    """
+    Get the predictions in a right format.
+
+    Args:
+        pred_dict (Dict): Prediction dictionary.
+        nms_threshold (float): Threshold for the NMS postprocess.
+
+    Returns:
+        Tuple: Tuple containing the following:
+            - (torch.Tensor): Bounding boxes
+            - (np.ndarray): Image data
+    """
+    # Convert each box to a tensor and stack them into an unique tensor
     boxes = [box.to_tensor() for box in pred_dict['detection']['bboxes']]
     boxes = torch.stack(boxes)
 
+    # Get the scores and labels as tensor
     scores = torch.as_tensor(pred_dict['detection']['scores'])
     labels = torch.as_tensor(pred_dict['detection']['label_ids'])
+
     image = np.array(pred_dict['img'])
 
-    fixed_boxes = torchvision.ops.batched_nms(boxes, scores, labels, threshold)
+    # Apply NMS to postprocess the bounding boxes
+    fixed_boxes = torchvision.ops.batched_nms(boxes, scores,
+                                              labels,nms_threshold)
     boxes = boxes[fixed_boxes, :]
 
     return boxes, image
 
-def predict_class(classifier, image, bboxes):
+def predict_class(classifier : torch.nn.Module, image : np.ndarray,
+                  bboxes : torch.Tensor) -> np.ndarray:
+    """
+    Predict the class of each detected object.
+
+    Args:
+        classifier (torch.nn.Module): Classifier model.
+        image (np.ndarray): Image data.
+        bboxes (torch.Tensor): Bounding boxes.
+
+    Returns:
+        np.ndarray: Array containing the predicted class for each object.
+    """
     preds = []
 
     for bbox in bboxes:
         img = image.copy()
         bbox = np.array(bbox).astype(int)
+
+        # Get the bounding box content
         cropped_img = PIL.Image.fromarray(img).crop(bbox)
         cropped_img = np.array(cropped_img)
 
+        # Apply transformations to the cropped image
         tran_image = transform_image(cropped_img, 224)
+        # Channels first
         tran_image = tran_image.transpose(2, 0, 1)
         tran_image = torch.as_tensor(tran_image, dtype=torch.float).unsqueeze(0)
-        print(tran_image.shape)
+
+        # Make prediction
         y_preds = classifier(tran_image)
         preds.append(y_preds.softmax(1).detach().numpy())
 

requirements.txt

@@ -1,4 +1,5 @@
 icevision[all]
 matplotlib
 effdet
+mmcv-full
 Pillow==8.4.0

utils.py

@@ -4,8 +4,8 @@ import numpy as np
 import cv2
 import torch
 
+from icevision.models.checkpoint import model_from_checkpoint
 from classifier import CustomViT
-from model import get_model
 
 def plot_img_no_mask(image : np.ndarray, boxes : torch.Tensor, labels):
     colors = {
@@ -67,7 +67,15 @@ def get_models(
             - (torch.nn.Module): Classifier model
     """
     print('Loading the detection model')
-    det_model = get_model(detection_ckpt)
+    checkpoint_and_model = model_from_checkpoint(
+                                detection_ckpt,
+                                model_name='ross.efficientdet',
+                                backbone_name='d0',
+                                img_size=512,
+                                classes=['Waste'],
+                                revise_keys=[(r'^model\.', '')])
+
+    det_model = checkpoint_and_model['model']
     det_model.eval()
 
     print('Loading the classifier model')