_siglip.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284
  1. # Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. # This file is based on https://github.com/Kwai-Keye/Keye/blob/main/keye-vl-8b-preview/image_processing_keye.py
  15. # Original header:
  16. # Licensed under the Apache License, Version 2.0 (the "License");
  17. # you may not use this file except in compliance with the License.
  18. # You may obtain a copy of the License at
  19. #
  20. # http://www.apache.org/licenses/LICENSE-2.0
  21. #
  22. # Unless required by applicable law or agreed to in writing, software
  23. # distributed under the License is distributed on an "AS IS" BASIS,
  24. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  25. # See the License for the specific language governing permissions and
  26. # limitations under the License.
  27. """Image processor class for Keye."""
  28. # TODO: Support videos
  29. import json
  30. import math
  31. from pathlib import Path
  32. from typing import Dict, List, Optional, Union
  33. import numpy as np
  34. from ......utils import logging
  35. from ..common import (
  36. BatchFeature,
  37. convert_to_rgb,
  38. make_batched_images,
  39. make_list_of_images,
  40. to_numpy_array,
  41. )
  42. _OPENAI_CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
  43. _OPENAI_CLIP_STD = [0.26862954, 0.26130258, 0.27577711]
  44. def adjust_size(size, patch_size):
  45. num_patches = size // patch_size
  46. if num_patches % 2 != 0:
  47. num_patches -= 1
  48. return num_patches * patch_size
  49. def smart_resize(
  50. height: int,
  51. width: int,
  52. factor: int = 28,
  53. min_pixels: int = 28 * 28 * 130,
  54. max_pixels: int = 28 * 28 * 1280,
  55. ):
  56. """Rescales the image so that the following conditions are met:
  57. 1. Both dimensions (height and width) are divisible by 'factor'.
  58. 2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
  59. 3. The aspect ratio of the image is maintained as closely as possible.
  60. """
  61. # if height < factor or width < factor:
  62. # raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")
  63. # if int(height < factor//4) + int(width < factor//4):
  64. # raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor//4}")
  65. if height < factor:
  66. logging.debug(
  67. f"smart_resize: height={height} < factor={factor}, reset height=factor"
  68. )
  69. width = round((width * factor) / height)
  70. height = factor
  71. if width < factor:
  72. logging.debug(
  73. f"smart_resize: width={width} < factor={factor}, reset width=factor"
  74. )
  75. height = round((height * factor) / width)
  76. width = factor
  77. if max(height, width) / min(height, width) > 200:
  78. raise ValueError(
  79. f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
  80. )
  81. h_bar = round(height / factor) * factor
  82. w_bar = round(width / factor) * factor
  83. if h_bar * w_bar > max_pixels:
  84. beta = math.sqrt((height * width) / max_pixels)
  85. h_bar = math.floor(height / beta / factor) * factor
  86. w_bar = math.floor(width / beta / factor) * factor
  87. elif h_bar * w_bar < min_pixels:
  88. beta = math.sqrt(min_pixels / (height * width))
  89. h_bar = math.ceil(height * beta / factor) * factor
  90. w_bar = math.ceil(width * beta / factor) * factor
  91. return h_bar, w_bar
  92. class SiglipImageProcessor(object):
  93. model_input_names = [
  94. "pixel_values",
  95. "image_grid_thw",
  96. "pixel_values_videos",
  97. "video_grid_thw",
  98. ]
  99. def __init__(
  100. self,
  101. do_resize: bool = True,
  102. resample: int = 3,
  103. do_rescale: bool = True,
  104. rescale_factor: Union[int, float] = 1 / 255,
  105. do_normalize: bool = True,
  106. image_mean: Optional[Union[float, List[float]]] = None,
  107. image_std: Optional[Union[float, List[float]]] = None,
  108. do_convert_rgb: bool = True,
  109. min_pixels: int = 147384,
  110. max_pixels: int = 28 * 28 * 3600,
  111. patch_size: int = 14,
  112. temporal_patch_size: int = 1,
  113. merge_size: int = 2,
  114. **kwargs,
  115. ) -> None:
  116. super().__init__()
  117. self.do_resize = do_resize
  118. self.resample = resample
  119. self.do_rescale = do_rescale
  120. self.rescale_factor = rescale_factor
  121. self.do_normalize = do_normalize
  122. self.image_mean = image_mean if image_mean is not None else _OPENAI_CLIP_MEAN
  123. self.image_std = image_std if image_std is not None else _OPENAI_CLIP_STD
  124. self.min_pixels = min_pixels
  125. self.max_pixels = max_pixels
  126. self.patch_size = patch_size
  127. self.temporal_patch_size = temporal_patch_size
  128. self.merge_size = merge_size
  129. self.size = {"min_pixels": min_pixels, "max_pixels": max_pixels} # not used
  130. self.do_convert_rgb = do_convert_rgb
  131. @classmethod
  132. def from_pretrained(cls, pretrained_model_dir):
  133. pretrained_model_dir = Path(pretrained_model_dir)
  134. image_processor_config_path = pretrained_model_dir / "preprocessor_config.json"
  135. with open(image_processor_config_path, "r", encoding="utf-8") as f:
  136. image_processor_config = json.load(f)
  137. return cls(**image_processor_config)
  138. def _preprocess(
  139. self,
  140. images,
  141. do_resize: Optional[bool] = None,
  142. do_rescale: Optional[bool] = None,
  143. rescale_factor: Optional[float] = None,
  144. do_normalize: Optional[bool] = None,
  145. image_mean: Optional[Union[float, List[float]]] = None,
  146. image_std: Optional[Union[float, List[float]]] = None,
  147. do_convert_rgb: Optional[bool] = None,
  148. ):
  149. images = make_list_of_images(images)
  150. if do_convert_rgb:
  151. images = [convert_to_rgb(image) for image in images]
  152. width, height = images[0].size
  153. resized_height, resized_width = height, width
  154. processed_images = []
  155. for image in images:
  156. if do_resize:
  157. resized_height, resized_width = smart_resize(
  158. height,
  159. width,
  160. factor=self.patch_size * self.merge_size,
  161. min_pixels=self.min_pixels,
  162. max_pixels=self.max_pixels,
  163. )
  164. image = image.resize(
  165. (resized_width, resized_height), resample=self.resample
  166. )
  167. image = to_numpy_array(image)
  168. if do_rescale:
  169. image = (image * rescale_factor).astype(np.float32)
  170. if do_normalize:
  171. image = image.astype(np.float32)
  172. image -= np.array(image_mean, dtype=np.float32)
  173. image /= np.array(image_std, dtype=np.float32)
  174. processed_images.append(image)
  175. patches = np.array(processed_images)
  176. patches = patches.transpose(0, 3, 1, 2)
  177. if patches.shape[0] == 1:
  178. patches = np.tile(patches, (self.temporal_patch_size, 1, 1, 1))
  179. channel = patches.shape[1]
  180. grid_t = patches.shape[0] // self.temporal_patch_size
  181. grid_h, grid_w = (
  182. resized_height // self.patch_size,
  183. resized_width // self.patch_size,
  184. )
  185. patches = patches.reshape(
  186. grid_t,
  187. self.temporal_patch_size,
  188. channel,
  189. grid_h,
  190. self.patch_size,
  191. grid_w,
  192. self.patch_size,
  193. )
  194. patches = patches.transpose(0, 3, 5, 2, 1, 4, 6)
  195. assert self.temporal_patch_size == 1
  196. flatten_patches = patches.reshape(
  197. grid_t * grid_h * grid_w, channel, self.patch_size, self.patch_size
  198. )
  199. return flatten_patches, (grid_t, grid_h, grid_w)
  200. def __call__(
  201. self,
  202. images,
  203. videos=None,
  204. do_resize: Optional[bool] = None,
  205. size: Optional[Dict[str, int]] = None,
  206. do_rescale: Optional[bool] = None,
  207. rescale_factor: Optional[float] = None,
  208. do_normalize: Optional[bool] = None,
  209. image_mean: Optional[Union[float, List[float]]] = None,
  210. image_std: Optional[Union[float, List[float]]] = None,
  211. do_convert_rgb: Optional[bool] = None,
  212. return_tensors=None,
  213. ):
  214. do_resize = do_resize if do_resize is not None else self.do_resize
  215. size = size if size is not None else self.size
  216. do_rescale = do_rescale if do_rescale is not None else self.do_rescale
  217. rescale_factor = (
  218. rescale_factor if rescale_factor is not None else self.rescale_factor
  219. )
  220. do_normalize = do_normalize if do_normalize is not None else self.do_normalize
  221. image_mean = image_mean if image_mean is not None else self.image_mean
  222. image_std = image_std if image_std is not None else self.image_std
  223. do_convert_rgb = (
  224. do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
  225. )
  226. if images is not None:
  227. images = make_batched_images(images)
  228. if videos is not None:
  229. raise NotImplementedError("Videos are not yet supported")
  230. if images is not None:
  231. pixel_values, vision_grid_thws = [], []
  232. for image in images:
  233. patches, image_grid_thw = self._preprocess(
  234. image,
  235. do_resize=do_resize,
  236. do_rescale=do_rescale,
  237. rescale_factor=rescale_factor,
  238. do_normalize=do_normalize,
  239. image_mean=image_mean,
  240. image_std=image_std,
  241. do_convert_rgb=do_convert_rgb,
  242. )
  243. pixel_values.extend(patches)
  244. vision_grid_thws.append(image_grid_thw)
  245. pixel_values = np.array(pixel_values)
  246. vision_grid_thws = np.array(vision_grid_thws)
  247. data = {"pixel_values": pixel_values, "image_grid_thw": vision_grid_thws}
  248. return BatchFeature(data=data, tensor_type=return_tensors)