Update backend/main.py

backend/main.py

@@ -400,48 +400,109 @@ def format_progress_message(stage, current, total, extras=None):
         return f"{base} - {', '.join(f'{k}: {v}' for k,v in extras.items())}"
     return base
 
-def crop_video(video_path: str, timestamp1: str, timestamp2: str, timestamp3: str, temp_dir: str, ffmpeg_name: str = 'ffmpeg'):
+def crop_video(process_id: str, video_path: str, timestamp1: str, timestamp2: str, 
+              timestamp3: str, temp_dir: str, ffmpeg_path: str = 'ffmpeg') -> tuple[str, str]:
     """
-    Crop the video into two clips:
-    - First clip: timestamp1 to timestamp2
-    - Second clip: timestamp2 to timestamp3
+    Crop the video into two clips with cancellation support
     """
-    ffmpeg_path = Path(__file__).parent / 'models' / ffmpeg_name
+    temp_dir_path = Path(temp_dir)
     
-    ts1 = time_to_seconds(timestamp1)
-    ts2 = time_to_seconds(timestamp2)
-    ts3 = time_to_seconds(timestamp3)
+    # Create temp directory if it doesn't exist
+    temp_dir_path.mkdir(parents=True, exist_ok=True)
 
-    first_clip_path = os.path.join(temp_dir, f"clip1_{uuid.uuid4()}.mp4")
-    second_clip_path = os.path.join(temp_dir, f"clip2_{uuid.uuid4()}.mp4")
+    # Generate temporary filenames
+    first_clip_path = temp_dir_path / f"clip1_{uuid.uuid4()}.mp4"
+    second_clip_path = temp_dir_path / f"clip2_{uuid.uuid4()}.mp4"
 
-    if not os.path.exists(ffmpeg_path):
-        raise FileNotFoundError(f"ffmpeg binary not found at {ffmpeg_path}")
+    def check_cancellation():
+        """Check if processing was cancelled (replace with your actual progress store)"""
+        # You'll need to import or access your PROGRESS_STORE here
+        if PROGRESS_STORE.get(process_id, {}).get('status') == 'cancelled':
+            raise asyncio.CancelledError("Processing cancelled by user during video cropping")
 
-    # Clip 1
-    dur1 = ts2 - ts1
-    command1 = [
-        ffmpeg_path, '-y', '-i', video_path,
-        '-ss', str(ts1), '-t', str(dur1),
-        '-c:v', 'libx264', '-preset', 'fast', '-crf', '23',
-        '-c:a', 'aac', first_clip_path
-    ]
+    def run_ffmpeg_with_cancel_check(command: list, output_file: Path) -> None:
+        """Run ffmpeg command with cancellation checks"""
+        try:
+            # Start the process
+            process = subprocess.Popen(
+                command,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                universal_newlines=True
+            )
 
-    # Clip 2
-    dur2 = ts3 - ts2
-    command2 = [
-        ffmpeg_path, '-y', '-i', video_path,
-        '-ss', str(ts2), '-t', str(dur2),
-        '-c:v', 'libx264', '-preset', 'fast', '-crf', '23',
-        '-c:a', 'aac', second_clip_path
+            # Poll process while checking for cancellation
+            while True:
+                check_cancellation()
+                if process.poll() is not None:  # Process finished
+                    break
+                time.sleep(0.5)  # Check every 500ms
+
+            # Check final status
+            if process.returncode != 0:
+                raise subprocess.CalledProcessError(
+                    process.returncode,
+                    command,
+                    output=process.stdout,
+                    stderr=process.stderr
+                )
+
+        except asyncio.CancelledError:
+            # Cleanup and terminate process
+            if process.poll() is None:  # Still running
+                process.terminate()
+                try:
+                    process.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    process.kill()
+            
+            # Remove partial output file
+            if output_file.exists():
+                output_file.unlink()
+            
+            raise
+
+    # Convert timestamps
+    ts1 = time_to_seconds(timestamp1)
+    ts2 = time_to_seconds(timestamp2)
+    ts3 = time_to_seconds(timestamp3)
+
+    # Build commands
+    commands = [
+        (
+            [
+                ffmpeg_path, '-y', '-i', video_path,
+                '-ss', str(ts1), '-t', str(ts2 - ts1),
+                '-c:v', 'libx264', '-preset', 'fast', '-crf', '23',
+                '-c:a', 'aac', str(first_clip_path)
+            ],
+            first_clip_path
+        ),
+        (
+            [
+                ffmpeg_path, '-y', '-i', video_path,
+                '-ss', str(ts2), '-t', str(ts3 - ts2),
+                '-c:v', 'libx264', '-preset', 'fast', '-crf', '23',
+                '-c:a', 'aac', str(second_clip_path)
+            ],
+            second_clip_path
+        )
     ]
 
-    logger.info("Running command: %s", ' '.join(command1))
-    subprocess.run(command1, check=True)
-    logger.info("Running command: %s", ' '.join(command2))
-    subprocess.run(command2, check=True)
+    try:
+        # Process both clips
+        for cmd, output_path in commands:
+            logger.info("Running command: %s", ' '.join(cmd))
+            run_ffmpeg_with_cancel_check(cmd, output_path)
+        
+        return str(first_clip_path), str(second_clip_path)
 
-    return first_clip_path, second_clip_path
+    except asyncio.CancelledError:
+        # Cleanup both files if either was cancelled
+        for path in [first_clip_path, second_clip_path]:
+            if path.exists():
+                path.unlink()
+        raise
 
 #################################################
 # Video Processing Loop
@@ -466,6 +527,8 @@ def process_freeplay(process_id: str, freeplay_video: str) -> float:
     prev_pose = None
 
     for sec in range(int(duration)):
+        if PROGRESS_STORE.get(process_id, {}).get('status') == 'cancelled':
+            raise asyncio.CancelledError("Processing cancelled")
         print(f"Processing freeplay frame {sec}")
         if PROGRESS_STORE[process_id]["status"] == "cancelled":
             break
@@ -514,6 +577,8 @@ def process_experiment(process_id: str, experiment_video: str, freeplay_movement
     prev_pose = None
 
     for sec in range(int(duration)):
+        if PROGRESS_STORE.get(process_id, {}).get('status') == 'cancelled':
+            raise asyncio.CancelledError("Processing cancelled")
         print(f"Processing experiment frame {sec}")
         if PROGRESS_STORE[process_id]["status"] == "cancelled":
             break
@@ -668,6 +733,7 @@ async def process_video_async(process_id: str, video_path: Path, session_dir: Pa
     try:
         freeplay_video, experiment_video = await asyncio.to_thread(
             crop_video,
+            process_id,
             str(video_path),
             timestamp1,
             timestamp2,
@@ -840,17 +906,6 @@ async def serve_frontend(full_path: str):
 
 if __name__ == "__main__":
     import uvicorn
-    # Pre-create cache dir
-    Path(os.environ['TORCH_HOME']).mkdir(exist_ok=True)
-    
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=7860,
-        workers=1,
-        log_level="info",
-        # Add these parameters
-        limit_concurrency=1,  # Reduce memory pressure
-        timeout_graceful_shutdown=30
-    )
+    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
+