feature: Implemented example images and cleaned the interface

app.py

@@ -27,15 +27,58 @@ def plot_img_no_mask(image, boxes):
     ax.imshow(image)
     fig.savefig("img.png", bbox_inches='tight')
 
+st.subheader('Upload Custom Image')
+
 image_file = st.file_uploader("Upload Images", type=["png","jpg","jpeg"])
 
+st.subheader('Example Images')
+
+example_imgs = [
+    'example_imgs/basura_4_2.jpg',
+    'example_imgs/basura_2.jpg',
+    'example_imgs/basura_3.jpg'
+]
+
+with st.container() as cont:
+    st.image(example_imgs[0], width=150, caption='1')
+    if st.button('Select Image', key='Image_1'):
+        image_file = example_imgs[0]
+
+with st.container() as cont:
+    st.image(example_imgs[1], width=150, caption='2')
+    if st.button('Select Image', key='Image_2'):
+        image_file = example_imgs[1]
+
+with st.container() as cont:
+    st.image(example_imgs[2], width=150, caption='2')
+    if st.button('Select Image', key='Image_3'):
+        image_file = example_imgs[2]
+
+st.subheader('Detection parameters')
+
+detection_threshold = st.slider('Detection threshold',
+                                min_value=0.0,
+                                max_value=1.0,
+                                value=0.5,
+                                step=0.1)
+
+nms_threshold = st.slider('NMS threshold',
+                        min_value=0.0,
+                        max_value=1.0,
+                        value=0.3,
+                        step=0.1)
+
+st.subheader('Prediction')
+
 if image_file is not None:
-    print(image_file)
     print('Getting predictions')
-    data = image_file.read()
-    pred_dict = predict(model, data)
+    if isinstance(image_file, str):
+        data = image_file
+    else:
+        data = image_file.read()
+    pred_dict = predict(model, data, detection_threshold)
     print('Fixing the preds')
-    boxes, image = prepare_prediction(pred_dict)
+    boxes, image = prepare_prediction(pred_dict, nms_threshold)
     print('Plotting')
     plot_img_no_mask(image, boxes)
 

model.py

@@ -36,9 +36,9 @@ def get_checkpoint(checkpoint_path):
 
     return fixed_state_dict
 
-def predict(model, image):
-    #img = PIL.Image.open(image)
-    img = PIL.Image.open(BytesIO(image))
+def predict(model, image, detection_threshold):
+    img = PIL.Image.open(image)
+    #img = PIL.Image.open(BytesIO(image))
     img = np.array(img)
     img = PIL.Image.fromarray(img)
     
@@ -52,7 +52,7 @@ def predict(model, image):
                                            transforms, 
                                            model,
                                            class_map=class_map,
-                                           detection_threshold=0.5,
+                                           detection_threshold=detection_threshold,
                                            return_as_pil_img=False,
                                            return_img=True,
                                            display_bbox=False,
@@ -61,7 +61,7 @@ def predict(model, image):
 
     return pred_dict
 
-def prepare_prediction(pred_dict):
+def prepare_prediction(pred_dict, threshold):
     boxes = [box.to_tensor() for box in pred_dict['detection']['bboxes']]
     boxes = torch.stack(boxes)
 
@@ -69,7 +69,7 @@ def prepare_prediction(pred_dict):
     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)
+    fixed_boxes = torchvision.ops.batched_nms(boxes, scores, labels, threshold)
     boxes = boxes[fixed_boxes, :]
 
     return boxes, image