Update app.py

app.py

@@ -17,12 +17,15 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/blob/main/best.pt'
+# IMPORTANT: Ensure the model file is a valid PyTorch model file.
+# If using a URL, it must directly point to the raw .pt file, not an HTML page.
+model_path = 'best.pt'  # Change this to a valid path or direct URL to your model file
 try:
     model = YOLO(model_path)
 except Exception as ex:
     st.error(f"Unable to load model from: {model_path}")
     st.error(ex)
+    st.stop()  # Stop the app if the model isn't loaded
 
 # -------------------------------
 # App Title and Description
@@ -37,7 +40,7 @@ Use this app to detect wildfires in images, videos, updating image URLs, or live
 tabs = st.tabs(["File Upload", "Image URL", "YouTube Live Stream"])
 
 # ===============================
-# Tab 1: File Upload Mode
+# Tab 1: File Upload Mode (with auto-detection for images)
 with tabs[0]:
     st.header("Detect Wildfire from Uploaded File")
     col_input, col_result = st.columns(2)
@@ -52,19 +55,26 @@ with tabs[0]:
     if uploaded_file:
         file_type = uploaded_file.type.split('/')[0]
         if file_type == "image":
+            # Load image and allow toggle between Original and Detection result
             image = PIL.Image.open(uploaded_file)
-            col_input.image(image, caption="Uploaded Image", use_column_width=True)
-            if st.button("Detect Wildfire", key="detect_file_image"):
+            display_option = st.radio("Display Mode", options=["Original", "Detection"], index=1)
+            
+            if display_option == "Original":
+                col_result.image(image, caption="Original Image", use_column_width=True)
+            else:
+                # Auto-detect without a button click
                 results = model.predict(image, conf=file_confidence)
                 annotated = results[0].plot()[:, :, ::-1]
                 col_result.image(annotated, caption="Detection Result", use_column_width=True)
                 with col_result.expander("Detection Details"):
                     for box in results[0].boxes:
                         col_result.write("Box (xywh):", box.xywh)
+                        
         elif file_type == "video":
             tfile = tempfile.NamedTemporaryFile(delete=False)
             tfile.write(uploaded_file.read())
             cap = cv2.VideoCapture(tfile.name)
+            # For videos, detection is triggered with a button due to processing cost
             if st.button("Detect Wildfire in Video", key="detect_file_video"):
                 while cap.isOpened():
                     ret, frame = cap.read()
@@ -114,7 +124,6 @@ with tabs[1]:
                 annotated = results[0].plot()[:, :, ::-1]
                 placeholder.image(annotated, channels="BGR", use_column_width=True)
                 time.sleep(1)  # update interval in seconds
-                # Rerun the loop by breaking out of script execution cycle.
                 st.experimental_rerun()
             except Exception as e:
                 col_result.error(f"Error fetching image: {e}")
@@ -149,7 +158,6 @@ with tabs[2]:
                     if not cap.isOpened():
                         st.error("Unable to open YouTube live stream.")
                     else:
-                        # Provide a stop button for live stream detection
                         stop_live = st.button("Stop Live Detection", key="yt_stop")
                         while cap.isOpened() and not stop_live:
                             ret, frame = cap.read()
@@ -160,7 +168,6 @@ with tabs[2]:
                             annotated = results[0].plot()[:, :, ::-1]
                             col_result.image(annotated, channels="BGR", use_column_width=True)
                             time.sleep(0.05)
-                            # Check for stop command during loop
                             stop_live = st.button("Stop Live Detection", key="yt_stop_loop")
                         cap.release()
         except Exception as e: