Update app.py

app.py

@@ -1,185 +1,160 @@
+import os
+import sys
+import tempfile
 import cv2
-import streamlit as st
+import requests
 from ultralytics import YOLO
-import time
-import numpy as np
-from datetime import datetime
-import pytz
+import streamlit as st
 
-# Page config and header
+# Set page configuration
 st.set_page_config(
-    page_title="Your Original App",
-    page_icon="🔥",
+    page_title="People Tracking with YOLO11-pose",
+    page_icon="👥",
     layout="wide",
     initial_sidebar_state="expanded"
 )
-st.title("Your Original App")
+st.title("People Tracking with YOLO11-pose")
 
-# --- Session State Initialization ---
-if "streams" not in st.session_state:
-    st.session_state.streams = []
-if "num_streams" not in st.session_state:
-    st.session_state.num_streams = 1
-if "confidence" not in st.session_state:
-    st.session_state.confidence = 0.30  # default 30%
-if "target_fps" not in st.session_state:
-    st.session_state.target_fps = 1.0   # default 1 FPS
+# Sidebar: Input method and settings
+st.sidebar.header("Input Settings")
+uploaded_file = st.sidebar.file_uploader("Upload Image/Video", type=["jpg", "jpeg", "png", "bmp", "webp", "mp4"])
+youtube_link = st.sidebar.text_input("YouTube Link (optional)", "")
+image_url = st.sidebar.text_input("Image URL (optional)", "")
+sensitivity = st.sidebar.slider("Sensitivity (Confidence Threshold)", 0.0, 1.0, 0.2, step=0.01)
+process_button = st.sidebar.button("Process Input")
 
-# --- Default URLs and Names for 10 Streams ---
-default_m3u8_urls = [
-    "https://publicstreamer4.cotrip.org/rtplive/070E27890CAM1RHS/playlist.m3u8",   # EB at i270
-    "https://publicstreamer2.cotrip.org/rtplive/070E27555CAM1RP1/playlist.m3u8",  # EB @ York St Denver
-    "https://publicstreamer1.cotrip.org/rtplive/225N00535CAM1RP1/playlist.m3u8",  # NB at Iliff Denver
-    "https://publicstreamer2.cotrip.org/rtplive/070W28220CAM1RHS/playlist.m3u8",  # WB Half Mile West of I225 Denver
-    "https://publicstreamer1.cotrip.org/rtplive/070W26805CAM1RHS/playlist.m3u8",  # 1 mile E of Kipling Denver
-    "https://publicstreamer4.cotrip.org/rtplive/076W03150CAM1RP1/playlist.m3u8",  # Main St Hudson
-    "https://publicstreamer2.cotrip.org/rtplive/070E27660CAM1NEC/playlist.m3u8",   # EB Colorado Blvd i70 Denver
-    "https://publicstreamer2.cotrip.org/rtplive/070W27475CAM1RHS/playlist.m3u8",   # E of Washington St Denver
-    "https://publicstreamer3.cotrip.org/rtplive/070W28155CAM1RHS/playlist.m3u8",   # WB Peroia St Underpass Denver
-    "https://publicstreamer3.cotrip.org/rtplive/070E11660CAM1RHS/playlist.m3u8"    # Grand Ave Glenwood
-]
-default_names = [
-    "EB at i270",
-    "EB @ York St Denver",
-    "NB at Iliff Denver",
-    "WB Half Mile West of I225 Denver",
-    "1 mile E of Kipling Denver",
-    "Main St Hudson",
-    "EB Colorado Blvd i70 Denver",
-    "E of Washington St Denver",
-    "WB Peroia St Underpass Denver",
-    "Grand Ave Glenwood"
-]
+# Define the video extensions for later use
+video_exts = [".mp4", ".mov", ".avi", ".webm"]
 
-# --- Sidebar Settings ---
-with st.sidebar:
-    st.header("Stream Settings")
-    # Custom configuration for stream 1 (optional)
-    custom_m3u8 = st.text_input("Custom M3U8 URL for Stream 1 (optional)", value="", key="custom_m3u8")
-    custom_name = st.text_input("Custom Webcam Name for Stream 1 (optional)", value="", key="custom_name")
-    
-    # Choose number of streams (1 to 10)
-    num_streams = st.selectbox("Number of Streams", list(range(1, 11)), index=0)
-    st.session_state.num_streams = num_streams
+def process_input(uploaded_file, youtube_link, image_url, sensitivity):
+    input_path = None
+    temp_files = []
 
-    # Global settings for confidence and processing rate.
-    confidence = float(st.slider("Confidence Threshold", 5, 100, 30)) / 100
-    st.session_state.confidence = confidence
-    fps_options = {
-        "1 FPS": 1,
-        "1 frame/2s": 0.5,
-        "1 frame/3s": 0.3333,
-        "1 frame/5s": 0.2,
-        "1 frame/15s": 0.0667,
-        "1 frame/30s": 0.0333
-    }
-    video_option = st.selectbox("Processing Rate", list(fps_options.keys()), index=3)
-    st.session_state.target_fps = fps_options[video_option]
-
-# Initialize or update stream state
-if len(st.session_state.streams) != st.session_state.num_streams:
-    st.session_state.streams = []
-    for i in range(st.session_state.num_streams):
-        if i == 0:
-            url = custom_m3u8.strip() if custom_m3u8.strip() else default_m3u8_urls[0]
-            display_name = custom_name.strip() if custom_name.strip() else default_names[0]
-        else:
-            url = default_m3u8_urls[i] if i < len(default_m3u8_urls) else ""
-            display_name = default_names[i] if i < len(default_names) else f"Stream {i+1}"
-        st.session_state.streams.append({
-            "current_m3u8_url": url,
-            "processed_frame": np.zeros((480, 640, 3), dtype=np.uint8),
-            "start_time": time.time(),
-            "processed_count": 0,
-            "detected_frames": [],
-            "last_processed_time": 0,
-            "stats_text": "Processing FPS: 0.00\nFrame Delay: 0.00 sec\nTensor Results: No detections",
-            "highest_match": 0.0,
-            "display_name": display_name
-        })
-else:
-    if st.session_state.num_streams > 0:
-        url = custom_m3u8.strip() if custom_m3u8.strip() else default_m3u8_urls[0]
-        display_name = custom_name.strip() if custom_name.strip() else default_names[0]
-        st.session_state.streams[0]["current_m3u8_url"] = url
-        st.session_state.streams[0]["display_name"] = display_name
+    # Input priority: YouTube link > Image URL > Uploaded file.
+    if youtube_link and youtube_link.strip():
+        try:
+            from pytubefix import YouTube
+            yt = YouTube(youtube_link)
+            stream = yt.streams.filter(file_extension='mp4', progressive=True).order_by("resolution").desc().first()
+            if not stream:
+                return None, None, None, "No suitable mp4 stream found."
+            temp_path = os.path.join(tempfile.gettempdir(), f"yt_{os.urandom(8).hex()}.mp4")
+            stream.download(output_path=tempfile.gettempdir(), filename=os.path.basename(temp_path))
+            input_path = temp_path
+            temp_files.append(input_path)
+        except Exception as e:
+            return None, None, None, f"Error downloading YouTube video: {str(e)}"
+    elif image_url and image_url.strip():
+        try:
+            response = requests.get(image_url, stream=True, timeout=10)
+            response.raise_for_status()
+            temp_path = os.path.join(tempfile.gettempdir(), f"img_{os.urandom(8).hex()}.jpg")
+            with open(temp_path, "wb") as f:
+                f.write(response.content)
+            input_path = temp_path
+            temp_files.append(input_path)
+        except Exception as e:
+            return None, None, None, f"Error downloading image: {str(e)}"
+    elif uploaded_file is not None:
+        # Save the uploaded file to a temporary file
+        ext = os.path.splitext(uploaded_file.name)[1]
+        with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
+            tmp.write(uploaded_file.read())
+            input_path = tmp.name
+            temp_files.append(input_path)
+    else:
+        return None, None, None, "Please provide an input."
 
-confidence = st.session_state.confidence
-target_fps = st.session_state.target_fps
+    ext = os.path.splitext(input_path)[1].lower()
+    output_path = None
 
-# --- Load Model ---
-model_path = 'https://huggingface.co/spaces/tstone87/ccr-colorado/resolve/main/best.pt'
-@st.cache_resource
-def load_model():
-    return YOLO(model_path)
+    # Load the YOLO model (ensure the model file is available in your repository)
+    model = YOLO("yolo11n-pose.pt")
 
-try:
-    model = load_model()
-except Exception as ex:
-    st.error(f"Model loading failed: {str(ex)}")
-    st.stop()
+    try:
+        if ext in video_exts:
+            # Video processing
+            cap = cv2.VideoCapture(input_path)
+            if not cap.isOpened():
+                return None, None, None, f"Cannot open video file: {input_path}"
+            
+            fps = cap.get(cv2.CAP_PROP_FPS)
+            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            
+            if fps <= 0 or width <= 0 or height <= 0:
+                return None, None, None, "Invalid video properties detected."
+            
+            output_path = os.path.join(tempfile.gettempdir(), f"out_{os.urandom(8).hex()}.mp4")
+            # Use 'mp4v' as codec
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+            
+            if not out.isOpened():
+                return None, None, None, "Video processing failed: No suitable encoder available."
+            
+            processed_frames = 0
+            while True:
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                    
+                # Process frame: convert to RGB, run YOLO, then annotate and convert back to BGR.
+                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                results = model.predict(source=frame_rgb, conf=sensitivity)[0]
+                annotated_frame = results.plot()
+                annotated_frame_bgr = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
+                
+                out.write(annotated_frame_bgr)
+                processed_frames += 1
+            
+            cap.release()
+            out.release()
+            temp_files.append(output_path)
+            
+            if processed_frames == 0:
+                return None, None, None, "No frames processed from video."
+                
+            if not os.path.exists(output_path) or os.path.getsize(output_path) < 1024:
+                return None, None, None, f"Output video created but too small ({os.path.getsize(output_path)} bytes) - processing failed."
+                
+            return output_path, None, output_path, f"Video processed successfully! ({processed_frames}/{frame_count} frames)"
 
-# --- Create Placeholders for Streams in a 2-Column Grid ---
-num_streams = st.session_state.num_streams
-feed_placeholders = []
-stats_placeholders = []
-cols = st.columns(2)
-for i in range(num_streams):
-    col_index = i % 2
-    if i >= 2 and col_index == 0:
-        cols = st.columns(2)
-    feed_placeholders.append(cols[col_index].empty())
-    stats_placeholders.append(cols[col_index].empty())
-if num_streams == 1:
-    _ = st.columns(2)
+        else:
+            # Image processing
+            results = model.predict(source=input_path, conf=sensitivity)[0]
+            annotated = results.plot()
+            output_path = os.path.join(tempfile.gettempdir(), f"out_{os.urandom(8).hex()}.jpg")
+            cv2.imwrite(output_path, annotated)
+            temp_files.append(output_path)
+            return output_path, output_path, None, "Image processed successfully!"
 
-def update_stream(i):
-    current_time = time.time()
-    sleep_time = 1.0 / target_fps
-    stream_state = st.session_state.streams[i]
-    if current_time - stream_state["last_processed_time"] >= sleep_time:
-        url = stream_state["current_m3u8_url"]
-        cap = cv2.VideoCapture(url)
-        if not cap.isOpened():
-            stats_placeholders[i].text("Failed to open M3U8 stream.")
-            return
-        ret, frame = cap.read()
-        cap.release()
-        if not ret:
-            stats_placeholders[i].text("Stream interrupted or ended.")
-            return
-        res = model.predict(frame, conf=confidence)
-        processed_frame = res[0].plot()[:, :, ::-1]
-        tensor_info = "No detections"
-        max_conf = 0.0
-        try:
-            boxes = res[0].boxes
-            if boxes is not None and len(boxes) > 0:
-                max_conf = float(boxes.conf.max())
-                tensor_info = f"Detections: {len(boxes)} | Max Confidence: {max_conf:.2f}"
-        except Exception as ex:
-            tensor_info = f"Error extracting detections: {ex}"
-        if max_conf >= stream_state["highest_match"]:
-            stream_state["highest_match"] = max_conf
-            stream_state["detected_frames"].append(processed_frame)
-        stream_state["processed_count"] += 1
-        stream_state["last_processed_time"] = current_time
-        mt_time = datetime.now(pytz.timezone('America/Denver')).strftime('%Y-%m-%d %H:%M:%S MT')
-        stream_state["processed_frame"] = processed_frame
-        stream_state["stats_text"] = (
-            f"Processing FPS: {stream_state['processed_count'] / (current_time - stream_state['start_time']):.2f}\n"
-            f"{tensor_info}\n"
-            f"Highest Match: {stream_state['highest_match']:.2f}"
-        )
-        feed_placeholders[i].image(
-            processed_frame,
-            caption=f"Stream {i+1} - {stream_state['display_name']} - {mt_time}",
-            use_container_width=True
-        )
-        stats_placeholders[i].text(stream_state["stats_text"])
+    except Exception as e:
+        return None, None, None, f"Processing error: {str(e)}"
+    
+    finally:
+        # Clean up temporary files except the final output
+        for f in temp_files[:-1]:
+            if f and os.path.exists(f):
+                try:
+                    os.remove(f)
+                except:
+                    pass
 
-# --- Update Streams and Rerun ---
-for i in range(num_streams):
-    update_stream(i)
-time.sleep(1.0 / target_fps)
-st.experimental_rerun()
+# When the user clicks "Process Input"
+if process_button:
+    out_file, out_img, out_vid, status = process_input(uploaded_file, youtube_link, image_url, sensitivity)
+    st.write(status)
+    if out_img:
+        st.image(out_img, caption="Annotated Output (Image)", use_column_width=True)
+    if out_vid:
+        st.video(out_vid)
+    if out_file:
+        with open(out_file, "rb") as f:
+            st.download_button(
+                label="Download Annotated Output",
+                data=f,
+                file_name=os.path.basename(out_file),
+                mime="video/mp4" if os.path.splitext(out_file)[1].lower() in video_exts else "image/jpeg"
+            )