| 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: |
| |
| 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, |
| ): |
| |
| 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) |
|
|
| |
| 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: |
| |
| 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: |
| |
| 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) |
| |
| 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)}") |
|
|
