infimm-hd / processing_infimm_hd.py

liuhaogeng

first commit

b0b3b00 3 months ago

No virus

15.4 kB

	# coding=utf-8
	# Copyright 2022 The HuggingFace Inc. team.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.
	"""
	Processor class for InfiMMHD.
	"""

	import random
	from typing import List, Optional, Tuple, Union
	import torch
	import torchvision.transforms.functional as F
	from PIL import Image
	from torchvision.transforms import (
	CenterCrop,
	Compose,
	InterpolationMode,
	Normalize,
	Resize,
	ToTensor,
	)

	from transformers import AutoTokenizer
	from transformers.image_processing_utils import ImageProcessingMixin
	from transformers.processing_utils import ProcessorMixin
	from transformers.tokenization_utils_base import BatchEncoding

	IMAGE_TOKEN = "<image>"
	END_OF_CHUNK_TOKEN = "<\|endofchunk\|>"
	PAD_TOKEN = "<PAD>"

	OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
	OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)


	def _convert_to_rgb(image):
	return image.convert("RGB")


	class ResizeKeepRatio:
	"""Resize and Keep Ratio

	Copy & paste from `timm`
	"""

	def __init__(
	self,
	size,
	longest=0.0,
	interpolation=InterpolationMode.BICUBIC,
	random_scale_prob=0.0,
	random_scale_range=(0.85, 1.05),
	random_aspect_prob=0.0,
	random_aspect_range=(0.9, 1.11),
	):
	if isinstance(size, (list, tuple)):
	self.size = tuple(size)
	else:
	self.size = (size, size)
	self.interpolation = interpolation
	self.longest = float(longest) # [0, 1] where 0 == shortest edge, 1 == longest
	self.random_scale_prob = random_scale_prob
	self.random_scale_range = random_scale_range
	self.random_aspect_prob = random_aspect_prob
	self.random_aspect_range = random_aspect_range

	@staticmethod
	def get_params(
	img,
	target_size,
	longest,
	random_scale_prob=0.0,
	random_scale_range=(0.85, 1.05),
	random_aspect_prob=0.0,
	random_aspect_range=(0.9, 1.11),
	):
	"""Get parameters"""
	source_size = img.size[::-1] # h, w
	h, w = source_size
	target_h, target_w = target_size
	ratio_h = h / target_h
	ratio_w = w / target_w
	ratio = max(ratio_h, ratio_w) * longest + min(ratio_h, ratio_w) * (
	1.0 - longest
	)
	if random_scale_prob > 0 and random.random() < random_scale_prob:
	ratio_factor = random.uniform(random_scale_range[0], random_scale_range[1])
	ratio_factor = (ratio_factor, ratio_factor)
	else:
	ratio_factor = (1.0, 1.0)
	if random_aspect_prob > 0 and random.random() < random_aspect_prob:
	aspect_factor = random.uniform(
	random_aspect_range[0], random_aspect_range[1]
	)
	ratio_factor = (
	ratio_factor[0] / aspect_factor,
	ratio_factor[1] * aspect_factor,
	)
	size = [round(x * f / ratio) for x, f in zip(source_size, ratio_factor)]
	return size

	def __call__(self, img):
	"""
	Args:
	img (PIL Image): Image to be cropped and resized.

	Returns:
	PIL Image: Resized, padded to at least target size, possibly cropped to exactly target size
	"""
	size = self.get_params(
	img,
	self.size,
	self.longest,
	self.random_scale_prob,
	self.random_scale_range,
	self.random_aspect_prob,
	self.random_aspect_range,
	)
	img = F.resize(img, size, self.interpolation)
	return img

	def __repr__(self):
	format_string = self.__class__.__name__ + "(size={0}".format(self.size)
	format_string += f", interpolation={self.interpolation})"
	format_string += f", longest={self.longest:.3f})"
	return format_string


	def image_transform(
	image_size: Union[int, Tuple[int, int]],
	mean: Optional[Tuple[float, ...]] = None,
	std: Optional[Tuple[float, ...]] = None,
	resize_mode: Optional[str] = None,
	interpolation: Optional[str] = None,
	):
	mean = mean or OPENAI_DATASET_MEAN
	if not isinstance(mean, (list, tuple)):
	mean = (mean,) * 3

	std = std or OPENAI_DATASET_STD
	if not isinstance(std, (list, tuple)):
	std = (std,) * 3

	interpolation = interpolation or "bicubic"
	assert interpolation in ["bicubic", "bilinear", "random"]
	# NOTE random is ignored for interpolation_mode, so defaults to BICUBIC for inference if set
	interpolation_mode = (
	InterpolationMode.BILINEAR
	if interpolation == "bilinear"
	else InterpolationMode.BICUBIC
	)

	resize_mode = resize_mode or "shortest"
	assert resize_mode in ("shortest", "longest", "squash")

	normalize = Normalize(mean=mean, std=std)

	assert resize_mode == "shortest"
	if not isinstance(image_size, (tuple, list)):
	image_size = (image_size, image_size)
	if image_size[0] == image_size[1]:
	# simple case, use torchvision built-in Resize w/ shortest edge mode (scalar size arg)
	transforms = [Resize(image_size[0], interpolation=interpolation_mode)]
	else:
	# resize shortest edge to matching target dim for non-square target
	transforms = [ResizeKeepRatio(image_size)]
	transforms += [CenterCrop(image_size)]

	transforms.extend(
	[
	_convert_to_rgb,
	ToTensor(),
	normalize,
	]
	)
	return Compose(transforms)


	def get_target_size(width, height, max_image_size, min_image_size):
	target_width = 0
	target_height = 0
	if width < min_image_size:
	target_width = min_image_size
	elif width > max_image_size:
	target_width = max_image_size

	if height < min_image_size:
	target_height = min_image_size
	elif height > max_image_size:
	target_height = max_image_size

	if target_width == 0:
	ratio = ((width - min_image_size) + int(0.5*min_image_size))//min_image_size
	target_width = ratio * min_image_size + min_image_size

	if target_height == 0:
	ratio = ((height - min_image_size) + int(0.5*min_image_size))//min_image_size
	target_height = ratio * min_image_size + min_image_size

	return target_width, target_height

	class EVAClipImageProcessor(ImageProcessingMixin):
	def __init__(self, **kwargs) -> None:
	super().__init__(**kwargs)
	self.image_processor = image_transform(image_size=448)
	self.img_size = 448

	def _prepare_images(self, batch: List[List[Image.Image]]) -> torch.Tensor:
	"""
	Convert images to tensors, reshape them, and stack them.
	Args:
	batch: A list of lists of images.
	Returns:
	preprocessed images (tensors) or None
	shape (B, T_img, F, C, H, W)
	None if no images in batch
	"""

	target_image_num = []
	target_shape = []
	for x in batch:
	width, height = x[0].size
	tar_wid, tar_hei = get_target_size(width, height, 1344, self.img_size)
	target_shape.append((tar_wid, tar_hei))
	target_image_num.append(int(tar_wid/self.img_size*tar_hei/self.img_size))

	images_per_example = max(target_image_num)
	batch_images = None
	image_mask = None
	sub_image_shape = None
	for iexample, example in enumerate(batch):
	for img in example:
	img_ori = img
	tar_wid, tar_hei = target_shape[iexample]
	img_new = img.resize((tar_wid, tar_hei), Image.BILINEAR)
	sub_images = [img_ori]

	for y in range(0, tar_hei, self.img_size):
	for x in range(0, tar_wid, self.img_size):
	sub_img = img_new.crop((x, y, x + self.img_size, y + self.img_size))
	sub_images.append(sub_img)

	for iimage, image in enumerate(sub_images):
	preprocessed = self.image_processor(image)
	if batch_images is None:
	batch_images = torch.zeros(
	(len(batch), images_per_example+1, 1) + preprocessed.shape,
	dtype=preprocessed.dtype,
	)
	batch_images[iexample, iimage, 0] = preprocessed
	if not torch.is_tensor(image_mask):
	image_mask = torch.zeros((len(batch), images_per_example+1), dtype=preprocessed.dtype)
	image_mask[iexample,:target_image_num[iexample]+1] = 1.0
	if not torch.is_tensor(sub_image_shape):
	sub_image_shape = torch.zeros((len(batch), 2), dtype=preprocessed.dtype)
	sub_image_shape[iexample, 0], sub_image_shape[iexample, 1] = tar_wid/self.img_size, tar_hei/self.img_size

	# if batch_images is not None:
	# batch_images = batch_images.to(
	# self.device, dtype=self.cast_dtype, non_blocking=True
	# )

	# if image_mask is not None:
	# image_mask = image_mask.to(
	# self.device, dtype=self.cast_dtype, non_blocking=True
	# )

	# if sub_image_shape is not None:
	# sub_image_shape = sub_image_shape.to(
	# self.device, dtype=self.cast_dtype, non_blocking=True
	# )
	return batch_images, image_mask, sub_image_shape

	def preprocess(self, imgpaths=None):
	if imgpaths is None or len(imgpaths) == 0:
	images = [(Image.new("RGB", (224, 224), color="black"))]
	else:
	images = [Image.open(fp) for fp in imgpaths]
	return self._prepare_images([images])


	class InfiMMHDProcessor(ProcessorMixin):
	r"""
	Constructs a InfiMMLlama2 processor which wraps a tokenizer and an image processor into a single processor.

	Args:
	image_processor (`EVAClipImageProcessor`):
	An instance of [`EVAClipImageProcessor`]. The image processor is a required input.
	tokenizer (`LlamaTokenizer`):
	An instance of [`LlamaTokenizer`]. The tokenizer is a required input.
	image_size (`int`, optional, defaults to 336): Image size (assuming a square image)
	"""

	attributes = ["tokenizer"]
	tokenizer_class = "LlamaTokenizer"

	def __init__(self, tokenizer=None, **kwargs):
	self.image_processor = EVAClipImageProcessor()
	if tokenizer is None:
	tokenizer = AutoTokenizer.from_pretrained("infimm-hd", verbose=False)

	super().__init__(tokenizer, tokenizer)

	def _prepare_text(
	self,
	batch: List[List[str]],
	padding="longest",
	truncation=True,
	max_length=2048,
	):
	"""
	Tokenize the text and stack them.
	Args:
	batch: A list of lists of strings.
	Returns:
	input_ids (tensor)
	shape (B, T_txt)
	attention_mask (tensor)
	shape (B, T_txt)
	"""
	batch = [b.strip() for b in batch]
	encodings = self.tokenizer(
	batch,
	padding=padding,
	truncation=truncation,
	return_tensors="pt",
	max_length=max_length,
	)
	input_ids, attention_mask = encodings["input_ids"], encodings["attention_mask"]
	# print(self.tokenizer.convert_ids_to_tokens(input_ids[]))
	return input_ids, attention_mask

	def __call__(
	self,
	prompts,
	) -> BatchEncoding:
	"""This method takes batched or non-batched prompts made of text and images and converts them into prompts that
	the model was trained on and prepares the image pixel values for the model to process.
	"""
	image_paths = self._extract_image_paths(prompts)
	images, image_mask, sub_image_shape = self.image_processor.preprocess(image_paths)
	prompts = self._replace_with_media_tokens(prompts)
	final_prompt = self.apply_template(prompts)
	# system_prompt = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions."
	# final_prompt = f"{system_prompt} USER: <image>" + prompts + " ASSISTANT:"
	input_ids, attention_mask = self._prepare_text([final_prompt])
	return BatchEncoding(
	data={
	"input_ids": input_ids,
	"attention_mask": attention_mask,
	"batch_images": images,
	"image_mask": image_mask,
	"subimage_shape": sub_image_shape,
	}
	)

	def _extract_image_paths(self, prompts):
	image_paths = []
	for round in prompts:
	if round["role"] != "user":
	continue
	for piece in round["content"]:
	if isinstance(piece, dict):
	image_paths.append(piece["image"])
	return image_paths

	def _replace_with_media_tokens(self, prompts):
	new_prompts = []
	is_first_img = True
	for round in prompts:
	if round["role"] != "user":
	new_prompts.append(round)
	new_content = []
	for piece in round["content"]:
	if isinstance(piece, dict):
	new_content.append(
	f"{IMAGE_TOKEN}" if is_first_img
	else f"{END_OF_CHUNK_TOKEN}{IMAGE_TOKEN}"
	)
	is_first_img = False
	else:
	new_content.append(piece)
	new_prompts.append({"role": "user", "content": "".join(new_content)})
	return new_prompts

	def apply_template(self, messages, task="generation"):
	prompt = self.tokenizer.apply_chat_template(
	messages,
	tokenize=False,
	add_generation_prompt=True if task == "generation" else False,
	)
	return prompt

	def batch_decode(self, args, *kwargs):
	"""
	This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
	refer to the docstring of this method for more information.
	"""
	return self.tokenizer.batch_decode(args, *kwargs)

	def decode(self, args, *kwargs):
	"""
	This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
	the docstring of this method for more information.
	"""
	return self.tokenizer.decode(args, *kwargs)

	@property
	def model_input_names(self):
	tokenizer_input_names = self.tokenizer.model_input_names
	image_processor_input_names = self.image_processor.model_input_names
	return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))