Spaces:
Paused
Paused
| import asyncio | |
| import logging | |
| from datetime import datetime | |
| from jobs.models import JobStatus | |
| from jobs.storage import get_job_storage | |
| from jobs.streaming import create_stream, remove_stream | |
| from inference import run_inference, run_grounded_sam2_tracking | |
| async def process_video_async(job_id: str) -> None: | |
| storage = get_job_storage() | |
| job = storage.get(job_id) | |
| if not job: | |
| return | |
| detection_path = None | |
| depth_path = None | |
| depth_error = None | |
| partial_success = False | |
| # Create stream for live view | |
| stream_queue = create_stream(job_id) | |
| try: | |
| # Run detection or segmentation first | |
| if job.mode == "segmentation": | |
| detection_path = await asyncio.to_thread( | |
| run_grounded_sam2_tracking, | |
| input_video_path=job.input_video_path, | |
| output_video_path=job.output_video_path, | |
| queries=job.queries, | |
| max_frames=None, | |
| segmenter_name=job.segmenter_name, | |
| job_id=job_id, | |
| stream_queue=stream_queue, | |
| step=job.step, | |
| num_maskmem=7, | |
| detector_name=job.detector_name, | |
| _ttfs_t0=job.ttfs_t0, | |
| ) | |
| else: | |
| detections_list = None | |
| # Unified inference pipeline (handles depth internally if enabled) | |
| result_pkg = await asyncio.to_thread( | |
| run_inference, | |
| job.input_video_path, | |
| job.output_video_path, | |
| job.queries, | |
| None, | |
| job.detector_name, | |
| job_id, | |
| job.depth_estimator_name, | |
| job.depth_scale, | |
| stream_queue, | |
| job.first_frame_detections, | |
| ) | |
| detection_path, detections_list = result_pkg | |
| # If depth was requested, checking if output path exists for depth | |
| # The unified pipeline creates 'output_video_path'. | |
| # If depth enabled, it might have written depth there? | |
| # Actually run_inference returns (video_path, detections). | |
| # If depth was ON, the video at video_path *has* depth overlays. | |
| # But the 'Depth Video' (heatmap only) is usually separate. | |
| # Our Plan says: "Unified loop... Write Frame to Disk". | |
| # If we want separate depth video, we need `run_inference` to handle it. | |
| # Or just update 'depth_path' to be the same main video if it's merged? | |
| # Let's keep it simple: If depth enabled, the main video IS the depth view (overlay). | |
| # Or if we want separate `depth_output_path`, we need `run_inference` to handle it. | |
| # Let's assume for now `run_inference` writes the main visualization path. | |
| if job.depth_estimator_name: | |
| # In unified mode, the main video contains the depth viz | |
| depth_path = detection_path | |
| logging.info("Depth estimation included in main video for job %s", job_id) | |
| # Mark as completed (with or without depth) | |
| storage.update( | |
| job_id, | |
| status=JobStatus.COMPLETED, | |
| completed_at=datetime.utcnow(), | |
| output_video_path=detection_path, | |
| depth_output_path=depth_path, | |
| partial_success=partial_success, | |
| depth_error=depth_error, | |
| ) | |
| except RuntimeError as exc: | |
| # Handle cancellation specifically | |
| if "cancelled" in str(exc).lower(): | |
| logging.info("Job %s was cancelled", job_id) | |
| storage.update( | |
| job_id, | |
| status=JobStatus.CANCELLED, | |
| completed_at=datetime.utcnow(), | |
| error="Cancelled by user", | |
| ) | |
| else: | |
| logging.exception("Background processing failed for job %s", job_id) | |
| storage.update( | |
| job_id, | |
| status=JobStatus.FAILED, | |
| completed_at=datetime.utcnow(), | |
| error=str(exc), | |
| ) | |
| except Exception as exc: | |
| logging.exception("Background processing failed for job %s", job_id) | |
| storage.update( | |
| job_id, | |
| status=JobStatus.FAILED, | |
| completed_at=datetime.utcnow(), | |
| error=str(exc), | |
| ) | |
| finally: | |
| remove_stream(job_id) | |