Update src/streamlit_app.py

src/streamlit_app.py

@@ -48,8 +48,6 @@ 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)
 
@@ -81,6 +79,11 @@ 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__))
 
@@ -95,12 +98,13 @@ iou_thres = 0.45
 device = torch.device("cpu")
 
 # 載入模型
-ckpt = torch.load(weight_path, map_location=device, weights_only=False)
-model = ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()
+@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
 
-# Streamlit 介面
-st.title("YOLOv7 Mask Detection")
-st.write("Detect whether a person is wearing a face mask or not.")
+model = load_model()
 
 option = st.radio("Select Input Method", ["Upload Image", "Image URL"])
 
@@ -108,21 +112,26 @@ if option == "Upload Image":
     uploaded_file = st.file_uploader("Please upload an image.", type=["jpg", "jpeg", "png"])
     
     if uploaded_file is not None:
-        st.image(uploaded_file, caption="Uploaded image", use_column_width=True)
-        if st.button("Start Detection"):
-            img = PILImage.create(uploaded_file)
+        # 直接處理上傳的檔案，不需要額外按鈕
+        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)
 
-
 elif option == "Image URL":
     url = st.text_input("Please input an image URL.")
     if url:
         try:
-            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)
+            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)
         except Exception as e:
             st.error(f"Problem reading image from URL: {url}")
             st.error(str(e))
 
+