feature: Objects classification

app.py

@@ -3,25 +3,55 @@ import matplotlib.pyplot as plt
 import numpy as np
 import cv2
 import PIL
+import torch
 
-from model import get_model, predict, prepare_prediction
+from classifier import CustomEfficientNet
+from model import get_model, predict, prepare_prediction, predict_class
 
 print('Creating the model')
 model = get_model('checkpoint.ckpt')
+print('Loading the classifier')
+classifier = CustomEfficientNet(target_size=7, pretrained=False)
+classifier.load_state_dict(torch.load('class_efficientB0_taco_7_class.pth'))
+
+def plot_img_no_mask(image, boxes, labels):
+    colors = {
+        0: (255,255,0),
+        1: (255, 0, 0),
+        2: (0, 0, 255),
+        3: (0,128,0),
+        4: (255,165,0),
+        5: (230,230,250),
+        6: (192,192,192)
+    }
+
+    texts = {
+        0: 'plastic',
+        1: 'dangerous',
+        2: 'carton',
+        3: 'glass',
+        4: 'organic',
+        5: 'rest',
+        6: 'other'
+    }
 
-def plot_img_no_mask(image, boxes):
     # Show image
     boxes = boxes.cpu().detach().numpy().astype(np.int32)
     fig, ax = plt.subplots(1, 1, figsize=(12, 6))
 
-    for i, box in enumerate(boxes):            
+    for i, box in enumerate(boxes):
+        color = colors[labels[i]]
+
         [x1, y1, x2, y2] = np.array(box).astype(int)
         # Si no se hace la copia da error en cv2.rectangle
         image = np.array(image).copy()
 
         pt1 = (x1, y1)
         pt2 = (x2, y2)
-        cv2.rectangle(image, pt1, pt2, (220,0,0), thickness=5)
+        cv2.rectangle(image, pt1, pt2, color, thickness=5)
+        cv2.putText(image, texts[labels[i]], (x1, y1-10),
+                    cv2.FONT_HERSHEY_SIMPLEX, 4, thickness=5, color=color)
+
 
     plt.axis('off')
     ax.imshow(image)
@@ -79,8 +109,11 @@ if image_file is not None:
     pred_dict = predict(model, data, detection_threshold)
     print('Fixing the preds')
     boxes, image = prepare_prediction(pred_dict, nms_threshold)
+
+    print('Predicting classes')
+    labels = predict_class(classifier, image, boxes)
     print('Plotting')
-    plot_img_no_mask(image, boxes)
+    plot_img_no_mask(image, boxes, labels)
 
     img = PIL.Image.open('img.png')
     st.image(img,width=750)

classifier.py

@@ -0,0 +1,45 @@
+import timm
+import torch.nn as nn
+
+import torch
+
+def get_efficientnet(model_name):
+    model = timm.create_model(model_name, pretrained=True)
+
+    return model
+
+class CustomEfficientNet(nn.Module):
+    """
+    This class defines a custom EfficientNet network.
+
+    Parameters
+    ----------
+    target_size : int
+        Number of units for the output layer.
+    pretrained : bool
+        Determine if pretrained weights are used.
+
+    Attributes
+    ----------
+    model : nn.Module
+        EfficientNet model.
+    """
+    def __init__(self, model_name : str = 'efficientnet_b0',
+                 target_size : int = 4, pretrained : bool = True):
+        super().__init__()
+        self.model = timm.create_model(model_name, pretrained=pretrained)
+        
+        # Modify the classifier layer
+        in_features = self.model.classifier.in_features
+        self.model.classifier = 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

@@ -72,4 +72,21 @@ def prepare_prediction(pred_dict, threshold):
     fixed_boxes = torchvision.ops.batched_nms(boxes, scores, labels, threshold)
     boxes = boxes[fixed_boxes, :]
 
-    return boxes, image
+    return boxes, image
+
+def predict_class(model, image, bboxes):
+    preds = []
+
+    for bbox in 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)
+        preds.append(y_preds.softmax(1).detach().numpy())
+
+    preds = np.concatenate(preds).argmax(1)
+
+    return preds