Refactor video frame extraction to improve PyNvCodec availability check

- Removed the try-except block for importing PyNvcFrameExtractor, simplifying the import logic.
- Updated the condition for initializing the PyNvcFrameExtractor in the VideoFrameExtractionStage to rely solely on the _PYNVC_AVAILABLE flag.
- Adjusted the handling of pixel format conversion in NvVideoDecoder to prepare for future updates to cvcuda.

Signed-off-by: Abhinav Garg <abhinavg@stanford.edu>

Author: abhinavg4

nemo_curator/stages/video/clipping/video_frame_extraction.py

@@ -24,17 +24,9 @@
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.stages.resources import Resources
 from nemo_curator.tasks.video import VideoTask
+from nemo_curator.utils.nvcodec_utils import _PYNVC_AVAILABLE, PyNvcFrameExtractor
 from nemo_curator.utils.operation_utils import make_pipeline_named_temporary_file
 
-try:
-    from nemo_curator.utils.nvcodec_utils import PyNvcFrameExtractor
-
-    _PYNVC_AVAILABLE = True
-except (ImportError, RuntimeError):
-    logger.warning("PyNvcFrameExtractor not available, PyNvCodec mode will fall back to FFmpeg")
-    PyNvcFrameExtractor = None
-    _PYNVC_AVAILABLE = False
-
 
 def get_frames_from_ffmpeg(
     video_file: Path,
@@ -123,7 +115,7 @@ def setup(self, worker_metadata: WorkerMetadata | None = None) -> None:  # noqa:
             worker_metadata (WorkerMetadata, optional): Information about the worker (provided by some backends)
         """
         if self.decoder_mode == "pynvc":
-            if _PYNVC_AVAILABLE and PyNvcFrameExtractor is not None:
+            if _PYNVC_AVAILABLE:
                 self.pynvc_frame_extractor = PyNvcFrameExtractor(
                     width=self.output_hw[1],
                     height=self.output_hw[0],

nemo_curator/utils/nvcodec_utils.py

@@ -22,21 +22,21 @@
 
 try:
     import cvcuda  # type: ignore[import-untyped]
-    import nvcv  # type: ignore[import-untyped]
     import pycuda.driver as cuda  # type: ignore[import-untyped]
     import PyNvVideoCodec as Nvc  # type: ignore[import-untyped]
 
     pixel_format_to_cvcuda_code = {
         Nvc.Pixel_Format.YUV444: cvcuda.ColorConversion.YUV2RGB,  # type: ignore[import-untyped]
         Nvc.Pixel_Format.NV12: cvcuda.ColorConversion.YUV2RGB_NV12,  # type: ignore[import-untyped]
     }
+    _PYNVC_AVAILABLE = True
 except (ImportError, RuntimeError):
     logger.warning("PyNvVideoCodec is not installed, some features will be disabled.")
     Nvc = None
     cvcuda = None
-    nvcv = None
     cuda = None
     pixel_format_to_cvcuda_code = {}
+    _PYNVC_AVAILABLE = False
 
 
 class FrameExtractionPolicy(enum.Enum):
@@ -268,7 +268,8 @@ def generate_decoded_frames(self) -> list[torch.Tensor]:
         for packet in self.nvDemux:
             list_frames = self.nvDec.Decode(packet)
             for decoded_frame in list_frames:
-                nvcv_tensor = nvcv.as_tensor(nvcv.as_image(decoded_frame.nvcv_image(), nvcv.Format.U8))
+                # TODO: Replace nvcv with cvcuda. Before that remove the use of nvcv_image. It's deprecated
+                nvcv_tensor = cvcuda.as_tensor(cvcuda.as_image(decoded_frame.nvcv_image(), cvcuda.Format.U8))
                 if nvcv_tensor.layout == "NCHW":
                     nchw_shape = nvcv_tensor.shape
                     nhwc_shape = (nchw_shape[0], nchw_shape[2], nchw_shape[3], nchw_shape[1])