Update app.py

app.py

@@ -40,25 +40,36 @@ def check_file(file, suffix=''):
 st.title("Hololive Waifu Classification")
 
 image = st.text_input('Image URL', '')
-imgsz = st.text_input('Image Size', 1280)
+st.info(
+    'Images for quick tesing:\n \n \n'
+    ' - https://i.imgur.com/tFZwWYw.jpg'
+    '\n \n \n'
+    ' - https://static.wikia.nocookie.net/omniversal-battlefield/images/b/bd/Council.jpg'
+    , icon="ℹ️")
+pretrained = st.selectbox('Select pre-trained', ('best.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)
 multi_label = st.selectbox('Multiple labels per box', (False, True))
 agnostic = st.selectbox('Class-agnostic', (False, True))
 amp = st.selectbox('Automatic Mixed Precision inference', (False, True))
-max_det  = st.text_input('Maximum number of detections per image', 1000)
+max_det  = st.number_input(label='Maximum number of detections per image', min_value=None, max_value=None, value=1000, step=1)
 
 if st.button('Excute'):
-    image = check_file(image)
-    input_image = Image.open(image)
-    model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
-    model.conf = float(conf)
-    model.max_det = int(max_det)
-    model.iou = float(iou)
-    model.agnostic = agnostic
-    model.multi_label = multi_label
-    model.amp = amp 
-    results = model(input_image, size=int(imgsz))
+    with st.spinner('Loading the image...'):
+        image = check_file(image)
+        input_image = Image.open(image)
+    with st.spinner('Loading the model...'):
+        model = torch.hub.load('ultralytics/yolov5', 'custom', path=pretrained)
+    with st.spinner('Updating configuration...'):
+        model.conf = float(conf)
+        model.max_det = int(max_det)
+        model.iou = float(iou)
+        model.agnostic = agnostic
+        model.multi_label = multi_label
+        model.amp = amp 
+    with st.spinner('Predicting...'):
+        results = model(input_image, size=int(imgsz))
     for img in results.render():
         st.image(img)
     st.write(results.pandas().xyxy[0])