U ;4 hs#@sddlmZddlZddlZddlZddlmZddlmZm Z ddl Z ddl m Z mZddlmZddlmZdd d gZerddlad and \aad dZGdddZGdddZGdd d e eZGdd d eeZdS)) 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_cv2rrQ/root/miniconda3/envs/cl2/lib/python3.8/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 @speZdZeddejddfdddddd d d d d ZeddddZddZ d dddZ ddddZ dS)r NFrz str | intzCallable | Nonez int | NonestrboolintNone) video_source transformmax_num_frames color_ordermultiprocessing channel_dimreturncCsPts td|tkrt||_||_||_||_|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 RuntimeErrorrNotImplementedErrorr2r4r/r3 open_videocapr0r1)r!r/r0r1r2r3r4rrr__init__Bs  zVideoDataset.__init__)r/c CsXt|tr"tj|s"td|tt|}W5QRX| sTtd||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: zFailed to open video: ) isinstancer+rpathisfiler6rr VideoCaptureZisOpened)r/r9rrrr8ps  zVideoDataset.open_videocCs|jr||jS|jS)zaReturn the cap. If multiprocessing, create a new one. Else return the one from construction time.)r3r8r/r9r rrr_get_capszVideoDataset._get_capr5cCs|tjS)z"Get the FPS of the capture device.)r?getr Z CAP_PROP_FPSr rrrget_fpsszVideoDataset.get_fpsrcCs^|\}}|std|jtjkr6t|tj}t |d|j }|j dk rZ| |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) r?readr6r2rRGBr ZcvtColorZ COLOR_BGR2RGBnpZmoveaxisr4r0)r!retframerrr get_frames zVideoDataset.get_frame) r'r(r)rrrEr: staticmethodr8r?rBrIrrrrr >s.c@sVeZdZdZddddZeddddZd dd d Zd d Zd ddddZ dS)r zU Video dataset from file. This class requires that OpenCV be installed. r.r@cOs8tj|f|||}|jdks.||jkr4||_dSr#)r r:get_num_framesr1)r!argskwargs num_framesrrrr:szVideoFileDataset.__init__zdict[str, str]c CstsiSddddddd}i}tvtb}|D]R\}}t}tj |d|}tj |}| ||dd}|r|||<| q8W5QRXW5QRX|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) rQ) rrtempfileTemporaryDirectoryitemsr Z VideoWriterrr<joinZVideoWriter_fourccrrelease) Z all_codecscodecsZtmp_dircodecextwriterfnameZfourccZnoviderrrrrget_available_codecss  z%VideoFileDataset.get_available_codecsr-cCs(t|tj}|dkr$td|S)zz Return the number of frames in a video file. Raises: RuntimeError: no frames found. rz0 frames found)r-r?rAr ZCAP_PROP_FRAME_COUNTr6)r!rNrrrrKszVideoFileDataset.get_num_framescCs|jSr#)r1r rrr__len__szVideoFileDataset.__len__r)indexr5cCs2|jdk r||jkrt|tj||S)z4 Fetch single data item from index. N)r1 IndexErrorr?setr ZCAP_PROP_POS_FRAMESrI)r!r^rrr __getitem__szVideoFileDataset.__getitem__N) r'r(r)r*r:rJr\rKr]rarrrrr s c@s*eZdZdZeddddZddZdS) 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. r-r@cCs:tsdSd}t|}|ds$q6|d7}|q |S)zOGet number of possible devices detected by OpenCV that can be used for capture.rrP)rr r>rDrV)Z num_devicesr9rrrget_num_devicess   zCameraDataset.get_num_devicesccs6d}|}|d7}|V|jdk r||jkrq2qdS)NrrP)rIr1)r!Z frame_countrHrrr__iter__s  zCameraDataset.__iter__N)r'r(r)r*rJrbrcrrrrr s ) __future__rrrrRcollections.abcrtypingrrnumpyrFZtorch.utils.datarrZmonai.utils.enumsrZmonai.utils.moduler __all__r rrrr r r rrrr s&     \: