First version

.gitattributes

@@ -25,3 +25,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text

model.py

@@ -0,0 +1,69 @@
+from icevision import *
+import collections
+import cv2
+import PIL
+import torch
+import numpy as np
+import torchvision
+
+MODEL_TYPE = models.ross.efficientdet
+
+def get_model(checkpoint_path):
+    extra_args = {}
+    backbone = MODEL_TYPE.backbones.d0
+    # The efficientdet model requires an img_size parameter
+    extra_args['img_size'] = 512
+
+    model = MODEL_TYPE.model(backbone=backbone(pretrained=True),
+                             num_classes=2, 
+                             **extra_args)
+    
+    ckpt = get_checkpoint(checkpoint_path)
+    model.load_state_dict(ckpt)
+
+    return model
+
+def get_checkpoint(checkpoint_path):
+    ckpt = torch.load('checkpoint.ckpt')
+
+    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, image):
+    img = PIL.Image.open(image)
+    class_map = ClassMap(classes=['Waste'])
+    transforms = tfms.A.Adapter([
+                    *tfms.A.resize_and_pad(512),
+                    tfms.A.Normalize()
+                ])
+
+    pred_dict  = MODEL_TYPE.end2end_detect(img,
+                                           transforms, 
+                                           model,
+                                           class_map=class_map,
+                                           detection_threshold=0.5,
+                                           return_as_pil_img=False,
+                                           return_img=True,
+                                           display_bbox=False,
+                                           display_score=False,
+                                           display_label=False)
+
+    return pred_dict
+
+def prepare_prediction(pred_dict):
+    boxes = [box.to_tensor() for box in pred_dict['detection']['bboxes']]
+    boxes = torch.stack(boxes)
+
+    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, 0.1)
+    boxes = boxes[fixed_boxes, :]
+
+    return boxes, image

requirements.txt

@@ -0,0 +1,2 @@
+icevision[full]
+matplotlib

web_app.py

@@ -0,0 +1,41 @@
+import streamlit as st
+import matplotlib.pyplot as plt
+import numpy as np
+import cv2
+import PIL
+
+from model import get_model, predict, prepare_prediction
+
+print('Creating the model')
+model = get_model('checkpoint.ckpt')
+
+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):            
+        [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)
+    plt.axis('off')
+    ax.imshow(image)
+    fig.savefig("img.png", bbox_inches='tight')
+
+image_file = st.file_uploader("Upload Images", type=["png","jpg","jpeg"])
+
+if image_file is not None:
+    print(image_file)
+    print('Getting predictions')
+    pred_dict = predict(model, image_file)
+    print('Fixing the preds')
+    boxes, image = prepare_prediction(pred_dict)
+    print('Plotting')
+    plot_img_no_mask(image, boxes)
+
+    img = PIL.Image.open('img.png')
+    st.image(img,width=750)