o e@sddlmZddlZddlmZddlmZddlmZm Z m Z m Z ddl Z ddl ZeGdddZGdd d Zd ddZ d!d"ddZd#ddZGdddZdS)$) annotationsN)deque) dataclass)Callable GeneratorOptionalTuplec@sReZdZUdZded<ded<ded<dZded<edd d ZedddZ dS) VideoInfoa A class to store video information, including width, height, fps and total number of frames. Attributes: width (int): width of the video in pixels height (int): height of the video in pixels fps (int): frames per second of the video total_frames (int, optional): total number of frames in the video, default is None Examples: ```python >>> import supervision as sv >>> video_info = sv.VideoInfo.from_video_path(video_path='video.mp4') >>> video_info VideoInfo(width=3840, height=2160, fps=25, total_frames=538) >>> video_info.resolution_wh (3840, 2160) ``` intwidthheightfpsN Optional[int] total_frames video_pathstrreturncCsvt|}|std|t|tj}t|tj}t|tj}t|tj }| t ||||S)NCould not open video at ) cv2 VideoCaptureisOpened Exceptionr getCAP_PROP_FRAME_WIDTHCAP_PROP_FRAME_HEIGHT CAP_PROP_FPSCAP_PROP_FRAME_COUNTreleaser )clsrvideor r r rr F/var/www/myenv/lib/python3.10/site-packages/supervision/utils/video.pyfrom_video_path,s zVideoInfo.from_video_pathTuple[int, int]cCs |j|jfSN)r r selfr r r! resolution_wh9s zVideoInfo.resolution_wh)rrrr )rr#) __name__ __module__ __qualname____doc____annotations__r classmethodr"propertyr'r r r r!r s    r c@s6eZdZdZdddd Zd d ZdddZddZdS) VideoSinka Context manager that saves video frames to a file using OpenCV. Attributes: target_path (str): The path to the output file where the video will be saved. video_info (VideoInfo): Information about the video resolution, fps, and total frame count. codec (str): FOURCC code for video format Example: ```python >>> import supervision as sv >>> video_info = sv.VideoInfo.from_video_path('source.mp4') >>> frames_generator = get_video_frames_generator('source.mp4') >>> with sv.VideoSink(target_path='target.mp4', video_info=video_info) as sink: ... for frame in frames_generator: ... sink.write_frame(frame=frame) ``` mp4v target_pathr video_infor codeccCs||_||_||_d|_dSr$)r1r2_VideoSink__codec_VideoSink__writer)r&r1r2r3r r r!__init__Us zVideoSink.__init__c Csvz tj|j|_Wn ty)}ztt|dtjd|_WYd}~nd}~wwt|j|j|j j |j j |_ |S)Nz. Defaulting to mp4v...r0) rVideoWriter_fourccr4_VideoSink__fourcc TypeErrorprintr VideoWriterr1r2r r'r5)r&er r r! __enter__[szVideoSink.__enter__frame np.ndarraycCs|j|dSr$)r5write)r&r>r r r! write_frameiszVideoSink.write_framecCs|jdSr$)r5r)r&exc_type exc_value exc_tracebackr r r!__exit__lszVideoSink.__exit__N)r0)r1rr2r r3r)r>r?)r(r)r*r+r6r=rArEr r r r!r/>s    r/ source_pathrstartr endrcCst|}|std|t|tj}|dur$||kr$tdt|d}|dur2t||n|}| tj ||||fS)NrzRequested frames are outboundr) rrrrr rrmaxminsetCAP_PROP_POS_FRAMES)rFrGrHrrr r r!_validate_and_setup_videops   rMstrider!Generator[np.ndarray, None, None]c csrt|||\}}}|} |\}}|r||krn|Vt|dD] }|}|s-nq#||7}q |dS)a= Get a generator that yields the frames of the video. Args: source_path (str): The path of the video file. stride (int): Indicates the interval at which frames are returned, skipping stride - 1 frames between each. start (int): Indicates the starting position from which video should generate frames end (Optional[int]): Indicates the ending position at which video should stop generating frames. If None, video will be read to the end. Returns: (Generator[np.ndarray, None, None]): A generator that yields the frames of the video. Examples: ```python >>> import supervision as sv >>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4'): ... ... ``` TrNN)rMreadrangegrabr) rFrOrGrHrframe_positionsuccessr>_r r r!get_video_frames_generator}s    rWr1callback'Callable[[np.ndarray, int], np.ndarray]NonecCsltj|d}t||d }tt|dD]\}}|||}|j|dqWddS1s/wYdS)aS Process a video file by applying a callback function on each frame and saving the result to a target video file. Args: source_path (str): The path to the source video file. target_path (str): The path to the target video file. callback (Callable[[np.ndarray, int], np.ndarray]): A function that takes in a numpy ndarray representation of a video frame and an int index of the frame and returns a processed numpy ndarray representation of the frame. Examples: ```python >>> import supervision as sv >>> def callback(scene: np.ndarray, index: int) -> np.ndarray: ... ... >>> process_video( ... source_path='...', ... target_path='...', ... callback=callback ... ) ``` )r)r1r2)rF)r>N)r r"r/ enumeraterWrA)rFr1rXsource_video_infosinkindexr> result_framer r r! process_videos  "r`c@s:eZdZdZddddZdd d Zdd d ZdddZdS) FPSMonitorzN A class for monitoring frames per second (FPS) to benchmark latency.  sample_sizer cCst|d|_dS)a Args: sample_size (int): The maximum number of observations for latency benchmarking. Examples: ```python >>> import supervision as sv >>> frames_generator = sv.get_video_frames_generator('source.mp4') >>> fps_monitor = sv.FPSMonitor() >>> for frame in frames_generator: ... # your processing code here ... fps_monitor.tick() ... fps = fps_monitor() ``` )maxlenN)rall_timestamps)r&rcr r r!r6szFPSMonitor.__init__rfloatcCs8|jsdS|jd|jd}|dkrt|j|SdS)z Computes and returns the average FPS based on the stored time stamps. Returns: float: The average FPS. Returns 0.0 if no time stamps are stored. gr)relen)r& taken_timer r r!__call__szFPSMonitor.__call__rZcCs|jtdS)zI Adds a new time stamp to the deque for FPS calculation. N)reappendtime monotonicr%r r r!tickszFPSMonitor.tickcCs|jdS)z< Clears all the time stamps from the deque. N)reclearr%r r r!resetszFPSMonitor.resetN)rb)rcr )rrf)rrZ)r(r)r*r+r6rjrnrpr r r r!ras    ra)rFrrGr rHr)rNrN) rFrrOr rGr rHrrrP)rFrr1rrXrYrrZ) __future__rrl collectionsr dataclassesrtypingrrrrrnumpynpr r/rMrWr`rar r r r!s   1 2 *(