Update app.py

app.py

@@ -36,7 +36,13 @@ def check_file(file, suffix=''):
             ).st_size > 0, f'File download failed: {url}'  # check
         return file
 
-
+def read_pretrain(path):
+    return torch.hub.load('ultralytics/yolov5', 'custom', path=path
+    
+default_pretrained = '2022.11.04-YOLOv5x6_1280-Hololive_Waifu_Classification.pt'
+with st.spinner('Loading the model...'):
+    read_pretrain(default_pretrained)
+    
 st.title("Hololive Waifu Classification")
 
 image = st.text_input('Image URL', '')
@@ -51,7 +57,7 @@ st.info(
     ' - https://megapx-assets.dcard.tw/images/65993ab1-fe08-43be-87cd-2ecd201cacbd/1280.jpeg'
     '\n \n \n'
     ' - https://img.esportsku.com/wp-content/uploads//2021/07/hololive-en.png')
-pretrained = st.selectbox('Select pre-trained', ('2022.11.01-YOLOv5x6_1280-Hololive_Waifu_Classification.pt', 'last.pt'))
+pretrained = st.selectbox('Select pre-trained', ('2022.11.04-YOLOv5x6_1280-Hololive_Waifu_Classification.pt', '2022.11.01-YOLOv5x6_1280-Hololive_Waifu_Classification.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)
@@ -64,9 +70,10 @@ if st.button('Excute'):
     with st.spinner('Loading the 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=os.path.join('pretrained', pretrained))
-        model = torch.hub.load('ultralytics/yolov5', 'custom', path=pretrained)
+    if default_pretrained != pretrained:
+        with st.spinner('Loading the model...'):
+            # model = torch.hub.load('ultralytics/yolov5', 'custom', path=os.path.join('pretrained', pretrained))
+            model = torch.hub.load('ultralytics/yolov5', 'custom', path=pretrained)
     with st.spinner('Updating configuration...'):
         model.conf = float(conf)
         model.max_det = int(max_det)