| from __future__ import annotations |
|
|
| import base64 |
| import logging |
| import math |
| import os |
| import sys |
| import time |
| import warnings |
| from functools import lru_cache |
| from io import BytesIO |
| import io |
|
|
| import requests |
| import torch |
| import torchvision |
| from packaging import version |
| from PIL import Image |
| from torchvision import transforms |
| from torchvision.transforms import InterpolationMode |
| from typing import Optional |
| import random |
|
|
|
|
| logger = logging.getLogger(__name__) |
|
|
| IMAGE_FACTOR = 32 |
| MIN_PIXELS = 4 * 32 * 32 |
| MAX_PIXELS = 50 * 32 * 32 |
| MAX_RATIO = 200 |
|
|
| VIDEO_MIN_PIXELS = 128 * 28 * 28 |
| VIDEO_MAX_PIXELS = 768 * 28 * 28 |
| FRAME_FACTOR = 2 |
| FPS = 2.0 |
| FPS_MIN_FRAMES = 4 |
| FPS_MAX_FRAMES = 768 |
|
|
| VIDEO_TOTAL_PIXELS = int(float(os.environ.get('VIDEO_MAX_PIXELS', 128000 * 28 * 28 * 0.9))) |
|
|
|
|
| def round_by_factor(number: int, factor: int) -> int: |
| return round(number / factor) * factor |
|
|
|
|
| def ceil_by_factor(number: int, factor: int) -> int: |
| return math.ceil(number / factor) * factor |
|
|
|
|
| def floor_by_factor(number: int, factor: int) -> int: |
| return math.floor(number / factor) * factor |
|
|
|
|
| def smart_resize( |
| height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS |
| ) -> tuple[int, int]: |
| if max(height, width) / min(height, width) > MAX_RATIO: |
| raise ValueError( |
| f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" |
| ) |
| h_bar = max(factor, round_by_factor(height, factor)) |
| w_bar = max(factor, round_by_factor(width, factor)) |
| if h_bar * w_bar > max_pixels: |
| beta = math.sqrt((height * width) / max_pixels) |
| h_bar = floor_by_factor(height / beta, factor) |
| w_bar = floor_by_factor(width / beta, factor) |
| elif h_bar * w_bar < min_pixels: |
| beta = math.sqrt(min_pixels / (height * width)) |
| h_bar = ceil_by_factor(height * beta, factor) |
| w_bar = ceil_by_factor(width * beta, factor) |
| if h_bar / w_bar > MAX_RATIO: |
| h_bar = w_bar * MAX_RATIO |
| elif w_bar / h_bar > MAX_RATIO: |
| w_bar = h_bar * MAX_RATIO |
| return h_bar, w_bar |
|
|
|
|
| def to_rgb(pil_image: Image.Image) -> Image.Image: |
| if pil_image.mode == 'RGBA': |
| white_background = Image.new("RGB", pil_image.size, (255, 255, 255)) |
| white_background.paste(pil_image, mask=pil_image.split()[3]) |
| return white_background |
| else: |
| return pil_image.convert("RGB") |
|
|
|
|
| def resize_bbox(bbox, w, h, scale=None): |
| x1, y1, x2, y2 = bbox |
| orig_center_x = (x1 + x2) / 2 |
| orig_center_y = (y1 + y2) / 2 |
| orig_w = x2 - x1 |
| orig_h = y2 - y1 |
|
|
| max_scale_x = min(orig_center_x / max(orig_w / 2, 1e-6), (w - orig_center_x) / max(orig_w / 2, 1e-6)) |
| max_scale_y = min(orig_center_y / max(orig_h / 2, 1e-6), (h - orig_center_y) / max(orig_h / 2, 1e-6)) |
| max_scale = min(max_scale_x, max_scale_y) |
|
|
| if scale is not None: |
| scale = min(max_scale, scale) |
| else: |
| scale = 1.0 |
| if max_scale > 1.0: |
| scale = random.uniform(1.0, max_scale) |
|
|
| new_w = orig_w * scale |
| new_h = orig_h * scale |
|
|
| new_x1 = int(math.floor(orig_center_x - new_w / 2)) |
| new_y1 = int(math.floor(orig_center_y - new_h / 2)) |
| new_x2 = int(math.ceil(orig_center_x + new_w / 2)) |
| new_y2 = int(math.ceil(orig_center_y + new_h / 2)) |
|
|
| new_x1 = max(0, new_x1) |
| new_y1 = max(0, new_y1) |
| new_x2 = min(w, new_x2) |
| new_y2 = min(h, new_y2) |
|
|
| if new_x1 >= new_x2 or new_y1 >= new_y2 or x2 > w or y2 > h: |
| return bbox |
|
|
| return [new_x1, new_y1, new_x2, new_y2] |
|
|
|
|
| def fetch_image(ele, size_factor: int = IMAGE_FACTOR) -> Image.Image: |
| if "image" in ele: |
| image = ele["image"] |
| else: |
| image = ele["image_url"] |
| image_obj = None |
| if isinstance(image, Image.Image): |
| image_obj = image |
| elif image.startswith("http://") or image.startswith("https://"): |
| response = requests.get(image, stream=True) |
| image_obj = Image.open(BytesIO(response.content)) |
| elif image.startswith("file://"): |
| image_obj = Image.open(image[7:]) |
| elif image.startswith("data:image"): |
| if "base64," in image: |
| _, base64_data = image.split("base64,", 1) |
| data = base64.b64decode(base64_data) |
| image_obj = Image.open(BytesIO(data)) |
| else: |
| image_obj = Image.open(image) |
|
|
| bbox = ele.get("bbox", None) |
| if bbox: |
| w, h = image_obj.width, image_obj.height |
| if len(bbox) == 6: |
| x1_prop, y1_prop, x2_prop, y2_prop, w_bar, h_bar = bbox |
| x1 = int(x1_prop / w_bar * w) |
| y1 = int(y1_prop / h_bar * h) |
| x2 = int(x2_prop / w_bar * w) |
| y2 = int(y2_prop / h_bar * h) |
| bbox = [x1, y1, x2, y2] |
|
|
| if ele.get("dynamic_crop", False): |
| bbox = resize_bbox(bbox, w, h) |
|
|
| image_obj = image_obj.crop(bbox) |
|
|
| if image_obj is None: |
| raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}") |
| image = to_rgb(image_obj) |
|
|
| if "hflip" in ele and ele["hflip"]: |
| image = image.transpose(Image.FLIP_LEFT_RIGHT) |
| if "rotate" in ele and ele["rotate"]: |
| angle = random.choice([-15, -10, -5, 5, 10, 15]) |
| image = image.rotate(angle) |
|
|
| if "resized_height" in ele and "resized_width" in ele: |
| resized_height, resized_width = smart_resize( |
| ele["resized_height"], |
| ele["resized_width"], |
| factor=size_factor, |
| ) |
| else: |
| width, height = image.size |
| min_pixels = ele.get("min_pixels", MIN_PIXELS) |
| max_pixels = ele.get("max_pixels", MAX_PIXELS) |
| if max_pixels is None: |
| max_pixels = MAX_PIXELS |
| resized_height, resized_width = smart_resize( |
| height, |
| width, |
| factor=size_factor, |
| min_pixels=min_pixels, |
| max_pixels=max_pixels, |
| ) |
| image = image.resize((resized_width, resized_height)) |
|
|
| return image |
|
|
|
|
| def smart_nframes(ele: dict, total_frames: int, video_fps: int | float) -> int: |
| assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`" |
| if "nframes" in ele: |
| nframes = round_by_factor(ele["nframes"], FRAME_FACTOR) |
| else: |
| fps = ele.get("fps", FPS) |
| min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR) |
| max_frames = floor_by_factor(ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR) |
| nframes = total_frames / video_fps * fps |
| if nframes > total_frames: |
| logger.warning(f"smart_nframes: nframes[{nframes}] > total_frames[{total_frames}]") |
| nframes = min(min(max(nframes, min_frames), max_frames), total_frames) |
| nframes = floor_by_factor(nframes, FRAME_FACTOR) |
| if not (FRAME_FACTOR <= nframes and nframes <= total_frames): |
| raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.") |
| return nframes |
|
|
|
|
| def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR, return_video_sample_fps: bool = False): |
| if isinstance(ele["video"], str): |
| video, audio, info = torchvision.io.read_video( |
| ele["video"], |
| start_pts=ele.get("video_start", 0.0), |
| end_pts=ele.get("video_end", None), |
| pts_unit="sec", |
| output_format="TCHW", |
| ) |
| total_frames, video_fps = video.size(0), info["video_fps"] |
| nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) |
| idx = torch.linspace(0, total_frames - 1, nframes).round().long() |
| sample_fps = nframes / max(total_frames, 1e-6) * video_fps |
| video = video[idx] |
|
|
| nframes, _, height, width = video.shape |
| min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS) |
| total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS) |
| max_pixels = max(min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05)) |
| max_pixels_supposed = ele.get("max_pixels", max_pixels) |
| max_pixels = min(max_pixels_supposed, max_pixels) |
| if "resized_height" in ele and "resized_width" in ele: |
| resized_height, resized_width = smart_resize(ele["resized_height"], ele["resized_width"], factor=image_factor) |
| else: |
| resized_height, resized_width = smart_resize(height, width, factor=image_factor, min_pixels=min_pixels, max_pixels=max_pixels) |
| video = transforms.functional.resize(video, [resized_height, resized_width], interpolation=InterpolationMode.BICUBIC, antialias=True).float() |
| if return_video_sample_fps: |
| return video, sample_fps |
| return video |
| else: |
| assert isinstance(ele["video"], (list, tuple)) |
| process_info = ele.copy() |
| process_info.pop("type", None) |
| process_info.pop("video", None) |
| images = [fetch_image({"image": video_element, **process_info}, size_factor=image_factor) for video_element in ele["video"]] |
| nframes = ceil_by_factor(len(images), FRAME_FACTOR) |
| if len(images) < nframes: |
| images.extend([images[-1]] * (nframes - len(images))) |
| if return_video_sample_fps: |
| return images, process_info.pop("fps", 2.0) |
| return images |
|
|
|
|
| def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]: |
| vision_infos = [] |
| if isinstance(conversations[0], dict): |
| conversations = [conversations] |
| for conversation in conversations: |
| for message in conversation: |
| if isinstance(message["content"], list): |
| for ele in message["content"]: |
| if ( |
| "image" in ele |
| or "image_url" in ele |
| or "video" in ele |
| or ele["type"] in ("image", "image_url", "video") |
| ): |
| vision_infos.append(ele) |
| return vision_infos |
|
|
|
|
| def process_vision_info( |
| conversations: list[dict] | list[list[dict]], |
| return_video_kwargs: bool = False, |
| ): |
| vision_infos = extract_vision_info(conversations) |
| image_inputs = [] |
| video_inputs = [] |
| video_sample_fps_list = [] |
| for vision_info in vision_infos: |
| if "image" in vision_info or "image_url" in vision_info: |
| image_inputs.append(fetch_image(vision_info)) |
| elif "video" in vision_info: |
| video_input, video_sample_fps = fetch_video(vision_info, return_video_sample_fps=True) |
| video_sample_fps_list.append(video_sample_fps) |
| video_inputs.append(video_input) |
| else: |
| raise ValueError("image, image_url or video should in content.") |
| if len(image_inputs) == 0: |
| image_inputs = None |
| if len(video_inputs) == 0: |
| video_inputs = None |
| if return_video_kwargs: |
| return image_inputs, video_inputs, {'fps': video_sample_fps_list} |
| return image_inputs, video_inputs |
|
|
