import base64 import configparser import functools import hashlib import io import math import os import tempfile import threading import time from collections import OrderedDict from dataclasses import dataclass from typing import Any, Literal, TypedDict from urllib.parse import parse_qs import numpy as np from PIL import Image class ImageInput(TypedDict): type: Literal['image'] image: Image.Image class VideoInput(TypedDict, total=False): type: Literal['video'] video: Any source_fps: float fps: float high_motion: bool max_tokens: int max_tokens_per_frame: int max_frames: int MediaInput = ImageInput | VideoInput _TAR_FD_CACHE_SIZE = int(os.getenv("GROUNDINGPI_TAR_FD_CACHE_SIZE", "16")) _TAR_FD_CACHE: OrderedDict[str, int] = OrderedDict() _TAR_FD_CACHE_LOCK = threading.Lock() _OSS_BUCKET_CACHE: dict[tuple[str, str], Any] = {} _OSS_BUCKET_CACHE_LOCK = threading.Lock() _MEDIA_IO_COUNT = 0 def _parse_tar_media_ref(path: str): if "@" not in path or ".tar@" not in path: return None tar_path, query = path.split("@", 1) params = parse_qs(query, keep_blank_values=True) offsets = params.get("offset") sizes = params.get("size") if not offsets or not sizes: return None try: offset = int(offsets[0]) size = int(sizes[0]) except ValueError: return None if offset < 0 or size < 1: raise ValueError(f"Invalid TAR media range: offset={offset}, size={size}") return tar_path, offset, size, params.get("member", [None])[0] def _open_tar_fd(path: str) -> int: fd = _TAR_FD_CACHE.pop(path, None) if fd is None: fd = os.open(path, os.O_RDONLY) _TAR_FD_CACHE[path] = fd while len(_TAR_FD_CACHE) > _TAR_FD_CACHE_SIZE: _, stale_fd = _TAR_FD_CACHE.popitem(last=False) try: os.close(stale_fd) except OSError: pass return fd def _read_local_range(path: str, offset: int, size: int) -> bytes: if _TAR_FD_CACHE_SIZE > 0: with _TAR_FD_CACHE_LOCK: data = os.pread(_open_tar_fd(path), size, offset) else: fd = os.open(path, os.O_RDONLY) try: data = os.pread(fd, size, offset) finally: os.close(fd) if len(data) != size: raise EOFError( f"Short TAR media read: path={path}, offset={offset}, " f"size={size}, got={len(data)}" ) return data def _load_oss_config() -> dict[str, str]: config_path = os.getenv("GROUNDINGPI_OSS_CONFIG_FILE") or os.getenv( "OSSUTIL_CONFIG_FILE" ) values: dict[str, str] = {} if config_path: parser = configparser.ConfigParser() if not parser.read(config_path): raise FileNotFoundError(f"OSS config file not found: {config_path}") section = parser["Credentials"] if parser.has_section("Credentials") else parser.defaults() values = {key.lower(): value.strip() for key, value in section.items()} return values def _get_oss_bucket(uri: str): try: import oss2 except ImportError as exc: raise ImportError( "OSS TAR range reads require oss2; install the Python oss2 package" ) from exc rest = uri[len("oss://"):] if "/" not in rest: raise ValueError(f"OSS URI has no object key: {uri}") bucket_name, key = rest.split("/", 1) cfg = _load_oss_config() endpoint = ( os.getenv("OSS_ENDPOINT") or cfg.get("endpoint") or "http://oss-cn-wulanchabu-internal.aliyuncs.com" ).strip() access_key_id = ( os.getenv("OSS_ACCESS_KEY_ID") or cfg.get("accesskeyid") or "" ).strip() access_key_secret = ( os.getenv("OSS_ACCESS_KEY_SECRET") or cfg.get("accesskeysecret") or "" ).strip() sts_token = (os.getenv("OSS_STS_TOKEN") or cfg.get("ststoken") or "").strip() if not access_key_id or not access_key_secret: raise ValueError( "OSS credentials are missing; set GROUNDINGPI_OSS_CONFIG_FILE or OSS environment variables" ) cache_key = (endpoint, bucket_name) with _OSS_BUCKET_CACHE_LOCK: bucket = _OSS_BUCKET_CACHE.get(cache_key) if bucket is None: auth = ( oss2.StsAuth(access_key_id, access_key_secret, sts_token) if sts_token else oss2.Auth(access_key_id, access_key_secret) ) bucket = oss2.Bucket(auth, endpoint, bucket_name) _OSS_BUCKET_CACHE[cache_key] = bucket return bucket, key def _maybe_log_media_io(kind: str, uri: str, size: int, seconds: float) -> None: global _MEDIA_IO_COUNT if os.getenv("LOCAL_RANK", "0") != "0": return _MEDIA_IO_COUNT += 1 every = max(1, int(os.getenv("GROUNDINGPI_MEDIA_IO_LOG_EVERY", "100"))) if _MEDIA_IO_COUNT > 3 and _MEDIA_IO_COUNT % every: return mib_per_second = size / max(seconds, 1e-9) / (1024 * 1024) source_id = hashlib.sha1(uri.encode("utf-8")).hexdigest()[:12] print( "[groundingpi_media_io] " f"kind={kind} source={source_id} bytes={size} " f"seconds={seconds:.4f} mib_per_second={mib_per_second:.2f}", flush=True, ) def _read_oss_bytes(uri: str, byte_range: tuple[int, int] | None = None) -> bytes: bucket, key = _get_oss_bucket(uri) start = time.perf_counter() obj = bucket.get_object(key, byte_range=byte_range) data = obj.read() elapsed = time.perf_counter() - start expected = None if byte_range is None else byte_range[1] - byte_range[0] + 1 if expected is not None and len(data) != expected: raise EOFError( f"Short OSS range read: uri={uri}, range={byte_range}, " f"expected={expected}, got={len(data)}" ) _maybe_log_media_io("oss_range" if byte_range else "oss_object", uri, len(data), elapsed) return data def load_media_bytes(source: str | bytes | bytearray | io.BytesIO) -> bytes: """Read local, local-TAR, raw OSS, or ranged OSS-TAR media bytes.""" if isinstance(source, bytes): return source if isinstance(source, bytearray): return bytes(source) if isinstance(source, io.BytesIO): return source.getvalue() if not isinstance(source, str): raise ValueError(f"Unsupported media byte source: {type(source)}") tar_ref = _parse_tar_media_ref(source) if tar_ref is not None: tar_path, offset, size, _ = tar_ref if tar_path.startswith("oss://"): return _read_oss_bytes(tar_path, (offset, offset + size - 1)) return _read_local_range(tar_path, offset, size) if source.startswith("oss://"): return _read_oss_bytes(source) with open(source, "rb") as handle: return handle.read() def load_image_source(source: Any) -> Image.Image: """Load an image from PIL, local path, TAR ref, OSS URI, or raw bytes.""" if isinstance(source, dict): source = source.get("image", source.get("path")) if isinstance(source, Image.Image): return source.convert("RGB") if isinstance(source, str) and source.startswith("data:"): return _to_pil(source) return Image.open(io.BytesIO(load_media_bytes(source))).convert("RGB") class VideoReader: """Small indexed video reader with dependency fallbacks. Frame lists and tensors are accepted directly. File inputs prefer decord so long videos are sampled without decoding every frame; OpenCV and torchvision are fallbacks for environments where decord is unavailable. """ def __init__(self, source: Any, source_fps: float | None = None): self.backend = None self.reader = None self.frames = None self.temp_path = None self.source_kind = "memory" self.source_bytes = 0 self.io_seconds = 0.0 self.source_fps = float(source_fps) if source_fps else None if isinstance(source, (list, tuple)): self.backend = 'frames' self.frames = list(source) self.frame_count = len(self.frames) else: try: import torch is_tensor = isinstance(source, torch.Tensor) except ImportError: is_tensor = False if isinstance(source, np.ndarray) or is_tensor: array = source.detach().cpu().numpy() if is_tensor else source if array.ndim != 4: raise ValueError( f'Video tensors must be 4D, got shape {array.shape}') if array.shape[-1] in (3, 4): pass elif array.shape[1] in (3, 4): array = array.transpose(0, 2, 3, 1) else: raise ValueError( 'Video tensors must be THWC or TCHW with 3/4 channels') self.backend = 'array' self.frames = array self.frame_count = int(array.shape[0]) elif isinstance(source, str): tar_ref = _parse_tar_media_ref(source) if tar_ref is not None or source.startswith("oss://"): start = time.perf_counter() data = load_media_bytes(source) self.io_seconds = time.perf_counter() - start self.source_bytes = len(data) self.source_kind = ( "oss_tar_range" if tar_ref is not None and source.startswith("oss://") else "local_tar_range" if tar_ref is not None else "oss_object" ) member = tar_ref[3] if tar_ref is not None else source suffix = os.path.splitext(member or "video.mp4")[1] or ".mp4" self._open_bytes(data, suffix=suffix) else: self.source_kind = "local_file" self._open_file(source) elif isinstance(source, (bytes, bytearray, io.BytesIO)): data = load_media_bytes(source) self.source_bytes = len(data) self._open_bytes(data) else: raise ValueError(f'Unsupported video data type: {type(source)}') if self.frame_count <= 0: raise ValueError('Video contains no frames') if not self.source_fps or not math.isfinite(self.source_fps) or self.source_fps <= 0: # A frame list without source_fps is treated as already sampled. self.source_fps = 1.0 def _open_bytes(self, data: bytes, suffix: str = ".mp4"): try: import decord self.reader = decord.VideoReader( io.BytesIO(data), ctx=decord.cpu(0), num_threads=1 ) self.backend = 'decord_bytes' self.frame_count = len(self.reader) if self.source_fps is None: self.source_fps = float(self.reader.get_avg_fps()) return except (ImportError, OSError, RuntimeError, TypeError): pass preferred_dir = os.getenv("GROUNDINGPI_MEDIA_TMP_DIR", "/dev/shm/groundingpi_media") try: os.makedirs(preferred_dir, exist_ok=True) temp_dir = preferred_dir except OSError: temp_dir = tempfile.gettempdir() with tempfile.NamedTemporaryFile( prefix="groundingpi_video_", suffix=suffix, dir=temp_dir, delete=False ) as handle: handle.write(data) self.temp_path = handle.name self._open_file(self.temp_path) def _open_file(self, path: str): try: import decord self.reader = decord.VideoReader(path, ctx=decord.cpu(0), num_threads=1) self.backend = 'decord' self.frame_count = len(self.reader) if self.source_fps is None: self.source_fps = float(self.reader.get_avg_fps()) return except (ImportError, OSError, RuntimeError): pass try: import cv2 reader = cv2.VideoCapture(path) if not reader.isOpened(): raise RuntimeError(f'OpenCV cannot open {path}') self.reader = reader self.backend = 'opencv' self.frame_count = int(reader.get(cv2.CAP_PROP_FRAME_COUNT)) if self.source_fps is None: self.source_fps = float(reader.get(cv2.CAP_PROP_FPS)) return except (ImportError, OSError, RuntimeError): pass try: from torchvision.io import read_video frames, _, info = read_video(path, pts_unit='sec', output_format='THWC') self.frames = frames.cpu().numpy() self.backend = 'array' self.frame_count = int(self.frames.shape[0]) if self.source_fps is None: self.source_fps = float(info.get('video_fps', 0.0)) return except Exception as exc: raise RuntimeError( f'Cannot decode video {path}; install decord, OpenCV, or a working torchvision video backend' ) from exc @property def duration(self) -> float: return self.frame_count / self.source_fps def get_frames(self, indices: list[int]) -> list[Image.Image]: if not indices: return [] indices = [min(max(int(i), 0), self.frame_count - 1) for i in indices] if self.backend in {'frames', 'array'}: return [_to_pil(self.frames[i]) for i in indices] if self.backend in {'decord', 'decord_bytes'}: arrays = self.reader.get_batch(indices).asnumpy() return [Image.fromarray(frame).convert('RGB') for frame in arrays] if self.backend == 'opencv': import cv2 frames = [] for index in indices: self.reader.set(cv2.CAP_PROP_POS_FRAMES, index) ok, frame = self.reader.read() if not ok: raise RuntimeError(f'Failed to decode video frame {index}') frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))) return frames raise RuntimeError(f'Unknown video backend: {self.backend}') def close(self): if self.backend == 'opencv' and self.reader is not None: self.reader.release() self.reader = None if self.temp_path is not None: try: os.remove(self.temp_path) except FileNotFoundError: pass self.temp_path = None def sample_video_indices(frame_count: int, source_fps: float, target_fps: float) -> list[int]: """Return monotonically increasing frame indices sampled by timestamps.""" target_fps = min(float(target_fps), float(source_fps)) duration = frame_count / source_fps sample_count = max(1, int(math.ceil(duration * target_fps))) times = np.arange(sample_count, dtype=np.float64) / target_fps indices = np.minimum(np.floor(times * source_fps).astype(np.int64), frame_count - 1) return np.unique(indices).tolist() def uniformly_limit_indices(indices: list[int], max_frames: int) -> list[int]: if len(indices) <= max_frames: return indices positions = np.linspace(0, len(indices) - 1, max_frames) return [indices[int(round(position))] for position in positions] def estimate_motion_score(frames: list[Image.Image], side: int = 96) -> float: """Mean normalized luminance change across sparse probe frames.""" if len(frames) < 2: return 0.0 probes = [] for frame in frames: gray = frame.convert('L') gray.thumbnail((side, side), Image.Resampling.BILINEAR) probes.append(np.asarray(gray, dtype=np.float32) / 255.0) scores = [float(np.mean(np.abs(curr - prev))) for prev, curr in zip(probes, probes[1:])] return float(np.mean(scores)) @dataclass class TransparentBgConfig: """The config of the transparent background.""" pattern: Literal["white", "black", "gray", "chessboard"] = "black" """The pattern of the transparent background.""" chessboard_square_size: int = 16 """The size of the squares in the chessboard background.""" chessboard_square_on_top_left: bool = True """Whether to start the chessboard with a white square on the top left.""" chessboard_white_value: int = 255 """The value of the white pixels in the background.""" chessboard_gray_value: int = 200 """The value of the gray pixels in the background.""" @functools.lru_cache(maxsize=256) def _create_chessboard_background( height: int, width: int, square_size: int, square_on_top_left: bool, white_value: int, gray_value: int, ) -> np.ndarray: """Create a chessboard background.""" bg = np.ones((height, width, 3), dtype=np.uint8) * white_value for y in range(0, height, square_size): for x in range(0, width, square_size): if (y // square_size + x // square_size) % 2 == ( 1 if square_on_top_left else 0): bg[y:y + square_size, x:x + square_size] = gray_value return bg def fill_transparent_bg_with( image: Image.Image, transparent_bg_config: TransparentBgConfig | None = None, ) -> Image.Image: """Composite a (possibly) transparent image onto a configured background. When ``transparent_bg_config`` is ``None``, the image is simply converted to RGB (preserving the historical behavior). Otherwise the alpha channel is alpha-composited over a background generated according to the config. """ if transparent_bg_config is None: return image.convert("RGB") if image.mode == "RGB": return image has_alpha = "A" in image.getbands() or "transparency" in image.info if not has_alpha: return image.convert("RGB") img = np.array(image.convert("RGBA")) height, width = img.shape[:2] bg_pattern = transparent_bg_config.pattern if bg_pattern == "white": bg = np.full((height, width, 3), 255, dtype=np.uint8) elif bg_pattern == "black": bg = np.zeros((height, width, 3), dtype=np.uint8) elif bg_pattern == "gray": bg = np.full((height, width, 3), 128, dtype=np.uint8) elif bg_pattern == "chessboard": bg = _create_chessboard_background( height, width, transparent_bg_config.chessboard_square_size, transparent_bg_config.chessboard_square_on_top_left, transparent_bg_config.chessboard_white_value, transparent_bg_config.chessboard_gray_value, ) else: raise ValueError(f"Invalid background pattern: {bg_pattern}") alpha = img[:, :, 3] img_rgb = img[:, :, :3] alpha_normalized = alpha.astype(np.float32) / 255.0 alpha_3d = np.stack([alpha_normalized] * 3, axis=2) result = alpha_3d * img_rgb + (1 - alpha_3d) * bg result = result.astype(np.uint8) return Image.fromarray(result) def navit_resize_image( width: int, height: int, patch_size: int, merge_kernel_size: int, in_patch_limit: int, patch_limit_on_one_side: int, fixed_output_tokens: int | None, ): # Apply the patch limits. s1 = math.sqrt( in_patch_limit / (max(1.0, width // patch_size) * max(1.0, height // patch_size))) s2 = patch_limit_on_one_side * patch_size / width s3 = patch_limit_on_one_side * patch_size / height scale = min(1.0, s1, s2, s3) new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale)) new_w = min(new_w, patch_limit_on_one_side * patch_size) new_h = min(new_h, patch_limit_on_one_side * patch_size) # Calculate the padding to make the height and width divisible by the merge kernel size and patch size. factor = merge_kernel_size * patch_size pad_height = (factor - new_h % factor) % factor pad_width = (factor - new_w % factor) % factor if fixed_output_tokens is not None: num_tokens = fixed_output_tokens else: # Calculate new dimensions after padding and patching token_height = (new_h + pad_height) // factor token_width = (new_w + pad_width) // factor assert token_height * merge_kernel_size <= patch_limit_on_one_side, ( f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}" ) assert token_width * merge_kernel_size <= patch_limit_on_one_side, ( f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}" ) num_tokens = token_height * token_width return { "num_tokens": num_tokens, "new_width": new_w, "new_height": new_h, "pad_width": pad_width, "pad_height": pad_height, "sampled_nframes": 1, } def navit_resize_video_frame( width: int, height: int, patch_size: int, merge_kernel_size: int, in_patch_limit: int, patch_limit_on_one_side: int, max_output_tokens: int, ): """Resize one video frame while strictly enforcing its LLM token cap.""" if max_output_tokens < 1: raise ValueError('max_output_tokens must be at least 1') high = max(1, min(int(in_patch_limit), int(max_output_tokens) * merge_kernel_size**2)) low = 1 best = None while low <= high: current = (low + high) // 2 candidate = navit_resize_image( width, height, patch_size, merge_kernel_size, current, patch_limit_on_one_side, None, ) if candidate['num_tokens'] <= max_output_tokens: best = candidate low = current + 1 else: high = current - 1 if best is None: # One merged 2x2 patch is the smallest representable visual grid. best = navit_resize_image( width, height, patch_size, merge_kernel_size, 1, patch_limit_on_one_side, None, ) if best['num_tokens'] > max_output_tokens: raise ValueError( f'Unable to resize frame under {max_output_tokens} tokens: {best}') return best def _to_pil( data: str | bytes | Image.Image | np.ndarray, transparent_bg_config: TransparentBgConfig | None = None, to_rgb: bool = True, ) -> Image.Image: """Load an image and (optionally) composite its transparent background. Args: data: A PIL Image, a base64 ``data:`` URL, a file path, or raw bytes. transparent_bg_config: The config used to fill the transparent background. ``None`` keeps the historical behavior of converting to RGB without compositing. to_rgb: If ``False`` the image is returned as-is (the ``transparent_bg_config`` is ignored). The caller is then expected to call :func:`fill_transparent_bg_with` later — e.g. after a resize. """ if isinstance(data, Image.Image): image = data elif isinstance(data, str): if data.startswith("data:"): raw_base64 = data.split(",")[1] image = Image.open(io.BytesIO(base64.b64decode(raw_base64))) else: image = Image.open(data) elif isinstance(data, bytes): image = Image.open(io.BytesIO(data)) elif isinstance(data, np.ndarray): if data.dtype != np.uint8: if np.issubdtype(data.dtype, np.floating) and data.max(initial=0) <= 1.0: data = data * 255.0 data = np.clip(data, 0, 255).astype(np.uint8) image = Image.fromarray(data) else: raise ValueError(f"Unsupported data type: {type(data)}") if not to_rgb: return image return fill_transparent_bg_with(image, transparent_bg_config) def ensure_media_type( media: MediaInput, transparent_bg_config: TransparentBgConfig | None = None, transparent_bg_fill_stage: Literal["before_resize", "after_resize"] = "before_resize", ) -> MediaInput: if media['type'] == 'image': media['image'] = _to_pil( media['image'], transparent_bg_config=transparent_bg_config, to_rgb=transparent_bg_fill_stage == "before_resize", ) return media else: raise ValueError(f"Unsupported media type: {media['type']}") def image_to_np( image: Image.Image, resize_to: tuple[int, int] | None = None, mode: str = "resize", raise_error_for_ill_resize: bool = True, transparent_bg_config: TransparentBgConfig | None = None, transparent_bg_fill_stage: Literal["before_resize", "after_resize"] = "before_resize", ) -> np.ndarray: """Convert an image to a numpy array. Args: content: The image to convert. resize_to: The size to resize the image to. mode: The mode to resize the image to. raise_error_for_ill_resize: Whether to raise an error for ill-sized resize. transparent_bg_config: The config of the transparent background. Only used when ``transparent_bg_fill_stage == "after_resize"`` (the caller is responsible for filling before resize otherwise). transparent_bg_fill_stage: When to composite the transparent background — before or after the resize step. Returns: A numpy array. """ assert isinstance(image, Image.Image), "image must be a PIL Image" if resize_to is not None: if mode == "resize": image = image.resize(resize_to, resample=Image.Resampling.BICUBIC) if transparent_bg_fill_stage == "after_resize": image = fill_transparent_bg_with(image, transparent_bg_config) elif mode == "rescale_and_pad_to_center": scale = min(resize_to[0] / image.width, resize_to[1] / image.height, 1.0) new_width = round(image.width * scale) new_height = round(image.height * scale) if new_width == 0 or new_height == 0: if raise_error_for_ill_resize: raise ValueError( f"Invalid resize to: {resize_to}, from image size: {image.size}" ) else: return np.zeros((resize_to[1], resize_to[0], 3), dtype=np.uint8) image = image.resize((new_width, new_height), resample=Image.Resampling.BICUBIC) if transparent_bg_fill_stage == "after_resize": image = fill_transparent_bg_with(image, transparent_bg_config) padding_left = (resize_to[0] - new_width) // 2 padding_right = resize_to[0] - new_width - padding_left padding_top = (resize_to[1] - new_height) // 2 padding_bottom = resize_to[1] - new_height - padding_top image = np.asarray(image) image = np.pad( image, ((padding_top, padding_bottom), (padding_left, padding_right), (0, 0)), mode="constant", constant_values=0, ) assert image.shape == (resize_to[1], resize_to[0], 3) elif mode == "rescale_and_pad_to_rightbottom": scale = min(resize_to[0] / image.width, resize_to[1] / image.height, 1.0) new_width = round(image.width * scale) new_height = round(image.height * scale) if new_width == 0 or new_height == 0: if raise_error_for_ill_resize: raise ValueError( f"Invalid resize to: {resize_to}, from image size: {image.size}" ) else: return np.zeros((resize_to[1], resize_to[0], 3), dtype=np.uint8) image = image.resize((new_width, new_height), resample=Image.Resampling.BICUBIC) if transparent_bg_fill_stage == "after_resize": image = fill_transparent_bg_with(image, transparent_bg_config) padding_right = resize_to[0] - new_width padding_bottom = resize_to[1] - new_height image = np.asarray(image) image = np.pad( image, ((0, padding_bottom), (0, padding_right), (0, 0)), mode="constant", constant_values=0, ) assert image.shape == (resize_to[1], resize_to[0], 3) else: raise ValueError(f"Invalid mode: {mode}") if isinstance(image, Image.Image): return np.asarray(image) else: return image def navit_patchify(pixel_values: np.ndarray, patch_size: int) -> dict[str, np.ndarray]: """Reshape the pixel values to a navit shape. Args: pixel_values: np.ndarray, shape (t, h, w, c) patch_size: int Returns: dict[str, np.ndarray] - patches: np.ndarray, shape (t * h//patch_size * w//patch_size, c, patch_size, patch_size) - grid_thw: np.ndarray, (t, h//patch_size, w//patch_size) """ T, H, W, C = pixel_values.shape assert C == 3, "pixel_values must have 3 channels" patches = pixel_values.reshape(T, H // patch_size, patch_size, W // patch_size, patch_size, C) # (T, H//patch_size, W//patch_size, C, patch_size, patch_size) patches = patches.transpose(0, 1, 3, 5, 2, 4) patches = patches.reshape(-1, C, patch_size, patch_size) grid_thw = np.array([T, H // patch_size, W // patch_size]) return {"pixel_values": patches, "grid_thw": grid_thw} def normalize(x: np.ndarray, mean, std_inv, pixels_dtype: np.dtype = np.float32) -> np.ndarray: """Normalize the image. Args: x: The image to normalize. The shape is (..., 3). The dtype is uint8. The range is [0, 255]. mean: The mean of the image. std_inv: The inverse of the std of the image. pixels_dtype: The dtype of the image. Returns: The normalized image. The shape is (..., 3). The dtype is determined by the pixels_dtype. """ x = (x / 255.0).astype(pixels_dtype) x -= mean x *= std_inv return x def _to_tensor(data, **kwargs): import torch if isinstance(data, np.ndarray): return torch.from_numpy(data).to(**kwargs) elif isinstance(data, torch.Tensor): return data.to(**kwargs) elif isinstance(data, list): return [_to_tensor(item, **kwargs) for item in data] elif isinstance(data, tuple): return tuple(_to_tensor(item, **kwargs) for item in data) elif isinstance(data, dict): return {k: _to_tensor(v, **kwargs) for k, v in data.items()} elif data is None: return None else: raise ValueError(f"Unsupported data type: {type(data)}")