Implemented gradio

app.py

@@ -1,4 +1,4 @@
-import streamlit as st
+import gradio as gr
 import PIL
 import torch
 
@@ -9,65 +9,47 @@ from model import get_model, predict, prepare_prediction, predict_class
 DET_CKPT = 'efficientDet_icevision.ckpt'
 CLASS_CKPT = 'class_ViT_taco_7_class.pth'
 
-det_model, classifier = get_models(DET_CKPT, CLASS_CKPT)
-
-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(det_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(share=True, enable_queue=True)

requirements.txt

@@ -1,5 +1,6 @@
 icevision[all]
 matplotlib
 effdet
+gradio
 streamlit==1.2.0
 Pillow==8.4.0