Add video input tab

app.py

@@ -2,6 +2,9 @@ import gradio as gr
 from PIL import Image
 import numpy as np
 from ultralytics import YOLO
+import cv2
+import tempfile
+import os
 
 # 載入模型
 model = YOLO('last.pt')  # 使用您的模型
@@ -25,22 +28,78 @@ def process_image(input_image):
     
     return result_pil
 
+def process_video(input_video):
+    # 創建臨時檔案來保存處理後的影片
+    temp_output = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
+    output_path = temp_output.name
+    temp_output.close()
+    
+    # 讀取輸入影片
+    cap = cv2.VideoCapture(input_video)
+    
+    # 獲取影片屬性
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    
+    # 創建影片寫入器
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+    
+    # 處理每一幀
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+            
+        # 使用模型進行預測
+        results = model.predict(frame)
+        result = results[0]
+        
+        # 獲取處理後的幀
+        result_frame = result.plot()
+        
+        # 寫入處理後的幀
+        out.write(result_frame)
+    
+    # 釋放資源
+    cap.release()
+    out.release()
+    
+    return output_path
+
 # 創建 Gradio 介面
 with gr.Blocks() as demo:
     gr.Markdown("# 磁磚檢測系統")
     
-    with gr.Row():
-        input_image = gr.Image(label="上傳圖片")
-        output_image = gr.Image(label="檢測結果")
-    
-    submit_btn = gr.Button("開始檢測")
-    
-    # 設置事件處理
-    submit_btn.click(
-        fn=process_image,
-        inputs=input_image,
-        outputs=output_image
-    )
+    with gr.Tabs():
+        with gr.TabItem("圖片檢測"):
+            with gr.Row():
+                input_image = gr.Image(label="上傳圖片")
+                output_image = gr.Image(label="檢測結果")
+            
+            submit_btn = gr.Button("開始檢測")
+            
+            # 設置事件處理
+            submit_btn.click(
+                fn=process_image,
+                inputs=input_image,
+                outputs=output_image
+            )
+        
+        with gr.TabItem("影片檢測"):
+            with gr.Row():
+                input_video = gr.Video(label="上傳影片")
+            
+            video_submit_btn = gr.Button("開始檢測")
+            output_file = gr.File(label="下載處理後的影片")
+            
+            # 設置事件處理
+            video_submit_btn.click(
+                fn=process_video,
+                inputs=input_video,
+                outputs=output_file
+            )
 
 if __name__ == "__main__":
     demo.launch()