Update src/streamlit_app.py

src/streamlit_app.py

@@ -48,6 +48,8 @@ def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scale
     return img, ratio, (dw, dh)
 
 def detect_modify(img0, model, device, conf=0.4, imgsz=640, conf_thres=0.25, iou_thres=0.45):
+    st.image(img0, caption="Your image", use_column_width=True)
+
     stride = int(model.stride.max())
     imgsz = check_img_size(imgsz, s=stride)
 
@@ -79,11 +81,6 @@ def detect_modify(img0, model, device, conf=0.4, imgsz=640, conf_thres=0.25, iou
 
     img0 = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB)
     st.image(img0, caption="Prediction Result", use_column_width=True)
-
-# Streamlit 介面
-st.title("YOLOv7 Mask Detection")
-st.write("Detect whether a person is wearing a face mask or not.")
-
 # 取得目前檔案所在目錄
 current_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -98,38 +95,29 @@ iou_thres = 0.45
 device = torch.device("cpu")
 
 # 載入模型
-@st.cache_resource  # 使用 cache 避免每次重新載入
-def load_model():
-    ckpt = torch.load(weight_path, map_location=device, weights_only=False)
-    model = ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()
-    return model
+ckpt = torch.load(weight_path, map_location=device, weights_only=False)
+model = ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()
 
-model = load_model()
+# Streamlit 介面
+st.title("YOLOv7 Mask Detection")
+st.write("Detect whether a person is wearing a face mask or not.")
 
 option = st.radio("Select Input Method", ["Upload Image", "Image URL"])
 
 if option == "Upload Image":
     uploaded_file = st.file_uploader("Please upload an image.", type=["jpg", "jpeg", "png"])
-    
     if uploaded_file is not None:
-        # 直接處理上傳的檔案，不需要額外按鈕
         img = PILImage.create(uploaded_file)
-        st.image(img, caption="Uploaded image", use_column_width=True)
-        
-        # 增加一個進度提示
-        with st.spinner('Running detection...'):
-            detect_modify(img, model, device, conf=conf, imgsz=imgsz, conf_thres=conf_thres, iou_thres=iou_thres)
+        detect_modify(img, model, device, conf=conf, imgsz=imgsz, conf_thres=conf_thres, iou_thres=iou_thres)
 
 elif option == "Image URL":
     url = st.text_input("Please input an image URL.")
     if url:
         try:
-            with st.spinner('Downloading and processing image...'):
-                response = requests.get(url)
-                response.raise_for_status()  # 檢查 http status
-                pil_img = PILImage.create(BytesIO(response.content))
-                st.image(pil_img, caption="Downloaded image", use_column_width=True)
-                detect_modify(pil_img, model, device, conf=conf, imgsz=imgsz, conf_thres=conf_thres, iou_thres=iou_thres)
+            response = requests.get(url)
+            response.raise_for_status()  # 檢查 http status
+            pil_img = PILImage.create(BytesIO(response.content))
+            detect_modify(pil_img, model, device, conf=conf, imgsz=imgsz, conf_thres=conf_thres, iou_thres=iou_thres)
         except Exception as e:
             st.error(f"Problem reading image from URL: {url}")
             st.error(str(e))