refactor: Using streamlit again

app.py

@@ -1,4 +1,4 @@
-import gradio as gr
+import streamlit as st
 import PIL
 import torch
 
@@ -9,47 +9,65 @@ from model import get_model, predict, prepare_prediction, predict_class
 DET_CKPT = 'efficientDet_icevision.ckpt'
 CLASS_CKPT = 'class_ViT_taco_7_class.pth'
 
-def waste_detector_interface(
-    image,
-    detection_threshold,
-    nms_threshold
-):  
-    det_model, classifier = get_models(DET_CKPT, CLASS_CKPT)
+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:
     print('Getting predictions')
-    pred_dict = predict(det_model, image, detection_threshold)
+    if isinstance(image_file, str):
+        data = image_file
+    else:
+        data = image_file.read()
+    pred_dict = predict(det_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, labels)
 
-    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)
+    img = PIL.Image.open('img.png')
+    st.image(img,width=750)