File size: 6,917 Bytes
267d7ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | """Image preprocessing for Unlimited-OCR compatible with MLX.
Handles image loading, tiling, normalization, and batch preparation.
"""
import math
from typing import List, Tuple, Optional
from io import BytesIO
import numpy as np
from PIL import Image, ImageOps
def load_image(image_path: str) -> Optional[Image.Image]:
"""Load an image with EXIF orientation correction."""
try:
image = Image.open(image_path)
corrected = ImageOps.exif_transpose(image)
return corrected.convert("RGB")
except Exception as e:
print(f"Error loading image {image_path}: {e}")
return None
def find_closest_aspect_ratio(
aspect_ratio: float,
target_ratios: List[Tuple[int, int]],
width: int,
height: int,
image_size: int,
) -> Tuple[int, int]:
"""Find the closest allowed aspect ratio for tiling."""
best_ratio_diff = float('inf')
best_ratio = (1, 1)
area = width * height
for ratio in target_ratios:
target_aspect = ratio[0] / ratio[1]
ratio_diff = abs(aspect_ratio - target_aspect)
if ratio_diff < best_ratio_diff:
best_ratio_diff = ratio_diff
best_ratio = ratio
elif ratio_diff == best_ratio_diff:
if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
best_ratio = ratio
return best_ratio
def dynamic_preprocess(
image: Image.Image,
min_num: int = 2,
max_num: int = 32,
image_size: int = 640,
use_thumbnail: bool = False,
) -> Tuple[List[Image.Image], Tuple[int, int]]:
"""Dynamically tile an image into patches.
Args:
image: Input PIL image
min_num: Minimum number of patches
max_num: Maximum number of patches
image_size: Size of each patch
use_thumbnail: Whether to include a thumbnail
Returns:
Tuple of (list of patch images, aspect ratio)
"""
orig_width, orig_height = image.size
aspect_ratio = orig_width / orig_height
# Generate valid aspect ratios
target_ratios = set()
for n in range(min_num, max_num + 1):
for i in range(1, n + 1):
for j in range(1, n + 1):
if min_num <= i * j <= max_num:
target_ratios.add((i, j))
target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
# Find best ratio
target_aspect_ratio = find_closest_aspect_ratio(
aspect_ratio, target_ratios, orig_width, orig_height, image_size
)
target_width = image_size * target_aspect_ratio[0]
target_height = image_size * target_aspect_ratio[1]
blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
# Resize and crop patches
resized_img = image.resize((target_width, target_height))
processed_images = []
for i in range(blocks):
col = i % target_aspect_ratio[0]
row = i // target_aspect_ratio[0]
box = (
col * image_size,
row * image_size,
(col + 1) * image_size,
(row + 1) * image_size,
)
split_img = resized_img.crop(box)
processed_images.append(split_img)
if use_thumbnail and len(processed_images) != 1:
thumbnail_img = image.resize((image_size, image_size))
processed_images.append(thumbnail_img)
return processed_images, target_aspect_ratio
def preprocess_image(
image: Image.Image,
base_size: int = 1024,
image_size: int = 640,
crop_mode: bool = True,
) -> Tuple[np.ndarray, np.ndarray, Optional[Tuple], np.ndarray]:
"""Preprocess an image for the model.
Args:
image: Input PIL image
base_size: Base image size for the global view (1024)
image_size: Tile size for patches (640)
crop_mode: Whether to use dynamic tiling
Returns:
Tuple of (patches_array, original_array, crop_shape, num_image_tokens)
"""
# Normalize transform
mean = np.array([0.5, 0.5, 0.5], dtype=np.float32)
std = np.array([0.5, 0.5, 0.5], dtype=np.float32)
def to_tensor(img: Image.Image) -> np.ndarray:
arr = np.array(img, dtype=np.float32) / 255.0
arr = (arr - mean) / std
return arr.transpose(2, 0, 1) # [C, H, W]
if crop_mode:
# Dynamic tiling
patches, crop_shape = dynamic_preprocess(
image, min_num=2, max_num=32, image_size=image_size
)
patches_arr = np.stack([to_tensor(p) for p in patches], axis=0) # [N, 3, 640, 640]
# Global view
orig_img = image.resize((base_size, base_size))
orig_arr = to_tensor(orig_img)[np.newaxis, ...] # [1, 3, 1024, 1024]
# Number of image tokens
n_patches = len(patches)
local_tokens = n_patches * (image_size // 16) ** 2 # Each patch → 40*40 area
# After SAM: 40*40=1600 tokens per patch → CLIP processes them
# After CLIP: 256 tokens per patch (1024/4=256?)
# Let's compute from the architecture: image_size=640, patch=16 → 40x40=1600
# SAM output: 1024-dim, 16x16 spatial (net_3 output)
# CLIP output: concat [CLIP[:, 1:], SAM flatten] → 2048, 256 spatial
# For seq_mask: each image patch contributes a region
# The model handles this internally; we just need to track the crop shape
return patches_arr, orig_arr, crop_shape
else:
# Single image without tiling (base mode)
orig_img = image.resize((base_size, base_size))
orig_arr = to_tensor(orig_img)[np.newaxis, ...] # [1, 3, 1024, 1024]
# No patches
patches_arr = np.zeros((0, 3, image_size, image_size), dtype=np.float32)
return patches_arr, orig_arr, (1, 1)
def build_input(
input_ids: List[int],
image_features_count: int,
) -> Tuple[List[int], np.ndarray]:
"""Build input with image placeholder tokens.
Args:
input_ids: Text token ids
image_features_count: Number of image feature vectors to insert
Returns:
Tuple of (extended_input_ids, seq_mask)
"""
# Image placeholder tokens use token ID 0 (or special image token)
IMAGE_TOKEN_ID = 0 # This model uses BOS token as placeholder
# Insert image tokens after the image placeholder in the prompt
# The model's conversation format uses <image> tag
extended_ids = []
seq_mask = [] # True where image features go
i = 0
while i < len(input_ids):
extended_ids.append(input_ids[i])
seq_mask.append(False)
# After BOS (token 0), insert image placeholder positions
if input_ids[i] == 0 and image_features_count > 0:
# Extend with image placeholder positions
for _ in range(image_features_count):
extended_ids.append(0)
seq_mask.append(True)
image_features_count = 0 # Only insert once
i += 1
return extended_ids, np.array(seq_mask, dtype=bool)
|