o Fi#@sddlmZddlZddlZddlZddlmZddlmZm Z ddl Z ddl m Z mZddlmZddlmZgdZerEddlad and \aad d ZGd ddZGdddZGddde eZGdddeeZdS)) annotationsN)Callable) TYPE_CHECKINGAny)DatasetIterableDataset) ColorOrder)optional_import) VideoDatasetVideoFileDataset CameraDatasetT)NNcCstd\aadS)zYImport cv2. Put it inside a function to avoid webcam lights blinking on ``import monai``.cv2N)r r has_cv2rrR/root/miniconda3/envs/fpt/lib/python3.10/site-packages/monai/data/video_dataset.py import_cv$src@s eZdZdZddZddZdS)SuppressStderrzQSuppress stderr. Useful as OpenCV (and dependencies) can produce a lot of output.cCsTttjd|_tj|_ttj|_ tj|_ t |j|j|jt_|S)Nw) openosdevnull errnull_filesysstderrfilenoold_stderr_fileno_undupdupold_stderr_fileno old_stderrdup2selfrrr __enter__.s zSuppressStderr.__enter__cGs2|jt_t|j|jt|j|jdSN) rrrrrrrcloser)r!_rrr__exit__7s zSuppressStderr.__exit__N)__name__ __module__ __qualname____doc__r"r&rrrrr+s rc@sTeZdZeddejddfdddZedddZddZ dddZ d ddZ dS)!r NFr video_source str | int transformCallable | Nonemax_num_frames int | None color_orderstrmultiprocessingbool channel_dimintreturnNonecCsPtstd|tvr t||_||_||_||_|s |||_ ||_ ||_ dS)a Base video dataset. Args: video_source: filename of video. transform: transform to be applied to each frame. max_num_frames: Max number of frames to iterate across. If `None` is passed, then the dataset will iterate until the end of the file. color_order: Color order to return frame. Default is RGB. multiprocessing: If `True`, open the video source on the fly. This makes things process-safe, which is useful when combined with a DataLoader with `num_workers>0`. However, when using with `num_workers==0`, it makes sense to use `multiprocessing=False`, as the source will then only be opened once, at construction, which will be faster in those circumstances. channel_dim: OpenCV reads with the channel as the last dimension. Use this flag to move it elsewhere. By default this is zero, so the channel dimension is moved to the front. Raises: RuntimeError: OpenCV not installed. NotImplementedError: Unknown color order. zOpenCV not installed.N) r RuntimeErrorrNotImplementedErrorr1r5r+r3 open_videocapr-r/)r!r+r-r/r1r3r5rrr__init__Bs   zVideoDataset.__init__cCslt|trtj|std|t t|}Wdn1s$wY| s4td||S)a0 Use OpenCV to open a video source from either file or capture device. Args: video_source: filename or index referring to capture device. Raises: RuntimeError: Source is a file but file not found. RuntimeError: Failed to open source. zVideo file does not exist: NzFailed to open video: ) isinstancer2rpathisfiler9rr VideoCaptureZisOpened)r+r<rrrr;ps  zVideoDataset.open_videocCs|jr ||jS|jS)zaReturn the cap. If multiprocessing, create a new one. Else return the one from construction time.)r3r;r+r<r rrr_get_capszVideoDataset._get_capcCs|tjS)z"Get the FPS of the capture device.)rBgetr Z CAP_PROP_FPSr rrrget_fpsszVideoDataset.get_fpsrcCs^|\}}|std|jtjkrt|tj}t |d|j }|j dur-| |S|S)zGet next frame. For a file, this will be the next frame, whereas for a camera source, it will be the next available frame.zFailed to read frame.N) rBreadr9r1rRGBr ZcvtColorZ COLOR_BGR2RGBnpZmoveaxisr5r-)r!retframerrr get_frames zVideoDataset.get_frame)r+r,r-r.r/r0r1r2r3r4r5r6r7r8)r+r,r7r6)r7r) r'r(r)rrrGr= staticmethodr;rBrDrKrrrrr >s .  r c@sDeZdZdZdddZedddZdd d Zd d ZdddZ dS)r zU Video dataset from file. This class requires that OpenCV be installed. r7r8cOsBtj|g|Ri||}|jdus||jkr||_dSdSr#)r r=get_num_framesr/)r!argskwargs num_framesrrrr=s  zVideoFileDataset.__init__dict[str, str]c CstsiSddddddd}i}tSt6}|D])\}}t}tj |d|}tj |}| ||dd}|rA|||<| qWdn1sPwYWd|SWd|S1shwY|S) zTry different codecs, see which are available. Returns a dictionary with of available codecs with codecs as keys and file extensions as values.z.mp4z.aviz.mjpeg)Zmp4vZX264ZH264ZMP42ZMJPGZDIVXtest) rUN) rrtempfileTemporaryDirectoryitemsr Z VideoWriterrr?joinZVideoWriter_fourccrrelease) Z all_codecscodecsZtmp_dircodecextwriterfnameZfourccZnoviderrrrrget_available_codecss.        z%VideoFileDataset.get_available_codecsr6cCs(t|tj}|dkrtd|S)zz Return the number of frames in a video file. Raises: RuntimeError: no frames found. rz0 frames found)r6rBrCr ZCAP_PROP_FRAME_COUNTr9)r!rQrrrrNszVideoFileDataset.get_num_framescCs|jSr#)r/r rrr__len__szVideoFileDataset.__len__indexrcCs2|jdur ||jkr t|tj||S)z4 Fetch single data item from index. N)r/ IndexErrorrBsetr ZCAP_PROP_POS_FRAMESrK)r!rbrrr __getitem__szVideoFileDataset.__getitem__N)r7r8)r7rRrL)rbr6r7r) r'r(r)r*r=rMr`rNrarerrrrr s    r c@s&eZdZdZed ddZddZdS) r a Video dataset from a capture device (e.g., webcam). This class requires that OpenCV be installed. Args: video_source: index of capture device. `get_num_devices` can be used to determine possible devices. transform: transform to be applied to each frame. max_num_frames: Max number of frames to iterate across. If `None` is passed, then the dataset will iterate infinitely. Raises: RuntimeError: OpenCV not installed. r7r6cCs<tsdSd} t|}|ds |S|d7}|q)zOGet number of possible devices detected by OpenCV that can be used for capture.rTrT)rr rArFrZ)Z num_devicesr<rrrget_num_devicess  zCameraDataset.get_num_devicesccs8d} |}|d7}|V|jdur||jkrdSq)NrTrT)rKr/)r!Z frame_countrJrrr__iter__s  zCameraDataset.__iter__NrL)r'r(r)r*rMrfrgrrrrr s   r ) __future__rrrrVcollections.abcrtypingrrnumpyrHZtorch.utils.datarrZmonai.utils.enumsrZmonai.utils.moduler __all__r rrrr r r rrrrs(    \: