Add support for NVDEC backend in video processing and enhance README documentation. Introduce a new script for extracting EasyCom sessions on specific GPUs, detailing session management and error handling. Update requirements.txt to include PyNvVideoCodec as an optional dependency. Refactor streaming_turn_model_directly.py to accommodate multiple decoder backends and improve frame indexing with ffprobe.

2026-08-18 08:41:54 +08:00
parent e4fb62ba60
commit 24cb92f1f4
5 changed files with 304 additions and 22 deletions
--- a/README.md
+++ b/README.md
@@ -651,6 +651,35 @@ python streaming_turn_model_directly.py \
  --output features/direct/Session_1.pt
 ```

+The direct reader defaults to Decord. On NVIDIA systems, install
+`PyNvVideoCodec` and select the NVDEC backend explicitly:
+
+```bash
+python -m pip install PyNvVideoCodec
+
+CUDA_VISIBLE_DEVICES=0 python streaming_turn_model_directly.py \
+  --model-id /path/to/Qwen2.5-Omni-3B \
+  --easycom-session-dir /path/to/Main/Video_Compressed/Session_1 \
+  --end 24 \
+  --mode open-stream-replay \
+  --memory-seconds 20 \
+  --video-fps 2 \
+  --frame-cache-size 64 \
+  --decoder pynvcodec \
+  --decoder-gpu-id 0 \
+  --flash-attention-2 \
+  --device cuda:0 \
+  --output features/direct/Session_1_24s_pynvcodec.pt
+```
+
+The NVDEC backend uses ffprobe for exact display-order PTS indexing, one
+reconfigurable `SimpleDecoder` for pixel decode, and the existing ffmpeg PCM
+audio path. `--decoder-gpu-id` is relative to `CUDA_VISIBLE_DEVICES`, so it is
+normally `0` even when the worker is assigned a nonzero host GPU. Decoder RGB
+conversion can differ slightly; validate a short Decord/PyNvVideoCodec pair
+before extracting every session, and do not mix decoder backends across
+train/test splits.
+
 Matching `states.shape[0]` to the target is necessary but not sufficient.
 The direct reader also validates EasyCom's 20-Hz source inventory, source PTS
 monotonicity, dense duration, crop bounds, audio sample count, and causal
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,3 +9,4 @@ tqdm
 scikit-learn

 # Optional on compatible NVIDIA GPUs: flash-attn
+# Optional NVDEC backend: PyNvVideoCodec
--- a/run_extract_easycom_2x4090_decord.sh
+++ b/run_extract_easycom_2x4090_decord.sh
@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+# Extract all 12 EasyCom sessions on host GPUs 3 and 4.
+# Exactly one Qwen/Decord worker runs on each GPU; sessions are processed
+# sequentially within each worker. The ordering balances the official 5-Hz
+# target lengths across the two GPUs.
+set -euo pipefail
+
+PROJECT_ROOT="${PROJECT_ROOT:-/data/wangx/Code/icassp2027-experiment}"
+DATASET_ROOT="${DATASET_ROOT:-/data/wangx/Datasets/easycom}"
+MODEL="${MODEL:-/data/wangx/Models/Qwen2.5-Omni-3B}"
+
+VIDEO_ROOT="${VIDEO_ROOT:-${DATASET_ROOT}/extracted_easycom/Main/Video_Compressed}"
+TARGET_ROOT="${TARGET_ROOT:-${DATASET_ROOT}/easycom_target_frame}"
+FEATURE_ROOT="${FEATURE_ROOT:-${DATASET_ROOT}/omni_features/open_stream_2fps_20s_decord}"
+LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_2x4090_decord}"
+
+mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}"
+
+for gpu in 3 4; do
+  nvidia-smi -i "${gpu}" --query-gpu=index,name,memory.total \
+    --format=csv,noheader
+done
+
+for sid in {1..12}; do
+  test -d "${VIDEO_ROOT}/Session_${sid}" || {
+    echo "Missing video directory: ${VIDEO_ROOT}/Session_${sid}" >&2
+    exit 1
+  }
+  test -s "${TARGET_ROOT}/Session_${sid}.npy" || {
+    echo "Missing target: ${TARGET_ROOT}/Session_${sid}.npy" >&2
+    exit 1
+  }
+done
+
+echo "Features: ${FEATURE_ROOT}"
+echo "Logs:     ${LOG_ROOT}"
+echo "GPU 3: Session_12 Session_5 Session_9 Session_8 Session_10 Session_1"
+echo "GPU 4: Session_6 Session_4 Session_3 Session_2 Session_7 Session_11"
+
+exec nice -n 5 env \
+  OMP_NUM_THREADS=4 \
+  MKL_NUM_THREADS=4 \
+  OPENBLAS_NUM_THREADS=4 \
+  PROJECT_ROOT="${PROJECT_ROOT}" \
+  MODEL="${MODEL}" \
+  VIDEO_ROOT="${VIDEO_ROOT}" \
+  TARGET_ROOT="${TARGET_ROOT}" \
+  FEATURE_ROOT="${FEATURE_ROOT}" \
+  LOG_ROOT="${LOG_ROOT}" \
+  GPUS="3 4" \
+  SESSIONS="12 6 5 4 9 3 8 2 10 7 1 11" \
+  MEMORY_SECONDS=20 \
+  NATIVE_BLOCK_SECONDS=2 \
+  VIDEO_FPS=2 \
+  FRAME_CACHE_SIZE=64 \
+  MODE=open-stream-replay \
+  DTYPE=bfloat16 \
+  DECODER=decord \
+  DECODER_GPU_ID=0 \
+  bash "${PROJECT_ROOT}/run_extract_easycom_direct.sh"
--- a/run_extract_easycom_direct.sh
+++ b/run_extract_easycom_direct.sh
@@ -18,6 +18,9 @@ NATIVE_BLOCK_SECONDS="${NATIVE_BLOCK_SECONDS:-2}"
 VIDEO_FPS="${VIDEO_FPS:-2}"
 MODE="${MODE:-open-stream-replay}"
 DTYPE="${DTYPE:-bfloat16}"
+DECODER="${DECODER:-decord}"
+DECODER_GPU_ID="${DECODER_GPU_ID:-0}"
+FRAME_CACHE_SIZE="${FRAME_CACHE_SIZE:-64}"
 # Space-separated host GPU ids, e.g. "0" or "0 1 2 3". One session per GPU at a time.
 GPUS="${GPUS:-0}"
 # Optional: only run selected session ids, e.g. SESSIONS="1 2 3"
@@ -69,6 +72,7 @@ run_one() {
    echo "target=${target}"
    echo "output=${output}"
    echo "memory_seconds=${MEMORY_SECONDS} video_fps=${VIDEO_FPS} mode=${MODE}"
+    echo "decoder=${DECODER} decoder_gpu_id=${DECODER_GPU_ID} frame_cache_size=${FRAME_CACHE_SIZE}"

    if ! CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \
      python streaming_turn_model_directly.py \
@@ -80,6 +84,9 @@ run_one() {
        --memory-seconds "${MEMORY_SECONDS}" \
        --native-block-seconds "${NATIVE_BLOCK_SECONDS}" \
        --video-fps "${VIDEO_FPS}" \
+        --frame-cache-size "${FRAME_CACHE_SIZE}" \
+        --decoder "${DECODER}" \
+        --decoder-gpu-id "${DECODER_GPU_ID}" \
        --dtype "${DTYPE}" \
        --flash-attention-2 \
        --device cuda:0 \
--- a/streaming_turn_model_directly.py
+++ b/streaming_turn_model_directly.py
@@ -95,7 +95,7 @@ class IndexedSourceClip:
    dense_end: float
    duration: float
    crop: tuple[int, int, int, int] | None
-    reader: Any
+    reader: Any | None
    source_frame_starts: np.ndarray
    source_frame_ends: np.ndarray
    source_fps: float
@@ -151,6 +151,76 @@ def _timestamp_array(reader: Any) -> tuple[np.ndarray, np.ndarray]:
    return starts, ends


+def _parse_ffprobe_rate(value: str, path: Path) -> float:
+    """Parse an ffprobe rational such as ``20/1``."""
+
+    try:
+        numerator, denominator = value.split("/", maxsplit=1)
+        rate = float(numerator) / float(denominator)
+    except (AttributeError, ValueError, ZeroDivisionError) as exc:
+        raise RuntimeError(f"Invalid ffprobe frame rate {value!r}: {path}") from exc
+    if not math.isfinite(rate) or rate <= 0:
+        raise RuntimeError(f"Non-positive ffprobe frame rate {rate}: {path}")
+    return rate
+
+
+def _ffprobe_video_index(
+    path: Path,
+) -> tuple[np.ndarray, np.ndarray, float, int, int]:
+    """Index display-order frame PTS without keeping a software decoder open."""
+
+    if shutil.which("ffprobe") is None:
+        raise RuntimeError("ffprobe is required by the pynvcodec backend")
+    completed = subprocess.run(
+        [
+            "ffprobe",
+            "-v",
+            "error",
+            "-select_streams",
+            "v:0",
+            "-show_frames",
+            "-show_entries",
+            (
+                "stream=width,height,avg_frame_rate:"
+                "frame=best_effort_timestamp_time,pkt_duration_time"
+            ),
+            "-of",
+            "json",
+            str(path),
+        ],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    if completed.returncode != 0:
+        raise RuntimeError(
+            f"Failed to index video with ffprobe: {path}: "
+            f"{completed.stderr.strip()}"
+        )
+    try:
+        payload = json.loads(completed.stdout)
+        stream = payload["streams"][0]
+        width = int(stream["width"])
+        height = int(stream["height"])
+        source_fps = _parse_ffprobe_rate(stream["avg_frame_rate"], path)
+        raw_starts = [
+            float(frame["best_effort_timestamp_time"])
+            for frame in payload["frames"]
+        ]
+    except (IndexError, KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
+        raise RuntimeError(f"Incomplete ffprobe frame index for {path}") from exc
+
+    if not raw_starts:
+        raise RuntimeError(f"Video contains no indexed frames: {path}")
+    starts = np.asarray(raw_starts, dtype=np.float64)
+    starts -= float(starts[0])
+    if np.any(np.diff(starts) < -1e-9):
+        raise RuntimeError(f"ffprobe video frame PTS values are not monotonic: {path}")
+    nominal = 1.0 / source_fps
+    ends = np.concatenate((starts[1:], starts[-1:] + nominal))
+    return starts, ends, source_fps, width, height
+
+
 def _decode_audio_exact(path: Path, duration: float) -> np.ndarray:
    """Decode, then pad/trim one clip to its dense video duration."""

@@ -203,6 +273,8 @@ class CausalSourceAV:
        output_fps: float,
        required_duration: float,
        frame_cache_size: int = 8,
+        decoder_backend: str = "decord",
+        decoder_gpu_id: int = 0,
    ) -> None:
        if not specs:
            raise ValueError("At least one source clip is required")
@@ -212,20 +284,40 @@ class CausalSourceAV:
            raise ValueError("required_duration must be positive")
        if frame_cache_size < 0:
            raise ValueError("frame_cache_size must be non-negative")
+        if decoder_backend not in {"decord", "pynvcodec"}:
+            raise ValueError(f"Unsupported decoder backend: {decoder_backend}")
+        if decoder_gpu_id < 0:
+            raise ValueError("decoder_gpu_id must be non-negative")
        if shutil.which("ffmpeg") is None:
            raise RuntimeError("ffmpeg is required to decode source audio")

+        decord = None
+        self._nvc = None
+        if decoder_backend == "decord":
            try:
                import decord
            except ImportError as exc:
                raise RuntimeError(
-                "Direct extraction requires decord; install "
-                "qwen-omni-utils[decord]"
+                    "The decord backend requires qwen-omni-utils[decord]"
                ) from exc
+        else:
+            try:
+                import PyNvVideoCodec as nvc
+            except ImportError as exc:
+                raise RuntimeError(
+                    "The pynvcodec backend requires PyNvVideoCodec; install it "
+                    "with `python -m pip install PyNvVideoCodec`"
+                ) from exc
+            self._nvc = nvc

        self.fps = float(output_fps)
        self.required_duration = float(required_duration)
        self.frame_cache_size = int(frame_cache_size)
+        self.decoder_backend = decoder_backend
+        self.decoder_gpu_id = int(decoder_gpu_id)
+        self._nvc_decoder = None
+        self._nvc_source_path: Path | None = None
+        self._nvc_decoder_reconfigurations = 0
        self._frame_cache: OrderedDict[tuple[int, int], Any] = OrderedDict()
        self._request_count = 0
        self._cache_hits = 0
@@ -246,6 +338,8 @@ class CausalSourceAV:
                f"{path.name}",
                flush=True,
            )
+            if decoder_backend == "decord":
+                assert decord is not None
                reader = decord.VideoReader(str(path))
                if not len(reader):
                    raise RuntimeError(f"No video frames decoded from {path}")
@@ -257,8 +351,11 @@ class CausalSourceAV:
                        f"Unexpected decoded frame shape {first_shape}: {path}"
                    )
                height, width = first_shape[:2]
+            else:
+                reader = None
+                starts, ends, source_fps, width, height = _ffprobe_video_index(path)

-            inferred_duration = len(reader) / source_fps
+            inferred_duration = len(starts) / source_fps
            duration = (
                float(spec.duration)
                if spec.duration is not None
@@ -389,7 +486,25 @@ class CausalSourceAV:
            return cached.copy()

        clip = self.clips[clip_index]
+        if self.decoder_backend == "decord":
+            if clip.reader is None:
+                raise RuntimeError(f"Missing decord reader for {clip.path}")
            array = clip.reader[frame_index].asnumpy()
+        else:
+            decoder = self._pynv_decoder_for_clip(clip)
+            decoded_frame = decoder[frame_index]
+            frame_tensor = torch.from_dlpack(decoded_frame)
+            if frame_tensor.ndim != 3 or frame_tensor.shape[-1] != 3:
+                raise RuntimeError(
+                    "PyNvVideoCodec RGB output must have HWC shape, got "
+                    f"{tuple(frame_tensor.shape)} for {clip.path}"
+                )
+            if frame_tensor.dtype != torch.uint8:
+                raise RuntimeError(
+                    "PyNvVideoCodec RGB output must be uint8, got "
+                    f"{frame_tensor.dtype} for {clip.path}"
+                )
+            array = frame_tensor.cpu().contiguous().numpy()
        if clip.crop is not None:
            x, y, width, height = clip.crop
            array = array[y : y + height, x : x + width]
@@ -401,6 +516,44 @@ class CausalSourceAV:
                self._frame_cache.popitem(last=False)
        return image.copy()

+    def _pynv_decoder_for_clip(self, clip: IndexedSourceClip) -> Any:
+        """Return one reusable NVDEC decoder configured for ``clip``."""
+
+        if self._nvc is None:
+            raise RuntimeError("PyNvVideoCodec was not initialized")
+        source_path = clip.path
+        if self._nvc_decoder is None:
+            self._nvc_decoder = self._new_pynv_decoder(source_path)
+            self._nvc_source_path = source_path
+        elif self._nvc_source_path != source_path:
+            torch.cuda.synchronize(self.decoder_gpu_id)
+            try:
+                self._nvc_decoder.reconfigure_decoder(str(source_path))
+            except Exception:
+                # A rebuild is safer than silently using stale decoder state.
+                self._nvc_decoder = self._new_pynv_decoder(source_path)
+            self._nvc_source_path = source_path
+            self._nvc_decoder_reconfigurations += 1
+
+        decoded_length = len(self._nvc_decoder)
+        if decoded_length != clip.num_frames:
+            raise RuntimeError(
+                "PyNvVideoCodec/ffprobe frame-count mismatch for "
+                f"{clip.path}: decoder={decoded_length}, ffprobe={clip.num_frames}"
+            )
+        return self._nvc_decoder
+
+    def _new_pynv_decoder(self, source_path: Path) -> Any:
+        assert self._nvc is not None
+        return self._nvc.SimpleDecoder(
+            str(source_path),
+            gpu_id=self.decoder_gpu_id,
+            use_device_memory=True,
+            need_scanned_stream_metadata=True,
+            decoder_cache_size=4,
+            output_color_type=self._nvc.OutputColorType.RGB,
+        )
+
    def frames(self, start: float, end: float) -> list[Any]:
        """Return causal global-grid frames in ``[start, end)``."""

@@ -451,7 +604,20 @@ class CausalSourceAV:

    def metadata(self) -> dict[str, Any]:
        return {
-            "decoder": "decord_with_frame_pts_and_ffmpeg_pcm_audio",
+            "decoder": (
+                "decord_with_frame_pts_and_ffmpeg_pcm_audio"
+                if self.decoder_backend == "decord"
+                else "pynvvideocodec_nvdec_with_ffprobe_pts_and_ffmpeg_pcm_audio"
+            ),
+            "decoder_backend": self.decoder_backend,
+            "decoder_gpu_id": (
+                self.decoder_gpu_id if self.decoder_backend == "pynvcodec" else None
+            ),
+            "decoder_version": (
+                getattr(self._nvc, "__version__", None)
+                if self.decoder_backend == "pynvcodec"
+                else None
+            ),
            "visual_sampling": "global_grid_latest_source_pts_not_after_tick",
            "interval_convention": "half_open_[start,end)",
            "output_video_fps": self.fps,
@@ -483,6 +649,7 @@ class CausalSourceAV:
                "odd_frame_duplications": self._odd_frame_duplications,
                "short_interval_fallbacks": self._short_interval_fallbacks,
                "tail_frame_reuses": self._tail_frame_reuses,
+                "decoder_reconfigurations": self._nvc_decoder_reconfigurations,
            },
        }

@@ -609,6 +776,18 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
    parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
    parser.add_argument("--frame-cache-size", type=int, default=8)
+    parser.add_argument(
+        "--decoder",
+        choices=("decord", "pynvcodec"),
+        default="decord",
+        help="Video pixel decoder; pynvcodec uses NVIDIA NVDEC",
+    )
+    parser.add_argument(
+        "--decoder-gpu-id",
+        type=int,
+        default=0,
+        help="GPU visible to this process used by PyNvVideoCodec",
+    )
    parser.add_argument("--verify-every", type=int, default=0)
    parser.add_argument("--verify-continuous-every", type=int, default=0)
    parser.add_argument("--no-progress", action="store_true")
@@ -742,6 +921,8 @@ def main() -> None:
        output_fps=args.video_fps,
        required_duration=required_duration,
        frame_cache_size=args.frame_cache_size,
+        decoder_backend=args.decoder,
+        decoder_gpu_id=args.decoder_gpu_id,
    )
    is_easycom = (
        args.easycom_session_dir is not None
@@ -772,6 +953,10 @@ def main() -> None:
                )
    _print_media_summary(reader, args.start, end)
    if args.media_only:
+        # Exercise the selected pixel decoder as well as its metadata path.
+        probe_end = min(end, args.start + max(args.step_seconds, 1.0 / reader.fps))
+        reader.frames(args.start, probe_end)
+        print(f"Decoded media probe successfully with {args.decoder}")
        return

    duration = end - args.start