import base64 import functools import io import math from dataclasses import dataclass from typing import Literal, TypedDict import numpy as np from PIL import Image class ImageInput(TypedDict): type: Literal['image'] image: Image.Image MediaInput = ImageInput @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 _to_pil( data: str | bytes | Image.Image, 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)) 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)}")