refactor: Using streamlit again

app.py

@@ -1,64 +1,75 @@
-import gradio as gr
-from gradio.networking import get_first_available_port
+import streamlit as st
 import PIL
 import torch
-import os
 
 from utils import plot_img_no_mask, get_models
 from classifier import CustomEfficientNet, CustomViT
 from model import get_model, predict, prepare_prediction, predict_class
 
-os.system('pkill -9 python')
-
 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)
-    print('Getting predictions')
-    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')
 
-    return plot_img_no_mask(image, boxes, labels)
+st.subheader('Upload Custom Image')
 
-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"),
-]
+image_file = st.file_uploader("Upload Images", type=["png","jpg","jpeg"])
 
-outputs = [
-    gr.outputs.Image(type="plot", label="Prediction"),
-]
+st.subheader('Example Images')
 
-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]
+example_imgs = [
+    'example_imgs/basura_4_2.jpg',
+    'example_imgs/basura_1.jpg',
+    'example_imgs/basura_3.jpg'
 ]
 
-gr.close_all()
-#port = get_first_available_port(7682, 9000)
+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]
 
-gr.Interface(
-    waste_detector_interface,
-    inputs,
-    outputs,
-    title=title,
-    description=description,
-    examples=examples,
-    theme="huggingface"
-).launch(share=True)
+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:
+    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)
+    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)
 
-os.system('python3 app.py')
+    img = PIL.Image.open('img.png')
+    st.image(img,width=750)

model.py

@@ -39,7 +39,11 @@ def get_checkpoint(checkpoint_path : str):
 
     return fixed_state_dict
 
-def predict(model : object, img : Union[str, BytesIO], detection_threshold : float):
+def predict(model : object, image : Union[str, BytesIO], detection_threshold : float):
+    img = PIL.Image.open(image)
+    #img = PIL.Image.open(BytesIO(image))
+    img = np.array(img)
+    img = PIL.Image.fromarray(img)
     class_map = ClassMap(classes=['Waste'])
     transforms = tfms.A.Adapter([
                     *tfms.A.resize_and_pad(512),

requirements.txt

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

utils.py

@@ -45,11 +45,10 @@ def plot_img_no_mask(image : np.ndarray, boxes : torch.Tensor, labels):
         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
+    fig.savefig("img.png", bbox_inches='tight')
 
 def get_models(
     detection_ckpt : str,