import random
from typing import Optional, Union, Dict, Any, List
from einops import rearrange, repeat
import torch
import math
import PIL.Image
import PIL.ImageSequence
import numpy as np
import PIL
from PIL import Image
from transformers.utils import TensorType, requires_backends, is_torch_dtype, is_torch_device
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers import AutoImageProcessor
from transformers.image_transforms import to_channel_dimension_format
from transformers.image_utils import (
ImageInput,
make_list_of_images,
valid_images,
is_torch_tensor,
is_batched,
to_numpy_array,
infer_channel_dimension_format,
ChannelDimension
)
from torchvision.ops.boxes import box_area
from torchvision.transforms import functional as F
from torchvision.transforms.transforms import InterpolationMode
from torchvision import transforms
def recursive_converter(converter, value):
if isinstance(value, list):
new_value = []
for v in value:
new_value += [recursive_converter(converter, v)]
return new_value
else:
return converter(value)
def box_iou(boxes1, area1, boxes2, eps=1e-5):
area2 = box_area(boxes2)
lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]
rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]
wh = (rb - lt).clamp(min=0) # [N,M,2]
inter = wh[:, :, 0] * wh[:, :, 1] # [N,M]
union = area1[:, None] + area2 - inter
iou = inter / (union+eps)
return iou, union
available_anchor_strategy = ['docowl', 'random', 'highest', 'last', 'llava']
grid_dict = {
'grid_33':[
(1,1),
(1,2),(2,1),
(1,3),(3,1),
(2,2),(1,4),(4,1),
(1,5),(5,1),
(1,6),(6,1),(2,3),(3,2),
(1,7),(7,1),
(4,2),(2,4),(1,8),(8,1),
(3,3),(1,9),(9,1)],
'grid_squ_3x3':[
(1,1),(2,2),(3,3)
],
'grid_squ_4':[
(2,2),(1,3),(1,4),(3,1),(4,1)
],
'grid_squ_6':[
(2,2),(1,3),(1,4),(3,1),(4,1), (2,3),(3,2)
],
'grid_squ_2':[
(2,1)
],
'grid_squ_9':[
(1,1),
(1,2),(2,1),
(1,3),(3,1),
(2,2),(1,4),(4,1),
(1,5),(5,1),
(1,6),(6,1),(2,3),(3,2),
(1,7),(7,1),
(4,2),(2,4),(1,8),(8,1),
(3,3),(1,9),(9,1)],
}
cut_prompt_template_dict = {
'v0': lambda img_token, h, w: f''.join([f"{img_token}" for i in range(h) for j in range(w)]),
'v1': lambda img_token, h, w: f'Cut to {h} rows {w} columns, '+ ' '.join([f"subimg({i},{j}){img_token}"for i in range(h) for j in range(w)]),
'v1_global': lambda img_token, h, w: f'Cut to {h} rows {w} columns with a global view, '+ ' '.join([f"subimg({i},{j}){img_token}"for i in range(h) for j in range(w)]+[f"global_view{img_token}"]),
'v2_global': lambda img_token, h, w: f'Cut to {h} rows {w} columns with a global view\n'+ '\n'.join([' '.join([f"subimg({i},{j}){img_token}" for j in range(w)]) for i in range(h)])+f"\nglobal_view{img_token}",
'v3': lambda img_token, h, w: f'<|start_cut|>{h}*{w}'+ ' '.join([f"{img_token}"for i in range(h) for j in range(w)])+'<|end_cut|>',
'v3_global': lambda img_token, h, w: f'<|start_cut|>{h}*{w}\n'+ '\n'.join([' '.join([f"{img_token}" for j in range(w)]) for i in range(h)])+f'\n{img_token}<|end_cut|>',
}
def anchor_rank(anchors, anchors_areas, input_image_size, eps=1e-5):
# anchors x1 y1 x2 y2
# image_size: (h, w)
# xyxy
input_image_bbox = torch.tensor([0, 0, input_image_size[1], input_image_size[0]]).unsqueeze(0)
boxes1 = anchors
boxes2 = input_image_bbox
boxes3 = anchors.clone()
# y2
boxes3[:,3] = input_image_size[0]/input_image_size[1]*anchors[:,2] # 用于算分辨率无关的iou
area1 = anchors_areas
iou, _ = box_iou(boxes1, area1, boxes2)
iou = iou.squeeze(1)
shape_iou, _ = box_iou(boxes1, area1, boxes3)
shape_iou = shape_iou.diag()
# 优先匹配形状接近 再匹配分辨率接近
index = torch.argmax(shape_iou*100+iou,dim=0)
return index
def select_best_resolution(anchors, anchors_areas, input_image_size): # TODO For a futher check
"""
Selects the best resolution from a list of possible resolutions based on the original size.
Args:
original_size (tuple): The original size of the image in the format (width, height).
possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), (width2, height2), ...].
Returns:
tuple: The best fit resolution in the format (width, height).
"""
original_size = (input_image_size[1], input_image_size[0])
possible_resolutions = [(_[2], _[3]) for _ in anchors] # xyxy -> w,h
original_width, original_height = original_size
best_fit = None
max_effective_resolution = 0
min_wasted_resolution = float('inf')
index = 0
for i, (width, height) in enumerate(possible_resolutions):
scale = min(width / original_width, height / original_height)
downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale)
effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height)
wasted_resolution = (width * height) - effective_resolution
if effective_resolution > max_effective_resolution or (effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution):
max_effective_resolution = effective_resolution
min_wasted_resolution = wasted_resolution
best_fit = (width, height)
index = i
return index
def build_cut_shape_indices(cut_shape):
# cut_shape: a list of (nh,nw)
cut_shape_indices = []
for shape in cut_shape:
n=shape[0]*shape[1]
indices = torch.cat([
repeat(torch.tensor(shape),'l -> n l',n=n),
torch.arange(n).unsqueeze(1)
], dim=1)
assert indices.shape[0] == n
assert indices.shape[1] == 3 # nh,nw,idx
cut_shape_indices.append(indices)
cut_shape_indices = torch.cat(cut_shape_indices,dim=0).long()
return cut_shape_indices
class AnchorResize(torch.nn.Module):
def __init__(self, image_size, anchors, interpolation=InterpolationMode.BILINEAR, antialias=None, anchor_strategy='docowl'):
super().__init__()
self.image_size = image_size
# xyxy
self.anchors = torch.tensor(
[[0, 0, _[1]*image_size[1], _[0]*image_size[0]]
for _ in anchors], requires_grad=False
)
self.anchor_areas = box_area(self.anchors)
self.interpolation = interpolation
self.antialias = antialias
self.anchor_strategy = anchor_strategy
assert self.anchor_strategy in available_anchor_strategy
def resize_global(self, img):
return F.resize(img, self.image_size, self.interpolation, max_size=None, antialias=self.antialias)
def forward(self, img, skip_resize=False):
"""
Args:
img (PIL Image or Tensor): Image to be scaled.
Returns:
PIL Image or Tensor: Rescaled image.
"""
if self.anchor_strategy == 'docowl':
selected_anchor = anchor_rank(self.anchors, self.anchor_areas, (img.size[1], img.size[0]))
elif self.anchor_strategy == 'random':
selected_anchor = random.randint(0,len(self.anchors)-1)
elif self.anchor_strategy == 'highest':
# 选面积最大的 在这个基础上 尽可能选最方正的
selected_anchor = torch.argmax(self.anchors[:,2]*self.anchors[:,3]*100-torch.abs(self.anchors[:,2]-self.anchors[:,3]))
elif self.anchor_strategy == 'last':
selected_anchor = len(self.anchors)-1
elif self.anchor_strategy == 'llava':
selected_anchor = select_best_resolution(self.anchors, self.anchor_areas, (img.size[1], img.size[0]))
else:
selected_anchor = None
assert selected_anchor is not None
target_size = self.anchors[selected_anchor][2:].tolist() # w,h
if skip_resize:
# for debug
return selected_anchor
return F.resize(img, [target_size[1],target_size[0]], self.interpolation, max_size=None, antialias=self.antialias), selected_anchor
def __repr__(self) -> str:
detail = f"(size={self.image_size}, anchor={self.anchors}, interpolation={self.interpolation.value}, antialias={self.antialias})"
return f"{self.__class__.__name__}{detail}"
class CutMixin:
def __init__(self, cut_cfg={"anchors": "grid_squ_6", "anchor_strategy": "docowl", "cut_prompt": "v3", "add_global": True, "cut_prob": 1.0}) -> None:
if cut_cfg is None:
self.cut_enable = False
return
else:
self.cut_enable = True
image_size = self.image_size
anchors = cut_cfg.get('anchors','grid_33')
anchor_strategy = cut_cfg.get('anchor_strategy','docowl')
cut_prompt = cut_cfg.get('cut_prompt','v0')
self.cut_prob = cut_cfg.get('cut_prob', 1.0)
self.force_shape_cut = cut_cfg.get('force_shape_cut', False)
force_shape_cut_anchors = cut_cfg.get('force_shape_cut_anchors', 'force_shape_cut_anchors')
self.add_global = cut_cfg.get('add_global', False)
# h,w
if isinstance(image_size, int):
image_size = (image_size, image_size)
self.image_size = image_size
if anchors in grid_dict:
anchors = grid_dict[anchors]
else:
anchors = eval(anchors)
self.anchors = [tuple(_) for _ in anchors]
self.anchor_max = max([max(_) for _ in self.anchors])
self.resizer = AnchorResize(image_size=image_size, anchors=anchors, interpolation=InterpolationMode.BICUBIC, anchor_strategy=anchor_strategy)
if force_shape_cut_anchors in grid_dict:
force_shape_cut_anchors = grid_dict[force_shape_cut_anchors]
else:
force_shape_cut_anchors = eval(force_shape_cut_anchors)
self.force_shape_cut_anchors = [tuple(_) for _ in force_shape_cut_anchors]
self.force_shape_cut_anchors_max = max([max(_) for _ in self.force_shape_cut_anchors])
self.old_resizer = transforms.Resize(image_size,interpolation=InterpolationMode.BICUBIC)
# 把image processor的缩放去掉 只保留后面的变换
self.image_transform = transforms.Compose(self.image_transform.transforms[1:])
if self.add_global:
self.cut_prompt_template = cut_prompt_template_dict[cut_prompt+'_global']
else:
self.cut_prompt_template = cut_prompt_template_dict[cut_prompt]
self.media_tokens = ["<|image|>", "<|video|>"]
def _process_image(self, images):
new_images = []
cut_shape = []
for image in images:
raw_image = image
image, selected_anchor = self.resizer(image)
image_input = self.image_transform(image) # h,w,3 -> 3,h,w
cut_shape.append((image_input.shape[1]//self.image_size[0], image_input.shape[2]//self.image_size[1])) # cut_h, cut_w
image_input = rearrange(image_input, 'C (num_h h) (num_w w) -> (num_h num_w) C h w', h=self.image_size[0], w=self.image_size[1])
new_images.append(image_input)
if self.add_global:
new_images.append(self.image_transform(self.resizer.resize_global(raw_image)).unsqueeze(0))
cut_shape.append((1,1))
new_images = torch.cat(new_images,dim=0)
cut_shape_indices = build_cut_shape_indices(cut_shape)
return new_images, cut_shape, cut_shape_indices
class mPLUGOwl3BatchFeature(BatchFeature):
r"""
Extend from BatchFeature for supporting various image size
"""
def __init__(self, data: Optional[Dict[str, Any]] = None, tensor_type: Union[None, str, TensorType] = None):
super().__init__(data)
self.convert_to_tensors(tensor_type=tensor_type)
def convert_to_tensors(self, tensor_type: Optional[Union[str, TensorType]] = None):
if tensor_type is None:
return self
is_tensor, as_tensor = self._get_is_as_tensor_fns(tensor_type)
def converter(value):
try:
if not is_tensor(value):
tensor = as_tensor(value)
return tensor
except: # noqa E722
if key == "overflowing_values":
raise ValueError("Unable to create tensor returning overflowing values of different lengths. ")
raise ValueError(
"Unable to create tensor, you should probably activate padding "
"with 'padding=True' to have batched tensors with the same length."
)
for key, value in self.items():
self[key] = recursive_converter(converter, value)
return self
def to(self, *args, **kwargs) -> "mPLUGOwl3BatchFeature":
requires_backends(self, ["torch"])
import torch
def cast_tensor(v):
# check if v is a floating point
if torch.is_floating_point(v):
# cast and send to device
return v.to(*args, **kwargs)
elif device is not None:
return v.to(device=device)
else:
return v
new_data = {}
device = kwargs.get("device")
# Check if the args are a device or a dtype
if device is None and len(args) > 0:
# device should be always the first argument
arg = args[0]
if is_torch_dtype(arg):
# The first argument is a dtype
pass
elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):
device = arg
else:
# it's something else
raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")
# We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`
for k, v in self.items():
new_data[k] = recursive_converter(cast_tensor, v)
self.data = new_data
return self
class mPLUGOwl3ImageProcessor(BaseImageProcessor, CutMixin):
model_input_names = ["pixel_values"]
def __init__(
self,
image_size,
mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5],
**kwargs):
super().__init__(**kwargs)
self.image_size = image_size
self.image_transform = transforms.Compose([
transforms.Resize((image_size, image_size), interpolation=Image.BICUBIC),
transforms.ToTensor(),
transforms.Normalize(mean, std),
])
CutMixin.__init__(self)
def preprocess(
self,
images: Union[Image.Image, List[Image.Image]],
cut_enable=True,
**kwargs
) -> mPLUGOwl3BatchFeature:
if isinstance(images, Image.Image):
images_list = [images]
else:
images_list = images
if self.cut_enable and cut_enable:
image_data, cut_shape, cut_shape_indices = self._process_image(images_list)
else:
image_data = [self.image_transform(self.resizer.resize_global(image)) for image in images_list]
image_data = torch.stack(image_data, dim=0)
cut_shape = cut_shape_indices = None
return mPLUGOwl3BatchFeature(data={'pixel_values': image_data, 'cut_shape':cut_shape, 'cut_shape_indices':cut_shape_indices})
def to_dict(self):
encoder_dict = super().to_dict()
pop_keys = ['image_transform', 'resizer', 'old_resizer', 'cut_prompt_template']
for pk in pop_keys:
encoder_dict.pop(pk, None)
return encoder_dict
AutoImageProcessor.register("mPLUGOwl3ImageProcessor", mPLUGOwl3ImageProcessor)