Update app.py

app.py

@@ -6,6 +6,8 @@ import threading
 import time
 import mediapipe as mp
 import pandas as pd
+from concurrent.futures import ThreadPoolExecutor
+import queue
 
 # === Setup ===
 OUTPUT_DIR = "captured_frames"
@@ -17,10 +19,41 @@ pose = mp.solutions.pose.Pose()
 state = {
     "cap": None,
     "frame": None,
+    "frame_rgb": None,  # Pre-converted RGB frame
     "play": False,
-    "video_path": None
+    "video_path": None,
+    "capture_queue": queue.Queue(),
+    "processing_thread": None
 }
 
+# Thread pool for background processing
+executor = ThreadPoolExecutor(max_workers=2)
+
+# === Background pose processing ===
+def process_pose_async(frame_bgr, filename, caption):
+    """Process pose estimation in background thread"""
+    try:
+        # Convert to RGB for MediaPipe
+        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+        results = pose.process(frame_rgb)
+        
+        coords = []
+        if results.pose_landmarks:
+            for lm in results.pose_landmarks.landmark:
+                coords.append((round(lm.x, 5), round(lm.y, 5), round(lm.z, 5)))
+        
+        # Add to dataframe
+        global df
+        new_row = pd.DataFrame([{
+            "filename": filename,
+            "caption": caption,
+            "pose_coords": coords
+        }])
+        df = pd.concat([df, new_row], ignore_index=True)
+        
+    except Exception as e:
+        print(f"Error processing pose: {e}")
+
 # === Load Video ===
 def load_video(video_file):
     try:
@@ -34,7 +67,12 @@ def load_video(video_file):
             state["cap"].release()
 
         state["cap"] = cv2.VideoCapture(video_path)
+        
+        # Set buffer size to reduce lag
+        state["cap"].set(cv2.CAP_PROP_BUFFERSIZE, 1)
+        
         state["frame"] = None
+        state["frame_rgb"] = None
         state["play"] = False
         return "✅ Video loaded successfully!"
     except Exception as e:
@@ -52,10 +90,19 @@ def play_video():
             if not ret:
                 state["play"] = False
                 break
+            
+            # Store both BGR and RGB versions
             state["frame"] = frame
-            time.sleep(0.1)
-
-    threading.Thread(target=stream).start()
+            state["frame_rgb"] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            
+            # Adaptive delay based on video FPS
+            fps = state["cap"].get(cv2.CAP_PROP_FPS)
+            if fps > 0:
+                time.sleep(1.0 / fps)
+            else:
+                time.sleep(0.033)  # ~30 FPS fallback
+
+    threading.Thread(target=stream, daemon=True).start()
     return "▶️ Playing..."
 
 # === Pause playback ===
@@ -65,40 +112,35 @@ def pause_video():
 
 # === Show current frame ===
 def show_frame():
-    if state["frame"] is not None:
-        return state["frame"][:, :, ::-1]  # BGR to RGB
+    if state["frame_rgb"] is not None:
+        return state["frame_rgb"]  # Already in RGB
     return None
 
-# === Capture frame (auto-pause) ===
+# === Fast capture frame (immediate pause + async processing) ===
 def capture_frame(caption):
     if state["frame"] is None:
         return "⚠️ No frame to capture.", None
 
-    # Auto-pause video
+    # IMMEDIATE pause - this is the key optimization
     state["play"] = False
-
-    frame = state["frame"]
+    
+    # Capture current frame immediately
+    frame_bgr = state["frame"].copy()  # Copy to avoid race conditions
+    frame_rgb = state["frame_rgb"].copy() if state["frame_rgb"] is not None else cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+    
+    # Generate filename and save immediately
     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)
-
-    coords = []
-    if results.pose_landmarks:
-        for lm in results.pose_landmarks.landmark:
-            coords.append((round(lm.x, 5), round(lm.y, 5), round(lm.z, 5)))
-
-    global df
-    df = pd.concat([df, pd.DataFrame([{
-        "filename": filename,
-        "caption": caption,
-        "pose_coords": coords
-    }])], ignore_index=True)
+    cv2.imwrite(path, frame_bgr)
+    
+    # Process pose estimation in background (non-blocking)
+    executor.submit(process_pose_async, frame_bgr, filename, caption)
+    
+    return f"✅ Captured & paused: {filename} (processing pose...)", frame_rgb
 
-    return f"✅ Captured & paused: {filename}", frame[:, :, ::-1]
+# === Show dataset info ===
+def show_dataset_info():
+    return f"📊 Dataset contains {len(df)} samples"
 
 # === Download CSV ===
 def download_csv():
@@ -108,36 +150,57 @@ def download_csv():
 
 # === Reset all ===
 def reset_all():
-    df.drop(df.index, inplace=True)
-    for f in os.listdir(OUTPUT_DIR):
-        os.remove(os.path.join(OUTPUT_DIR, f))
+    global df
+    df = pd.DataFrame(columns=["filename", "caption", "pose_coords"])
+    
+    # Clean up files
+    try:
+        for f in os.listdir(OUTPUT_DIR):
+            file_path = os.path.join(OUTPUT_DIR, f)
+            if os.path.isfile(file_path):
+                os.remove(file_path)
+    except Exception as e:
+        print(f"Error cleaning files: {e}")
+    
+    # Reset video state
     if state["cap"]:
         state["cap"].release()
-    state.update({"video_path": None, "cap": None, "frame": None, "play": False})
+    state.update({
+        "video_path": None, 
+        "cap": None, 
+        "frame": None, 
+        "frame_rgb": None,
+        "play": False
+    })
     return "🔁 Reset done.", None
 
 # === UI ===
-with gr.Blocks() as app:
-    gr.Markdown("## 🏹 Archery Pose Dataset Tool (Auto-Pause on Capture)")
+with gr.Blocks(title="Fast Archery Pose Capture") as app:
+    gr.Markdown("## 🏹 Archery Pose Dataset Tool (Optimized for Speed)")
+    gr.Markdown("⚡ **Optimized**: Instant capture with background pose processing")
 
     video_input = gr.Video(label="🎞️ Upload Video")
-    load_btn = gr.Button("📂 Load Video")
-    status = gr.Textbox(label="Status")
+    load_btn = gr.Button("📂 Load Video", variant="primary")
+    status = gr.Textbox(label="Status", interactive=False)
 
     with gr.Row():
-        play_btn = gr.Button("▶️ Play")
-        pause_btn = gr.Button("⏸️ Pause")
-        show_btn = gr.Button("🖼️ Show Current Frame")
+        play_btn = gr.Button("▶️ Play", variant="secondary")
+        pause_btn = gr.Button("⏸️ Pause", variant="secondary")
+        show_btn = gr.Button("🖼️ Show Current Frame", variant="secondary")
 
-    image_output = gr.Image(label="Current Frame")
-    caption_input = gr.Textbox(label="Caption")
-    capture_btn = gr.Button("📸 Capture & Pause")
+    image_output = gr.Image(label="Current Frame", height=400)
+    
+    with gr.Row():
+        caption_input = gr.Textbox(label="Caption", placeholder="Describe the pose...")
+        capture_btn = gr.Button("📸 Capture & Pause", variant="primary", size="lg")
 
     with gr.Row():
+        info_btn = gr.Button("📊 Dataset Info")
         download_btn = gr.Button("📥 Download CSV")
-        reset_btn = gr.Button("🔄 Reset")
+        reset_btn = gr.Button("🔄 Reset", variant="stop")
     
     csv_file = gr.File(label="📄 Dataset CSV")
+    dataset_info = gr.Textbox(label="Dataset Info", interactive=False)
 
     # Bind actions
     load_btn.click(load_video, inputs=video_input, outputs=status)
@@ -145,7 +208,12 @@ with gr.Blocks() as app:
     pause_btn.click(pause_video, outputs=status)
     show_btn.click(show_frame, outputs=image_output)
     capture_btn.click(capture_frame, inputs=caption_input, outputs=[status, image_output])
+    info_btn.click(show_dataset_info, outputs=dataset_info)
     download_btn.click(download_csv, outputs=csv_file)
     reset_btn.click(reset_all, outputs=[status, image_output])
+    
+    # Auto-refresh frame display while playing
+    app.load(lambda: None)  # Initialize
 
-app.launch()
+if __name__ == "__main__":
+    app.launch(share=False, server_name="0.0.0.0", server_port=7860)