Update app.py

app.py

@@ -2,65 +2,87 @@ import gradio as gr
 import cv2
 import os
 import uuid
-import mediapipe as mp
-import pandas as pd
 import threading
 import time
+import mediapipe as mp
+import pandas as pd
 
-# Setup
+# === Setup ===
 OUTPUT_DIR = "captured_frames"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 df = pd.DataFrame(columns=["filename", "caption", "pose_coords"])
 
 pose = mp.solutions.pose.Pose()
+
 state = {
-    "video_path": None,
     "cap": None,
     "frame": None,
     "play": False,
-    "thread": None
+    "video_path": None
 }
 
+# === Load Video ===
 def load_video(video_file):
-    if state["cap"]:
-        state["cap"].release()
-    state["video_path"] = video_file.name
-    state["cap"] = cv2.VideoCapture(state["video_path"])
-    state["play"] = False
-    return "Video loaded. Use play/pause to navigate."
-
+    try:
+        if hasattr(video_file, "name"):
+            video_path = video_file.name
+        else:
+            video_path = video_file
+        state["video_path"] = video_path
+
+        if state["cap"]:
+            state["cap"].release()
+
+        state["cap"] = cv2.VideoCapture(video_path)
+        state["frame"] = None
+        state["play"] = False
+        return "✅ Video loaded successfully!"
+    except Exception as e:
+        return f"❌ Error loading video: {e}"
+
+# === Play video in background ===
 def play_video():
-    def play():
-        while state["cap"].isOpened() and state["play"]:
+    if not state["cap"]:
+        return "⚠️ Load a video first."
+    state["play"] = True
+
+    def stream():
+        while state["cap"] and state["cap"].isOpened() and state["play"]:
             ret, frame = state["cap"].read()
             if not ret:
+                state["play"] = False
                 break
             state["frame"] = frame
-            time.sleep(0.1)  # ~10 FPS
+            time.sleep(0.1)
 
-    state["play"] = True
-    state["thread"] = threading.Thread(target=play)
-    state["thread"].start()
-    return "Playing video..."
+    threading.Thread(target=stream).start()
+    return "▶️ Playing..."
 
+# === Pause playback ===
 def pause_video():
     state["play"] = False
-    return "Video paused."
+    return "⏸️ Paused."
 
+# === Show current frame ===
 def show_frame():
     if state["frame"] is not None:
         return state["frame"][:, :, ::-1]  # BGR to RGB
     return None
 
+# === Capture frame (auto-pause) ===
 def capture_frame(caption):
     if state["frame"] is None:
-        return "No frame to capture.", None
+        return "⚠️ No frame to capture.", None
+
+    # Auto-pause video
+    state["play"] = False
 
     frame = state["frame"]
     filename = f"{uuid.uuid4().hex[:8]}.jpg"
     path = os.path.join(OUTPUT_DIR, filename)
     cv2.imwrite(path, frame)
 
+    # Pose estimation
     frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
     results = pose.process(frame_rgb)
 
@@ -76,13 +98,15 @@ def capture_frame(caption):
         "pose_coords": coords
     }])], ignore_index=True)
 
-    return f"Frame captured: {filename}", frame[:, :, ::-1]
+    return f"✅ Captured & paused: {filename}", frame[:, :, ::-1]
 
+# === Download CSV ===
 def download_csv():
-    csv_path = os.path.join(OUTPUT_DIR, "pose_dataset.csv")
-    df.to_csv(csv_path, index=False)
-    return csv_path
+    path = os.path.join(OUTPUT_DIR, "pose_dataset.csv")
+    df.to_csv(path, index=False)
+    return path
 
+# === Reset all ===
 def reset_all():
     df.drop(df.index, inplace=True)
     for f in os.listdir(OUTPUT_DIR):
@@ -90,36 +114,38 @@ def reset_all():
     if state["cap"]:
         state["cap"].release()
     state.update({"video_path": None, "cap": None, "frame": None, "play": False})
-    return "App reset.", None
+    return "🔁 Reset done.", None
 
+# === UI ===
 with gr.Blocks() as app:
-    gr.Markdown("## 🏹 Archery Dataset Tool (Play, Pause, Pose + Annotation)")
+    gr.Markdown("## 🏹 Archery Pose Dataset Tool (Auto-Pause on Capture)")
 
-    video_input = gr.Video(label="Upload Video")
-    load_btn = gr.Button("Load Video")
+    video_input = gr.Video(label="🎞️ Upload Video")
+    load_btn = gr.Button("📂 Load Video")
+    status = gr.Textbox(label="Status")
 
     with gr.Row():
         play_btn = gr.Button("▶️ Play")
         pause_btn = gr.Button("⏸️ Pause")
-        show_btn = gr.Button("🔁 Show Current Frame")
+        show_btn = gr.Button("🖼️ Show Current Frame")
 
-    image_box = gr.Image(label="Current Frame")
+    image_output = gr.Image(label="Current Frame")
     caption_input = gr.Textbox(label="Caption")
-    capture_btn = gr.Button("📸 Capture & Annotate Frame")
-
-    download_btn = gr.Button("📥 Download CSV")
-    csv_file = gr.File(label="Pose Dataset CSV")
+    capture_btn = gr.Button("📸 Capture & Pause")
 
-    reset_btn = gr.Button("🔄 Reset")
-    status = gr.Textbox(label="Status")
+    with gr.Row():
+        download_btn = gr.Button("📥 Download CSV")
+        reset_btn = gr.Button("🔄 Reset")
+    
+    csv_file = gr.File(label="📄 Dataset CSV")
 
-    # Event handlers
+    # Bind actions
     load_btn.click(load_video, inputs=video_input, outputs=status)
     play_btn.click(play_video, outputs=status)
     pause_btn.click(pause_video, outputs=status)
-    show_btn.click(show_frame, outputs=image_box)
-    capture_btn.click(capture_frame, inputs=caption_input, outputs=[status, image_box])
+    show_btn.click(show_frame, outputs=image_output)
+    capture_btn.click(capture_frame, inputs=caption_input, outputs=[status, image_output])
     download_btn.click(download_csv, outputs=csv_file)
-    reset_btn.click(reset_all, outputs=[status, image_box])
+    reset_btn.click(reset_all, outputs=[status, image_output])
 
 app.launch()