Upload app.py

app.py

@@ -45,7 +45,7 @@ st.info(
     ' - https://i.imgur.com/tFZwWYw.jpg'
     '\n \n \n'
     ' - https://static.wikia.nocookie.net/omniversal-battlefield/images/b/bd/Council.jpg')
-pretrained = st.selectbox('Select pre-trained', ('best.pt', 'last.pt'))
+pretrained = st.selectbox('Select pre-trained', ('2022.11.01-YOLOv5x6_1280-Hololive_Waifu_Classification.pt', 'last.pt'))
 imgsz = st.number_input(label='Image Size', min_value=None, max_value=None, value=1280, step=1)
 conf = st.slider(label='Confidence threshold', min_value=0.0, max_value=1.0, value=0.25, step=0.01)
 iou = st.slider(label='IoU threshold', min_value=0.0, max_value=1.0, value=0.45, step=0.01)
@@ -56,10 +56,10 @@ max_det  = st.number_input(label='Maximum number of detections per image', min_v
 
 if st.button('Excute'):
     with st.spinner('Loading the image...'):
-        image = check_file(image)
-        input_image = Image.open(image)
+        image_path = check_file(image)
+        input_image = Image.open(image_path)
     with st.spinner('Loading the model...'):
-        model = torch.hub.load('ultralytics/yolov5', 'custom', path=pretrained)
+        model = torch.hub.load('ultralytics/yolov5', 'custom', path=os.path.join('pretrained', pretrained))
     with st.spinner('Updating configuration...'):
         model.conf = float(conf)
         model.max_det = int(max_det)
@@ -71,4 +71,5 @@ if st.button('Excute'):
         results = model(input_image, size=int(imgsz))
     for img in results.render():
         st.image(img)
-    st.write(results.pandas().xyxy[0])
+    st.write(results.pandas().xyxy[0])
+    os.remove(image_path)