| | """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" |
| |
|
