Update backend/main.py

backend/main.py

@@ -21,8 +21,6 @@ import ffmpeg
 import torch
 import torchvision.transforms as T
 from ultralytics import YOLO
-import ultralytics.nn.tasks
-import ultralytics.nn.modules.conv
 import mediapipe as mp
 from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks, Form, Request
 from fastapi.responses import FileResponse, StreamingResponse, JSONResponse
@@ -41,9 +39,7 @@ torch.serialization.add_safe_globals([
     torch.nn.modules.activation.SiLU,
     torch.nn.modules.container.ModuleList,
     torch.nn.modules.upsampling.Upsample,
-    torch.nn.modules.pooling.MaxPool2d,
-    ultralytics.nn.tasks.DetectionModel,
-    ultralytics.nn.modules.conv.Conv
+    torch.nn.modules.pooling.MaxPool2d
 ])
 
 
@@ -404,109 +400,48 @@ 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(process_id: str, video_path: str, timestamp1: str, timestamp2: str, 
-              timestamp3: str, temp_dir: str, ffmpeg_path: str = 'ffmpeg') -> tuple[str, str]:
+def crop_video(video_path: str, timestamp1: str, timestamp2: str, timestamp3: str, temp_dir: str, ffmpeg_name: str = 'ffmpeg'):
     """
-    Crop the video into two clips with cancellation support
+    Crop the video into two clips:
+    - First clip: timestamp1 to timestamp2
+    - Second clip: timestamp2 to timestamp3
     """
-    temp_dir_path = Path(temp_dir)
+    ffmpeg_path = Path(__file__).parent / 'models' / ffmpeg_name
     
-    # Create temp directory if it doesn't exist
-    temp_dir_path.mkdir(parents=True, exist_ok=True)
-
-    # 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"
-
-    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")
-
-    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
-            )
-
-            # 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
-        )
+    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")
+
+    if not os.path.exists(ffmpeg_path):
+        raise FileNotFoundError(f"ffmpeg binary not found at {ffmpeg_path}")
+
+    # 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
     ]
 
-    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)
+    # 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
+    ]
 
-    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
+    logger.info("Running command: %s", ' '.join(command1))
+    subprocess.run(command1, check=True)
+    logger.info("Running command: %s", ' '.join(command2))
+    subprocess.run(command2, check=True)
+
+    return first_clip_path, second_clip_path
 
 #################################################
 # Video Processing Loop
@@ -531,8 +466,6 @@ 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
@@ -581,8 +514,6 @@ 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
@@ -737,7 +668,6 @@ 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,
@@ -910,6 +840,17 @@ async def serve_frontend(full_path: str):
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
-
+    # 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
+    )