| """Image processor class for Kimi-K2.5. |
| """ |
|
|
| import json |
| from typing import Any, Dict, Optional, Union |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
| from transformers.image_processing_utils import (BaseImageProcessor, |
| BatchFeature) |
| from transformers.utils import TensorType |
|
|
| from .media_utils import (MediaInput, VideoChunkInput, _to_tensor, |
| ensure_media_type, get_video_meta, image_to_np, |
| navit_patchify, navit_resize_image, |
| navit_resize_video, normalize, |
| real_sample_fps_and_max_num_frames, timestamp_as_str) |
|
|
| try: |
| from mecord import VideoReader |
| except ImportError: |
| VideoReader = None |
|
|
|
|
| def resampling(video_bytes: bytes, |
| sample_indices: list[int], |
| key_indices=None, |
| frame_time_info=None, |
| num_threads=4) -> str: |
| video = VideoReader(video_bytes, |
| num_threads=num_threads, |
| frame_time_info=frame_time_info, |
| key_indices=key_indices) |
| |
| frames = video[sample_indices] |
| frames = [Image.fromarray(frame) for frame in frames] |
| return frames |
|
|
|
|
| class KimiK25VisionProcessor(BaseImageProcessor): |
| model_type = "kimi_k25" |
|
|
| def __init__( |
| self, |
| media_proc_cfg: dict, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.media_proc_cfg = media_proc_cfg |
| self.num_frames_per_chunk = media_proc_cfg[ |
| 'temporal_merge_kernel_size'] |
|
|
| def media_tokens_calculator(self, media: MediaInput): |
| media = ensure_media_type(media) |
| ret = self.get_resize_config(media) |
| return ret['num_tokens'] |
|
|
| @classmethod |
| def make_chunk_prompt(cls, timestamp_text: str) -> str: |
| return f"{timestamp_text}<|media_begin|>video<|media_content|><|media_pad|><|media_end|>" |
|
|
| def split_video_chunks(self, |
| video_url: str | bytes) -> list[list[Image.Image]]: |
| |
| video_spec = get_video_meta(video_url) |
| sample_fps = min(self.media_proc_cfg['sample_fps'], video_spec.fps) |
| sampled_nframes = max( |
| round(video_spec.num_frames * sample_fps / video_spec.fps), 1) |
| frame_inds = np.linspace(0, video_spec.num_frames - 1, |
| sampled_nframes).round().astype(int) |
| frame_inds = frame_inds.tolist() |
| sampled_frame_ids = [] |
| temporal_merge_kernel_size = self.media_proc_cfg[ |
| "temporal_merge_kernel_size"] |
| num_chunks = 0 |
| chunk_timestamp = [] |
| for i in range(0, len(frame_inds), temporal_merge_kernel_size): |
| sampled_frame_ids.extend(frame_inds[i:i + |
| temporal_merge_kernel_size]) |
| start_time = frame_inds[i] / float(video_spec.fps) |
| timestamp_text = timestamp_as_str( |
| start_time, self.media_proc_cfg["timestamp_mode"]) |
| chunk_timestamp.append(timestamp_text) |
| num_chunks += 1 |
|
|
| sampled_frames = resampling(video_url, sampled_frame_ids) |
| chunks = [] |
| for chunk_id in range(num_chunks): |
| chunk = sampled_frames[chunk_id * |
| temporal_merge_kernel_size:(chunk_id + 1) * |
| temporal_merge_kernel_size] |
| chunks.append( |
| VideoChunkInput(type="video_chunk", |
| video_chunk=chunk, |
| prompt=self.make_chunk_prompt( |
| chunk_timestamp[chunk_id]))) |
| return chunks |
|
|
| def get_resize_config(self, media_input: MediaInput) -> dict: |
| if media_input['type'] == 'image': |
| w, h = media_input['image'].size |
| ret = navit_resize_image( |
| w, h, self.media_proc_cfg['patch_size'], |
| self.media_proc_cfg['merge_kernel_size'], |
| self.media_proc_cfg['in_patch_limit'], |
| self.media_proc_cfg['patch_limit_on_one_side'], |
| self.media_proc_cfg['fixed_output_tokens']) |
| return ret |
| elif media_input['type'] == 'video_chunk': |
| frame = media_input['video_chunk'][0] |
| width, height = frame.size |
| num_frames = len(media_input["video_chunk"]) |
| fps = 1.0 |
|
|
| sample_fps, max_num_frames_each_video = real_sample_fps_and_max_num_frames( |
| media_input["type"], |
| self.media_proc_cfg['sample_fps'], |
| self.media_proc_cfg['max_num_frames_each_video'], |
| ) |
|
|
| in_patch_limit_each_frame = self.media_proc_cfg[ |
| 'in_patch_limit_each_frame'] |
| if in_patch_limit_each_frame is None: |
| in_patch_limit_each_frame = self.media_proc_cfg[ |
| 'in_patch_limit'] |
|
|
| ret = navit_resize_video( |
| width, |
| height, |
| num_frames, |
| fps, |
| sample_fps, |
| self.media_proc_cfg['patch_size'], |
| self.media_proc_cfg['merge_kernel_size'], |
| in_patch_limit_each_frame, |
| self.media_proc_cfg['patch_limit_on_one_side'], |
| self.media_proc_cfg['in_patch_limit_video'], |
| max_num_frames_each_video, |
| self.media_proc_cfg['fixed_output_tokens'], |
| ) |
| return ret |
| else: |
| raise ValueError("Unsupported type: {}".format( |
| media_input['type'])) |
|
|
| def resize_image(self, image: Image.Image, new_width: int, new_height: int, |
| pad_width: int, pad_height: int) -> np.ndarray: |
| image_np = image_to_np(image, (new_width, new_height), "resize") |
| image_np = np.pad( |
| image_np, |
| ((0, pad_height), (0, pad_width), (0, 0)), |
| mode="constant", |
| constant_values=0, |
| ) |
| return image_np |
|
|
| def preprocess( |
| self, |
| medias: list[MediaInput], |
| return_tensors: Optional[Union[str, TensorType]] = None, |
| ) -> BatchFeature: |
| """ |
| Preprocess a atom vision input (images/video_chunk) into model-ready tensors. |
| |
| Args: |
| medias: List of MediaInput. |
| return_tensors: Desired output format ('pt', 'np', 'tf', or None). |
| |
| Returns: |
| BatchFeature containing 'pixel_values' and 'grid_thws' tensors. |
| """ |
| if not isinstance(medias, list): |
| medias = [medias] |
| if medias: |
| pixel_values = [] |
| for item in medias: |
| item = ensure_media_type(item) |
| resize_config = self.get_resize_config(item) |
| new_width, new_height, pad_width, pad_height = resize_config[ |
| 'new_width'], resize_config['new_height'], resize_config[ |
| 'pad_width'], resize_config['pad_height'] |
| if item['type'] == 'image': |
| image = item['image'] |
| image_np = self.resize_image(image, new_width, new_height, |
| pad_width, pad_height) |
| pixel_values.append(np.expand_dims(image_np, axis=0)) |
| elif item['type'] == 'video_chunk': |
| pixels = [] |
| for frame in item['video_chunk']: |
| frame_np = self.resize_image(frame, new_width, |
| new_height, pad_width, |
| pad_height) |
| pixels.append(frame_np) |
| pixel_values.append(np.stack(pixels, axis=0)) |
| else: |
| raise ValueError("Unsupported type: {}".format( |
| item['type'])) |
| normalized_pixel_values = [] |
| image_std_inv = 1.0 / np.array(self.media_proc_cfg['image_std']) |
| image_mean = np.array(self.media_proc_cfg['image_mean']) |
| for pixels in pixel_values: |
| pixels = normalize(pixels, image_mean, image_std_inv) |
| pixels_and_thw = navit_patchify( |
| pixels, |
| self.media_proc_cfg['patch_size'], |
| ) |
| normalized_pixel_values.append(pixels_and_thw) |
|
|
| pixel_values = torch.cat([ |
| _to_tensor(pixel_value['pixel_values']) |
| for pixel_value in normalized_pixel_values |
| ]) |
| grid_thws = torch.cat([ |
| _to_tensor(pixel_value['grid_thw'], |
| dtype=torch.int64).unsqueeze(0) |
| for pixel_value in normalized_pixel_values |
| ]) |
|
|
| data = { |
| 'pixel_values': pixel_values, |
| 'grid_thws': grid_thws, |
| } |
|
|
| else: |
| data = {} |
|
|
| return BatchFeature(data=data, tensor_type=return_tensors) |
|
|
| def __repr__(self): |
| return f"KimiK25VisionProcessor(media_proc_cfg={self.media_proc_cfg})" |
|
|
| def to_dict(self) -> Dict[str, Any]: |
| output = super().to_dict() |
| output["media_proc_cfg"] = self.media_proc_cfg |
| if "media_processor" in output: |
| del output["media_processor"] |
| return output |
|
|
| @classmethod |
| def from_dict(cls, config_dict: Dict[str, Any], **kwargs): |
| config = config_dict.copy() |
| media_proc_cfg = config.pop("media_proc_cfg", {}) |
| return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs) |
|
|
| def to_json_string(self): |
| dictionary = self.to_dict() |
| for key, value in dictionary.items(): |
| if hasattr(value, 'tolist'): |
| dictionary[key] = value.tolist() |
| return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" |
|
|
