image_processing_mplugowl3.py

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)