Spaces:
Runtime error
Runtime error
File size: 2,010 Bytes
f890c24 cd4c90e 9fbf078 cd4c90e de2e31f 3fa54be cd4c90e de2e31f 9fbf078 f890c24 ba92502 f890c24 ba92502 f890c24 ba92502 f890c24 ba92502 f890c24 95d9d45 f890c24 d27c41a f890c24 6438514 f890c24 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import streamlit as st
import PIL
import torch
from utils import plot_img_no_mask, get_models
from model import predict, prepare_prediction, predict_class
DET_CKPT = 'efficientDet_icevision.ckpt'
CLASS_CKPT = 'class_ViT_taco_7_class.pth'
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:
det_model, classifier = get_models(DET_CKPT, CLASS_CKPT)
print('Getting predictions')
pred_dict = predict(det_model, image_file, 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) |