dexverse-cli 0.0.2__py3-none-any.whl

dexverse/convert/image_video_encode.py

@@ -0,0 +1,449 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2021-2025 DexForce Technology Co., Ltd.
+#
+# All rights reserved.
+# ----------------------------------------------------------------------------
+"""image_video_encode.py — RGB / depth 图像与视频（MP4 / AVI）编解码工具。
+
+由 ``dexverse.convert.hdf5_to_raw`` 复用，提供把 MP4 字节流解码回 RGB 帧的
+能力（``decode_mp4_bytes_to_rgb_images``）。
+"""
+import numpy as np
+from PIL import Image
+from typing import List, Optional
+import av
+from fractions import Fraction
+import io
+import os
+from dataclasses import dataclass
+import logging
+
+logger = logging.getLogger(__name__)
+
+# List of hardware encoders to probe for auto-selection. Availability depends on the platform and FFmpeg build.
+# Determines the order of preference for auto-selection when vcodec="auto" is used.
+HW_ENCODERS = [
+    "h264_videotoolbox",  # macOS
+    "hevc_videotoolbox",  # macOS
+    "h264_nvenc",  # NVIDIA GPU
+    "hevc_nvenc",  # NVIDIA GPU
+    "h264_vaapi",  # Linux Intel/AMD
+    "h264_qsv",  # Intel Quick Sync
+]
+
+VALID_VIDEO_CODECS = {"h264", "hevc", "libsvtav1", "auto"} | set(HW_ENCODERS)
+
+
+def _get_codec_options(
+        vcodec: str,
+        g: Optional[int] = 2,
+        crf: Optional[int] = 30,
+        preset: Optional[int] = None,
+) -> dict:
+    """Build codec-specific options dict for video encoding."""
+    options = {}
+
+    # GOP size (keyframe interval) - supported by VideoToolbox and software encoders
+    if g is not None and (vcodec in ("h264_videotoolbox", "hevc_videotoolbox") or vcodec not in HW_ENCODERS):
+        options["g"] = str(g)
+
+    # Quality control (codec-specific parameter names)
+    if crf is not None:
+        if vcodec in ("h264", "hevc", "libsvtav1"):
+            options["crf"] = str(crf)
+        elif vcodec in ("h264_videotoolbox", "hevc_videotoolbox"):
+            quality = max(1, min(100, int(100 - crf * 2)))
+            options["q:v"] = str(quality)
+        elif vcodec in ("h264_nvenc", "hevc_nvenc"):
+            options["rc"] = "constqp"
+            options["qp"] = str(crf)
+        elif vcodec in ("h264_vaapi",):
+            options["qp"] = str(crf)
+        elif vcodec in ("h264_qsv",):
+            options["global_quality"] = str(crf)
+
+    # Preset (only for libsvtav1)
+    if vcodec == "libsvtav1":
+        options["preset"] = str(preset) if preset is not None else "12"
+
+    return options
+
+
+def detect_available_hw_encoders() -> List[str]:
+    """Probe PyAV/FFmpeg for available hardware video encoders."""
+    available = []
+    for codec_name in HW_ENCODERS:
+        try:
+            av.codec.Codec(codec_name, "w")
+            available.append(codec_name)
+        except Exception:  # nosec B110
+            logger.debug("HW encoder '%s' not available", codec_name)  # nosec B110
+    return available
+
+
+def resolve_vcodec(vcodec: str) -> str:
+    """Validate vcodec and resolve 'auto' to best available HW encoder, fallback to libsvtav1."""
+    if vcodec not in VALID_VIDEO_CODECS:
+        raise ValueError(f"Invalid vcodec '{vcodec}'. Must be one of: {sorted(VALID_VIDEO_CODECS)}")
+    if vcodec != "auto":
+        logger.info(f"Using video codec: {vcodec}")
+        return vcodec
+    available = detect_available_hw_encoders()
+    for encoder in HW_ENCODERS:
+        if encoder in available:
+            logger.info(f"Auto-selected video codec: {encoder}")
+            return encoder
+    logger.info("No hardware encoder available, falling back to software encoder 'libsvtav1'")
+    return "libsvtav1"
+
+
+def encode_rgb_images_as_mp4_bytes_impl(
+        images: List['Image.Image'],
+        fps: int = 30,
+        vcodec: str = "libsvtav1",
+        pix_fmt: str = "yuv420p",
+        g: Optional[int] = 2,
+        crf: Optional[int] = 30,
+        fast_decode: int = 0,
+        log_level: Optional[int] = av.logging.WARNING,
+        preset: Optional[int] = None
+) -> Optional[bytes]:
+    if not images:
+        return None
+
+    # Resolve vcodec
+    vcodec = resolve_vcodec(vcodec)
+
+    # Get image size from the first image
+    width, height = images[0].size
+
+    # Encoders/pixel formats incompatibility check
+    if (vcodec == "libsvtav1" or vcodec == "hevc") and pix_fmt == "yuv444p":
+        logger.warning(
+            f"Incompatible pixel format 'yuv444p' for codec {vcodec}, auto-selecting format 'yuv420p'"
+        )
+        pix_fmt = "yuv420p"
+
+    # Define video codec options
+    video_options = _get_codec_options(vcodec, g, crf, preset)
+
+    if fast_decode:
+        key = "svtav1-params" if vcodec == "libsvtav1" else "tune"
+        value = f"fast-decode={fast_decode}" if vcodec == "libsvtav1" else "fastdecode"
+        video_options[key] = value
+
+    # Set logging level
+    if log_level is not None:
+        logging.getLogger("libav").setLevel(log_level)
+
+    # Create a BytesIO object to hold the MP4 data
+    output_buffer = io.BytesIO()
+
+    # Create container with the BytesIO object
+    container = av.open(output_buffer, 'w', format='mp4')
+
+    # Add video stream
+    stream = container.add_stream(vcodec, rate=fps, options=video_options)
+    stream.width = width
+    stream.height = height
+    stream.pix_fmt = pix_fmt
+    stream.time_base = Fraction(1, fps)
+
+    # Encode frames
+    for i, img in enumerate(images):
+        # Convert PIL Image to AV VideoFrame
+        frame = av.VideoFrame.from_image(img.convert('RGB'))
+        frame.pts = i
+        frame.time_base = Fraction(1, fps)
+
+        # Encode frame
+        for packet in stream.encode(frame):
+            container.mux(packet)
+
+    # Flush encoder
+    for packet in stream.encode():
+        container.mux(packet)
+
+    # Reset logging level
+    if log_level is not None:
+        av.logging.restore_default_callback()
+
+    # Close container to ensure all data is written
+    container.close()
+
+    # Get the bytes from the buffer
+    output_buffer.seek(0)
+    output_bytes = output_buffer.getvalue()
+
+    return output_bytes
+
+
+@dataclass(frozen=True)
+class EncodeImageIntoMp4Option(object):
+    fps: int = 30
+    vcodec: str = "libsvtav1"
+    pix_fmt: str = "yuv420p"
+    g: Optional[int] = 2
+    crf: Optional[int] = 30
+    fast_decode: int = 0
+    log_level: Optional[int] = av.logging.WARNING
+    preset: Optional[int] = None
+
+
+def encode_rgb_images_as_mp4_bytes(
+        images: List['Image.Image'],
+        option_in: Optional[EncodeImageIntoMp4Option] = None
+) -> Optional[bytes]:
+    """Encode a list of RGB images into MP4 bytes.
+
+    Args:
+        images: List of PIL Image objects to encode.
+        option_in: Optional EncodeImageIntoMp4Option object with encoding parameters.
+            If None, default values will be used.
+
+    Returns:
+        MP4-encoded bytes, or None if no images were provided.
+    """
+    if option_in is None:
+        option_in = EncodeImageIntoMp4Option()
+
+    return encode_rgb_images_as_mp4_bytes_impl(
+        images=images,
+        fps=option_in.fps,
+        vcodec=option_in.vcodec,
+        pix_fmt=option_in.pix_fmt,
+        g=option_in.g,
+        crf=option_in.crf,
+        fast_decode=option_in.fast_decode,
+        log_level=option_in.log_level,
+        preset=option_in.preset
+    )
+
+
+def decode_mp4_bytes_to_rgb_images(encoded_bytes: bytes) -> List['Image.Image']:
+    # Create a BytesIO object from the encoded bytes
+    input_buffer = io.BytesIO(encoded_bytes)
+
+    # Open the buffer as a container
+    container = av.open(input_buffer, 'r')
+
+    # Iterate through frames
+    images: List['Image.Image'] = list()
+    for frame in container.decode(video=0):
+        # Convert AV VideoFrame to PIL Image
+        img = frame.to_image()
+        images.append(img)
+
+    # Done
+    container.close()
+    return images
+
+
+def encode_depth_images_as_avi_bytes_impl(
+        images: List[np.ndarray],
+        fps: int = 30,
+        vcodec: str = "ffv1",
+        pix_fmt: str = "gray16le",
+        log_level: Optional[int] = av.logging.WARNING,
+) -> Optional[bytes]:
+    """Encode a list of uint16 depth images into a video bytes.
+
+    Args:
+        images: List of uint16 numpy arrays (depth images).
+        fps: Frames per second for the video.
+        vcodec: Video codec to use. Default is 'ffv1' (lossless).
+        pix_fmt: Pixel format. Default is 'gray16le' (16-bit grayscale, little-endian).
+        log_level: PyAV logging level.
+
+    Returns:
+        Encoded video bytes, or None if no images were provided.
+    """
+    if images is None or len(images) == 0:
+        return None
+
+    # Check data type
+    for img in images:
+        if img.dtype != np.uint16:
+            return None
+
+    # Input checked
+    height, width = images[0].shape
+    if log_level is not None:
+        logging.getLogger("libav").setLevel(log_level)
+
+    output_buffer = io.BytesIO()
+    container = av.open(output_buffer, 'w', format='avi')
+
+    stream = container.add_stream(vcodec, rate=fps)
+    stream.width = width
+    stream.height = height
+    stream.pix_fmt = pix_fmt
+
+    for i, img in enumerate(images):
+        assert img.dtype == np.uint16
+        frame = av.VideoFrame.from_ndarray(img, format=pix_fmt)
+        frame.pts = i
+        frame.time_base = Fraction(1, fps)
+
+        for packet in stream.encode(frame):
+            container.mux(packet)
+
+    for packet in stream.encode():
+        container.mux(packet)
+
+    if log_level is not None:
+        av.logging.restore_default_callback()
+
+    container.close()
+    output_buffer.seek(0)
+    return output_buffer.getvalue()
+
+
+@dataclass(frozen=True)
+class EncodeDepthImageIntoAVIOption(object):
+    fps: int = 30
+    vcodec: str = "ffv1"
+    pix_fmt: str = "gray16le"
+    log_level: Optional[int] = av.logging.WARNING
+
+
+def encode_depth_images_as_avi_bytes(
+        images: List[np.ndarray],
+        option_in: Optional[EncodeDepthImageIntoAVIOption] = None) -> Optional[bytes]:
+    """Encode a list of depht images into avi bytes.
+
+    Args:
+        images: List of depth images objects to encode.
+        option_in: Optional EncodeDepthImageIntoAVIOption object with encoding parameters.
+                If None, default values will be used.
+
+    Returns:
+        avi-encoded bytes, or None if no images were provided.
+    """
+    if option_in is None:
+        option_in = EncodeDepthImageIntoAVIOption()
+    return encode_depth_images_as_avi_bytes_impl(images=images, fps=option_in.fps, vcodec=option_in.vcodec,
+                                                 pix_fmt=option_in.pix_fmt, log_level=option_in.log_level)
+
+
+def encode_metric_depth_images_as_avi_bytes(
+        images: List[np.ndarray],
+        option_in: Optional[EncodeDepthImageIntoAVIOption] = None) -> Optional[bytes]:
+    # Convert the input
+    depth_images_uint16: List[np.ndarray] = list()
+    for img in images:
+        if img.dtype == np.float32:
+            img_mm = img * 1000.0
+            depth_images_uint16.append(img_mm.astype(np.uint16))
+        elif img.dtype == np.uint16:
+            depth_images_uint16.append(img)
+        else:
+            return None
+
+    return encode_depth_images_as_avi_bytes(depth_images_uint16, option_in)
+
+
+def decode_avi_bytes_to_depth_images(
+        encoded_bytes: bytes,
+        pix_fmt: str = "gray16le",
+) -> List[np.ndarray]:
+    """Decode video bytes back to uint16 depth images.
+
+    Args:
+        encoded_bytes: Encoded video bytes.
+        pix_fmt: Pixel format used during encoding. Default is 'gray16le'.
+
+    Returns:
+        List of uint16 numpy arrays (depth images).
+    """
+    # Make input buffer
+    input_buffer = io.BytesIO(encoded_bytes)
+    container = av.open(input_buffer, 'r')
+
+    # Load the images
+    images: List[np.ndarray] = []
+    for frame in container.decode(video=0):
+        img = frame.to_ndarray(format=pix_fmt)
+        images.append(img)
+
+    # Done
+    container.close()
+    return images
+
+
+def encode_rgb_image_path_into_mp4(
+        image_path: List[str],
+        mp4_output_path: str,
+        option_in: Optional[EncodeImageIntoMp4Option] = None):
+    """Encode images from file paths into an MP4 file.
+
+    Args:
+        image_path: List of paths to image files.
+        mp4_output_path: Path to save the output MP4 file.
+        option_in: Optional EncodeImageIntoMp4Option object with encoding parameters.
+            If None, default values will be used.
+    """
+    # Read images from file paths
+    images = []
+    for path in image_path:
+        with Image.open(path) as img:
+            images.append(img.convert('RGB'))
+
+    # Encode images to MP4 bytes
+    encoded_bytes = encode_rgb_images_as_mp4_bytes(images, option_in)
+
+    if encoded_bytes is None:
+        logger.warning("No images to encode")
+        return
+
+    # Create the output directory if it doesn't exist
+    output_dir = os.path.dirname(mp4_output_path)
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+
+    # Write the MP4 bytes to file
+    with open(mp4_output_path, 'wb') as f:
+        f.write(encoded_bytes)
+
+    logger.info(f"Encoded {len(image_path)} images into {mp4_output_path}")
+
+
+def decode_mp4_file_into_images_directory(mp4_file_path: str, save_image_dir: str):
+    """Decode an MP4 file into a directory of images.
+
+    Args:
+        mp4_file_path: Path to the MP4 file to decode.
+        save_image_dir: Directory to save the decoded images.
+    """
+    # Create the save directory if it doesn't exist
+    os.makedirs(save_image_dir, exist_ok=True)
+
+    # Read the MP4 file as bytes
+    with open(mp4_file_path, 'rb') as f:
+        mp4_bytes = f.read()
+
+    # Decode the MP4 bytes to images
+    images = decode_mp4_bytes_to_rgb_images(mp4_bytes)
+
+    # Save each image with a sequential filename
+    for i, img in enumerate(images):
+        img_path = os.path.join(save_image_dir, f"frame_{i:06d}.png")
+        img.save(img_path)
+
+    # Done
+    logger.info(f"Decoded MP4 file {mp4_file_path} into {save_image_dir}")
+
+
+__all__ = [
+    "encode_rgb_images_as_mp4_bytes_impl",
+    "EncodeImageIntoMp4Option",
+    "encode_rgb_images_as_mp4_bytes",
+    "decode_mp4_bytes_to_rgb_images",
+    "encode_depth_images_as_avi_bytes_impl",
+    "EncodeDepthImageIntoAVIOption",
+    "encode_depth_images_as_avi_bytes",
+    "encode_metric_depth_images_as_avi_bytes",
+    "decode_avi_bytes_to_depth_images",
+    "encode_rgb_image_path_into_mp4",
+    "decode_mp4_file_into_images_directory",
+]

dexverse/downloader/__init__.py

@@ -0,0 +1,26 @@
+"""Download core: manager, backends, verify, state."""
+from __future__ import annotations
+
+from .manager import DownloadManager, DownloadReport, get_backend, repair
+from .base import Backend, Credentials, FileResult
+from .s5cmd_backend import S5CmdBackend, S5CmdNotAvailable, generate_s5cmd_lines
+from .python_s3_backend import PythonS3Backend
+from .http_backend import HttpBackend
+from .verify import verify_files, summarize
+
+__all__ = [
+    "DownloadManager",
+    "DownloadReport",
+    "get_backend",
+    "repair",
+    "Backend",
+    "Credentials",
+    "FileResult",
+    "S5CmdBackend",
+    "S5CmdNotAvailable",
+    "generate_s5cmd_lines",
+    "PythonS3Backend",
+    "HttpBackend",
+    "verify_files",
+    "summarize",
+]

dexverse/downloader/base.py

@@ -0,0 +1,101 @@
+"""Backend contract + shared credential resolution.
+
+A *backend* takes the manifest, the subset of files that still need fetching,
+the out root, and returns a per-object success/failure outcome. The download
+manager owns skip/resume/verify; backends only move bytes.
+"""
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Protocol
+
+from ..manifest import FileSpec, Manifest
+
+
+@dataclass
+class Credentials:
+    access_key: Optional[str]
+    secret_key: Optional[str]
+    session_token: Optional[str]
+    region: Optional[str]
+    endpoint: Optional[str]
+
+    @classmethod
+    def from_manifest(cls, manifest: Manifest) -> "Credentials":
+        """Prefer manifest credentials; fall back to AWS_* env vars."""
+        storage = manifest.storage
+        endpoint = storage.endpoint if storage else None
+        region = (storage.region if storage else None) or os.getenv("AWS_REGION") \
+            or os.getenv("AWS_DEFAULT_REGION") or "us-east-1"
+
+        ak = sk = st = None
+        if storage and storage.credentials and storage.credentials.is_populated():
+            ak = storage.credentials.accessKey
+            sk = storage.credentials.secretKey
+            st = storage.credentials.sessionToken or None
+        else:
+            ak = os.getenv("AWS_ACCESS_KEY_ID")
+            sk = os.getenv("AWS_SECRET_ACCESS_KEY")
+            st = os.getenv("AWS_SESSION_TOKEN") or None
+
+        return cls(
+            access_key=ak, secret_key=sk, session_token=st,
+            region=region, endpoint=endpoint,
+        )
+
+    def as_env(self) -> Dict[str, str]:
+        """A dict suitable for merging into a subprocess environment."""
+        env: Dict[str, str] = {}
+        if self.access_key:
+            env["AWS_ACCESS_KEY_ID"] = self.access_key
+        if self.secret_key:
+            env["AWS_SECRET_ACCESS_KEY"] = self.secret_key
+        if self.session_token:
+            env["AWS_SESSION_TOKEN"] = self.session_token
+        if self.region:
+            env["AWS_REGION"] = self.region
+            env["AWS_DEFAULT_REGION"] = self.region
+        return env
+
+
+@dataclass
+class FileResult:
+    objectKey: str
+    ok: bool
+    error: Optional[str] = None
+    localSize: Optional[int] = None
+    expectedSize: Optional[int] = None
+
+
+class ProgressReporter(Protocol):
+    """Optional byte-progress sink supplied by interactive callers."""
+
+    def start(self, *, total_bytes: Optional[int], total_files: int) -> None:
+        ...
+
+    def advance(self, bytes_count: int) -> None:
+        ...
+
+    def complete_file(self) -> None:
+        ...
+
+    def finish(self) -> None:
+        ...
+
+
+class Backend(Protocol):
+    name: str
+
+    def download(
+        self,
+        manifest: Manifest,
+        files: List[FileSpec],
+        out: Path,
+        *,
+        parallel: int,
+        log_path: Optional[Path] = None,
+        progress: Optional[ProgressReporter] = None,
+    ) -> Dict[str, FileResult]:
+        ...