src/streamlit_app.py

import streamlit as st
import cv2
from ultralytics import YOLO
import tempfile
import os
from pathlib import Path

# 設定頁面配置
st.set_page_config(
    page_title="YOLO 物體偵測",
    page_icon="🎯",
    layout="wide"
)

@st.cache_resource
def load_model():
    """載入 YOLO 模型"""
    try:
        # 使用 YOLOv8n 模型（最小版本，適合 Hugging Face Spaces）
        model = YOLO("yolov8n.pt")
        return model
    except Exception as e:
        st.error(f"❌ 載入模型失敗: {e}")
        return None

def _display_detected_frames(conf, model, st_frame, image, frame_count):
    """處理並顯示偵測結果"""
    # 調整影像大小（更小以節省記憶體）
    height, width = image.shape[:2]
    new_width = 640  # 縮小尺寸
    new_height = int(new_width * height / width)
    image = cv2.resize(image, (new_width, new_height))
    
    # 執行物體偵測
    results = model.predict(image, conf=conf, verbose=False)
    
    # 繪製偵測結果
    annotated_image = results[0].plot()
    
    # 顯示結果
    st_frame.image(
        annotated_image,
        caption=f'偵測結果 - 第 {frame_count} 幀',
        channels="BGR",
        use_container_width=True
    )
    
    return results

# 主要介面
st.title("🎯 YOLO 物體偵測應用程式")
st.markdown("*適用於 Hugging Face Spaces 的輕量版本*")
st.markdown("---")

# 載入模型
with st.spinner("正在載入 YOLO 模型..."):
    model = load_model()

if model is None:
    st.error("模型載入失敗，請重新整理頁面再試")
    st.stop()
else:
    st.success("✅ 模型載入成功！")

# 側邊欄設定
st.sidebar.header("⚙️ 偵測設定")
confidence_threshold = st.sidebar.slider(
    "信心度閾值",
    min_value=0.1,
    max_value=1.0,
    value=0.5,
    step=0.05,
    help="只顯示信心度高於此閾值的偵測結果"
)

# 處理幀數限制（節省資源）
max_frames = st.sidebar.number_input(
    "最大處理幀數",
    min_value=10,
    max_value=300,
    value=100,
    step=10,
    help="限制處理的最大幀數以節省資源"
)

# 顯示使用說明
with st.sidebar.expander("📖 使用說明"):
    st.write("""
    1. 上傳影片檔案（建議小於50MB）
    2. 調整信心度閾值
    3. 設定最大處理幀數
    4. 點擊開始偵測
    """)

# 主要內容區域
col1, col2 = st.columns([1, 1])

with col1:
    st.header("📤 上傳影片")
    source_video = st.file_uploader(
        "選擇影片檔案",
        type=['mp4', 'avi', 'mov', 'mkv'],
        help="建議檔案大小小於 50MB，支援格式: MP4, AVI, MOV, MKV"
    )
    
    if source_video:
        file_size = lenimport streamlit as st
import cv2
from ultralytics import YOLO
import tempfile
import os
import requests
from pathlib import Path

# 設定頁面配置
st.set_page_config(
    page_title="YOLO 物體偵測",
    page_icon="🎯",
    layout="wide"
)

@st.cache_resource
def load_model():
    """載入或下載 YOLO 模型"""
    model_path = "src/yolov8n.pt"  # 使用 YOLOv8n 作為替代，因為 yolo12n.pt 可能不存在
    
    # 如果模型檔案不存在，會自動從 ultralytics 下載
    try:
        model = YOLO(model_path)
        st.success(f"✅ 成功載入模型: {model_path}")
        return model
    except Exception as e:
        st.error(f"❌ 載入模型失敗: {e}")
        return None

def _display_detected_frames(conf, model, st_frame, image, frame_count):
    """處理並顯示偵測結果"""
    # 調整影像大小
    height, width = image.shape[:2]
    new_width = 720
    new_height = int(new_width * height / width)
    image = cv2.resize(image, (new_width, new_height))
    
    # 執行物體偵測
    results = model.predict(image, conf=conf, verbose=False)
    
    # 繪製偵測結果
    annotated_image = results[0].plot()
    
    # 顯示結果
    st_frame.image(
        annotated_image,
        caption=f'偵測結果 - 第 {frame_count} 幀',
        channels="BGR",
        use_container_width=True
    )
    
    return results

# 主要介面
st.title("🎯 YOLO 物體偵測應用程式")
st.markdown("---")

# 載入模型
with st.spinner("正在載入 YOLO 模型..."):
    model = load_model()

if model is None:
    st.stop()

# 側邊欄設定
st.sidebar.header("⚙️ 偵測設定")
confidence_threshold = st.sidebar.slider(
    "信心度閾值",
    min_value=0.1,
    max_value=1.0,
    value=0.6,
    step=0.05,
    help="只顯示信心度高於此閾值的偵測結果"
)

# 顯示模型資訊
with st.sidebar.expander("📊 模型資訊"):
    st.write(f"模型名稱: {model.model_name if hasattr(model, 'model_name') else 'YOLOv8'}")
    st.write(f"信心度閾值: {confidence_threshold}")

# 主要內容區域
col1, col2 = st.columns([1, 1])

with col1:
    st.header("📤 上傳影片")
    source_video = st.file_uploader(
        "選擇影片檔案",
        type=['mp4', 'avi', 'mov', 'mkv', 'wmv'],
        help="支援格式: MP4, AVI, MOV, MKV, WMV"
    )

with col2:
    st.header("📊 處理狀態")
    status_placeholder = st.empty()
    progress_bar = st.progress(0)

# 如果有上傳影片
if source_video:
    st.markdown("---")
    st.header("📹 原始影片預覽")
    st.video(source_video)
    
    # 執行按鈕
    if st.button("🚀 開始物體偵測", type="primary"):
        try:
            # 建立臨時檔案
            with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tfile:
                tfile.write(source_video.read())
                temp_file_path = tfile.name
            
            # 開啟影片檔案
            vid_cap = cv2.VideoCapture(temp_file_path)
            
            if not vid_cap.isOpened():
                st.error("❌ 無法開啟影片檔案")
                st.stop()
            
            # 獲取影片資訊
            total_frames = int(vid_cap.get(cv2.CAP_PROP_FRAME_COUNT))
            fps = vid_cap.get(cv2.CAP_PROP_FPS)
            
            status_placeholder.info(f"📊 影片資訊: {total_frames} 幀, {fps:.1f} FPS")
            
            # 建立結果顯示區域
            st.markdown("---")
            st.header("🎯 即時偵測結果")
            st_frame = st.empty()
            
            frame_count = 0
            detection_results = []
            
            # 處理每一幀
            while vid_cap.isOpened():
                success, image = vid_cap.read()
                
                if success:
                    frame_count += 1
                    
                    # 更新進度條
                    progress = frame_count / total_frames
                    progress_bar.progress(progress)
                    
                    # 執行偵測並顯示結果
                    results = _display_detected_frames(
                        confidence_threshold, 
                        model, 
                        st_frame, 
                        image, 
                        frame_count
                    )
                    
                    # 儲存偵測統計
                    num_detections = len(results[0].boxes) if results[0].boxes is not None else 0
                    detection_results.append(num_detections)
                    
                    # 更新狀態
                    status_placeholder.success(
                        f"✅ 處理進度: {frame_count}/{total_frames} "
                        f"({progress:.1%}) - 偵測到 {num_detections} 個物體"
                    )
                    
                    # 每10幀暫停一下，避免處理太快
                    if frame_count % 10 == 0:
                        import time
                        time.sleep(0.1)
                        
                else:
                    break
            
            # 清理資源
            vid_cap.release()
            os.unlink(temp_file_path)
            
            # 顯示最終統計
            st.markdown("---")
            st.header("📈 偵測統計")
            
            col1, col2, col3, col4 = st.columns(4)
            with col1:
                st.metric("總幀數", total_frames)
            with col2:
                st.metric("處理幀數", frame_count)
            with col3:
                st.metric("平均偵測數", f"{sum(detection_results)/len(detection_results):.1f}")
            with col4:
                st.metric("最大偵測數", max(detection_results) if detection_results else 0)
            
            st.success("🎉 影片處理完成！")
            
        except Exception as e:
            st.error(f"❌ 處理影片時發生錯誤: {str(e)}")
            
            # 清理臨時檔案
            try:
                if 'temp_file_path' in locals():
                    os.unlink(temp_file_path)
            except:
                pass

else:
    st.info("👆 請先上傳一個影片檔案開始偵測")

# 頁腳
st.markdown("---")
st.markdown(
    """
    <div style='text-align: center'>
        <p>🤖 由 YOLO 模型驅動的物體偵測應用程式</p>
        <p><small>支援多種影片格式，即時物體偵測與標註</small></p>
    </div>
    """, 
    unsafe_allow_html=True
)