417 lines
16 KiB
Python
417 lines
16 KiB
Python
|
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)
|