feature: Implemented gradio

app.py

@@ -1,121 +1,55 @@
-import streamlit as st
-import matplotlib.pyplot as plt
-import numpy as np
-import cv2
+import gradio as gr
 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
 
-print('Creating the model')
-model = get_model('efficientDet_icevision.ckpt')
-print('Loading the classifier')
-classifier = CustomViT(target_size=7, pretrained=False)
-classifier.load_state_dict(torch.load('class_ViT_taco_7_class.pth', map_location='cpu'))
-# Set eval mode to deactivate dropout and BN layers 
-classifier.eval()
+DET_CKPT = 'efficientDet_icevision.ckpt'
+CLASS_CKPT = 'class_ViT_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'
-    }
-
-    # 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):
-        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, 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)
-    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_1.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:
+def waste_detector_interface(
+    image,
+    detection_threshold,
+    nms_threshold
+):  
+    det_model, classifier = get_models(DET_CKPT, CLASS_CKPT)
     print('Getting predictions')
-    if isinstance(image_file, str):
-        data = image_file
-    else:
-        data = image_file.read()
-    pred_dict = predict(model, data, detection_threshold)
+    pred_dict = predict(det_model, image, 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, labels)
 
-    img = PIL.Image.open('img.png')
-    st.image(img,width=750)
+    return plot_img_no_mask(image, boxes, labels)
+
+inputs = [
+    gr.inputs.Image(type="pil", label="Original Image"),
+    gr.inputs.Number(default=0.5, label="detection_threshold"),
+    gr.inputs.Number(default=0.5, label="nms_threshold"),
+]
+
+outputs = [
+    gr.outputs.Image(type="plot", label="Prediction"),
+]
+
+title = 'Waste Detection'
+description = 'Demo for waste object detection. It detects and classify waste in images according to which rubbish bin the waste should be thrown. Upload an image or click an image to use.'
+examples = [
+    ['example_imgs/basura_4_2.jpg', 0.5, 0.5],
+    ['example_imgs/basura_1.jpg', 0.5, 0.5],
+    ['example_imgs/basura_3.jpg', 0.5, 0.5]
+]
+
+gr.Interface(
+    waste_detector_interface,
+    inputs,
+    outputs,
+    title=title,
+    description=description,
+    examples=examples,
+    theme="huggingface",
+).launch(debug=True, enable_queue=True)

model.py

@@ -1,4 +1,5 @@
 from io import BytesIO
+from typing import Union
 from icevision import *
 import collections
 import PIL
@@ -12,7 +13,7 @@ import icevision.models.ross.efficientdet
 
 MODEL_TYPE = icevision.models.ross.efficientdet
 
-def get_model(checkpoint_path):
+def get_model(checkpoint_path : str):
     extra_args = {}
     backbone = MODEL_TYPE.backbones.d0
     # The efficientdet model requires an img_size parameter
@@ -27,8 +28,8 @@ def get_model(checkpoint_path):
 
     return model
 
-def get_checkpoint(checkpoint_path):
-    ckpt = torch.load('checkpoint.ckpt', map_location=torch.device('cpu'))
+def get_checkpoint(checkpoint_path : str):
+    ckpt = torch.load(checkpoint_path, map_location=torch.device('cpu'))
 
     fixed_state_dict = collections.OrderedDict()
 
@@ -38,15 +39,7 @@ def get_checkpoint(checkpoint_path):
 
     return fixed_state_dict
 
-def predict(model, image, detection_threshold):
-    if isinstance(image, str):
-        img = PIL.Image.open(image)
-    else:
-        img = PIL.Image.open(BytesIO(image))
-        
-    img = np.array(img)
-    img = PIL.Image.fromarray(img)
-    
+def predict(model : object, img : Union[str, BytesIO], detection_threshold : float):
     class_map = ClassMap(classes=['Waste'])
     transforms = tfms.A.Adapter([
                     *tfms.A.resize_and_pad(512),

utils.py

@@ -0,0 +1,79 @@
+from typing import Tuple
+import matplotlib.pyplot as plt
+import numpy as np
+import cv2
+import torch
+
+from  classifier import CustomViT
+from model import get_model
+
+def plot_img_no_mask(image : np.ndarray, 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'
+    }
+
+    # 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):
+        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, 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)
+
+    return fig
+
+def get_models(
+    detection_ckpt : str,
+    classifier_ckpt : str
+) -> Tuple[torch.nn.Module, torch.nn.Module]:
+    """
+    Get the detection and classifier models
+
+    Args:
+        detection_ckpt (str): Detection model checkpoint
+        classifier_ckpt (str): Classifier model checkpoint
+
+    Returns:
+        tuple: Tuple containing:
+            - (torch.nn.Module): Detection model
+            - (torch.nn.Module): Classifier model
+    """
+    print('Loading the detection model')
+    det_model = get_model(detection_ckpt)
+    det_model.eval()
+
+    print('Loading the classifier model')
+    classifier = CustomViT(target_size=7, pretrained=False)
+    classifier.load_state_dict(torch.load(classifier_ckpt, map_location='cpu'))
+    classifier.eval()
+
+    return det_model, classifier