"""
Processor class for Molmo.
"""
from typing import Optional
import PIL
from PIL import Image
try:
from typing import Unpack
except ImportError:
from typing_extensions import Unpack
import re
from typing import List, Optional, Union
import numpy as np
import torch
import torchvision.transforms.functional as F
from transformers import AutoTokenizer
from transformers.image_utils import ImageInput
from transformers.processing_utils import (ProcessingKwargs, ProcessorMixin,
TextKwargs)
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
from transformers.utils import logging
logger = logging.get_logger(__name__)
IGNORE_INDEX = -100
DEFAULT_PAD_TOKEN_INDEX = 0
IMAGE_TOKEN_INDEX = -200
DEFAULT_IMAGE_TOKEN = "<image>"
# For Objects
DEFAULT_OBJECT_TOKEN = "<obj<i>>"
DEFAULT_OBJECT_FEATURE_TOKEN = "<objfeat>"
DEFAULT_OBJECT_INDEX = -300
# For Grounding
DEFAULT_GROUNDING_START = "<ground>"
DEFAULT_GROUNDING_END = "</ground>"
DEFAULT_GROUNDING_OBJECTS_START = "<objects>"
DEFAULT_GROUNDING_OBJECTS_END = "</objects>"
def xyxy_to_xywh(boxes):
"""
Convert boxes from xywh to xyxy format.
Parameters:
boxes (numpy.ndarray): An array of shape (N, 4) where N is the number of boxes.
Each box is represented as [x, y, x, y].
Returns:
numpy.ndarray: An array of shape (N, 4) where each box is represented as [x_min, y_min, w, h].
"""
boxes = np.array(boxes)
x_min, y_min, x_max, y_max = (
boxes[:, 0],
boxes[:, 1],
boxes[:, 2],
boxes[:, 3],
)
w = x_max - x_min
h = y_max - y_min
return np.stack([x_min, y_min, w, h], axis=1)
def xywh_to_xyxy(boxes):
"""
Convert boxes from xywh to xyxy format.
Parameters:
boxes (numpy.ndarray): An array of shape (N, 4) where N is the number of boxes.
Each box is represented as [x, y, width, height].
Returns:
numpy.ndarray: An array of shape (N, 4) where each box is represented as [x_min, y_min, x_max, y_max].
"""
boxes = np.array(boxes)
x, y, width, height = (
boxes[:, 0],
boxes[:, 1],
boxes[:, 2],
boxes[:, 3],
)
x_max = x + width
y_max = y + height
return np.stack([x, y, x_max, y_max], axis=1)
def expand2square(pil_img, background_color):
width, height = pil_img.size
if width == height:
return pil_img
elif width > height:
result = Image.new(pil_img.mode, (width, width), background_color)
result.paste(pil_img, (0, (width - height) // 2))
return result
else:
result = Image.new(pil_img.mode, (height, height), background_color)
result.paste(pil_img, ((height - width) // 2, 0))
return result
def pad_boxes(gt_boxes, old_size):
old_w, old_h = old_size
gt_boxes = np.array(gt_boxes).astype(np.float32)
# Calculate the padding added
if old_w > old_h:
pad_top = (old_w - old_h) // 2
pad_bottom = old_w - old_h - pad_top
pad_left, pad_right = 0, 0
else:
pad_left = (old_h - old_w) // 2
pad_right = old_h - old_w - pad_left
pad_top, pad_bottom = 0, 0
# Adjust the boxes for padding
gt_boxes[:, 0] += pad_left # x
gt_boxes[:, 1] += pad_top # y
return gt_boxes
def resize_boxes(gt_boxes, old_size, new_size):
old_w, old_h = old_size
new_h, new_w = new_size
gt_boxes = np.array(gt_boxes).astype(np.float32)
# Calculate scale factors
scale_x = new_w / max(old_w, old_h)
scale_y = new_h / max(old_w, old_h)
# Resize the boxes
gt_boxes[:, 0] *= scale_x # x
gt_boxes[:, 1] *= scale_y # y
gt_boxes[:, 2] *= scale_x # w
gt_boxes[:, 3] *= scale_y # h
return gt_boxes
def split_special_strings(input_string: str, special_strings: list[str] = None):
"""Split the input string into a list of strings, keeping the special strings.
Args:
input_string (str): The input string to split.
Example:
input_string = "<image>\n<obj0><objfeat><obj1><objfeat>\n I am happy today."
output = ['<image>', '\n<obj0>', '<objfeat>', '<obj1>', '<objfeat>', '\n I am happy today.']
Returns:
list: A list of strings, with the special strings separated from the rest of the input string.
"""
# Create a regex pattern to match the special strings
pattern = "|".join(map(re.escape, special_strings))
# Split the input string using the pattern, keeping the special strings in the result
split_list = re.split(f"({pattern})", input_string)
# Remove empty strings from the list
split_list = [s for s in split_list if s]
return split_list
def tokenizer_image_object_token(prompt, tokenizer):
bos_token_id = tokenizer.bos_token_id
split_tokens = [DEFAULT_IMAGE_TOKEN, DEFAULT_OBJECT_FEATURE_TOKEN]
chunks = split_special_strings(prompt, split_tokens)
input_encode = [bos_token_id]
for chunk in chunks:
if chunk == DEFAULT_IMAGE_TOKEN:
input_encode.append(IMAGE_TOKEN_INDEX)
elif chunk == DEFAULT_OBJECT_FEATURE_TOKEN:
input_encode.append(DEFAULT_OBJECT_INDEX)
else:
input_encode.extend(tokenizer.encode(chunk, add_special_tokens=False))
return input_encode
class ChatRexProcessor(ProcessorMixin):
attributes = ["image_processor", "tokenizer"]
image_processor_class = "AutoImageProcessor"
tokenizer_class = "AutoTokenizer"
def __init__(self, image_processor = None, tokenizer : AutoTokenizer = None, **kwargs):
# self.image_processor = image_processor
# self.tokenizer = tokenizer
super().__init__(image_processor, tokenizer)
self._special_tokens = None
self.template = dict(
SYSTEM=('A chat between a curious user and an artificial '
'intelligence assistant. The assistant gives '
'helpful, detailed, and polite answers to the '
'user\'s questions. {system}\n '),
INSTRUCTION=('USER: {input} ASSISTANT:'),
SEP='\n')
def process(
self,
image: Union[str, Image.Image],
bbox: List[List[int]],
question: str,
):
"""Prepare input data for inference.
Args:
image (Union[str, Image.Image]): The image to process.
bbox (List[List[int]]): A list of bounding boxes for the image. Each bounding box should
be in order of [x, y, x , y].
question (str): The question to ask about the image.
"""
data_dict = {}
# step1 load image
if type(image) == str:
image = Image.open(image).convert("RGB")
ori_w, ori_h = F.get_image_size(image)
image = expand2square(
image,
tuple(int(x * 255) for x in self.image_processor.image_mean),
)
pad_w, pad_h = F.get_image_size(image)
image_aux = self.image_processor.preprocess(image, return_tensors="pt")[
"pixel_values"
][0]
resize_h, resize_w = image_aux.shape[-2:]
data_dict["pixel_values_aux"] = image_aux.unsqueeze(0)
image = image_aux.clone()
image = torch.nn.functional.interpolate(
image[None],
size=[336, 336],
mode="bilinear",
align_corners=False,
)[0]
data_dict["pixel_values"] = image.unsqueeze(0)
# step2 load boxes
bbox= xyxy_to_xywh(bbox)
bbox = pad_boxes(bbox, (ori_w, ori_h))
bbox = resize_boxes(bbox, (pad_w, pad_h), (resize_h, resize_w))
data_dict["gt_boxes"] = torch.tensor(xywh_to_xyxy(bbox)).unsqueeze(0)
# step3 prepare question
total_num_boxes = len(bbox)
obj_tokens = [
DEFAULT_OBJECT_TOKEN.replace("<i>", str(i)) for i in range(total_num_boxes)
]
obj_tokens = (
DEFAULT_OBJECT_FEATURE_TOKEN.join(obj_tokens) + DEFAULT_OBJECT_FEATURE_TOKEN
)
question = question.replace(DEFAULT_IMAGE_TOKEN, "")
question = DEFAULT_IMAGE_TOKEN + "\n" + obj_tokens + "\n" + question
inputs = ""
inputs += self.template["INSTRUCTION"].format(input=question, round=1)
# step4 tokenize question
input_ids = tokenizer_image_object_token(inputs, self.tokenizer)
data_dict["input_ids"] = torch.tensor(input_ids).unsqueeze(0)
return data_dict
ChatRexProcessor.register_for_auto_class()