| from PIL import Image |
| from itertools import zip_longest |
| import os |
| import requests |
| from urllib.parse import urlparse |
|
|
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
| import cv2 |
| import numpy as np |
|
|
| from model.llava.mm_utils import tokenizer_image_token |
|
|
| from model.SAM.utils.transforms import ResizeLongestSide |
| from dataset.utils.blip_processors import Blip2ImageEvalProcessor |
|
|
| from tools.prompting import write_list as _write_list |
|
|
| from tools.utils import ( |
| DEFAULT_IM_END_TOKEN, |
| DEFAULT_IM_START_TOKEN, |
| DEFAULT_IMAGE_TOKEN, |
| IMAGE_TOKEN_INDEX, |
| ) |
|
|
|
|
| def is_url(image): |
| try: |
| result = urlparse(image) |
| return all([result.scheme, result.netloc]) |
| except: |
| return False |
|
|
|
|
| class Chat: |
| IMG_MEAN = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1) |
| IMG_STD = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1) |
| IMG_SIZE = 1024 |
| IGNORE_LABEL = 255 |
| WHITE_BAND_WIDTH = 10 |
|
|
| def __init__( |
| self, |
| model, |
| tokenizer, |
| qformer_tokenizer, |
| args, |
| device="cuda:0", |
| ): |
| self.device = device |
| self.model = model |
| self.global_enc_processor = Blip2ImageEvalProcessor() |
| self.tokenizer = tokenizer |
| self.qformer_tokenizer = qformer_tokenizer |
| image_size = args.image_size |
| self.transform = ResizeLongestSide(image_size) |
| self.use_mm_start_end = args.use_mm_start_end |
| |
| try: |
| first_param = next(self.model.parameters()) |
| self.runtime_dtype = first_param.dtype |
| self.runtime_device = first_param.device |
| except Exception: |
| self.runtime_dtype = torch.bfloat16 |
| self.runtime_device = torch.device(self.device) |
| try: |
| ge_param = next( |
| self.base_model.model.grounding_encoder.image_encoder.parameters() |
| ) |
| self.grounding_dtype = ge_param.dtype |
| self.grounding_device = ge_param.device |
| except Exception: |
| self.grounding_dtype = self.runtime_dtype |
| self.grounding_device = self.runtime_device |
|
|
| def _load_image_to_rgb(self, image): |
| if is_url(image): |
| response = requests.get(image) |
| image_np = np.frombuffer(response.content, np.uint8) |
| image_np = cv2.imdecode(image_np, cv2.IMREAD_COLOR) |
| else: |
| image_np = cv2.imread(image) |
| if image_np is None: |
| raise ValueError(f"Unable to load image: {image}") |
| return cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB) |
|
|
| def _global_enc_from_array(self, image_np: np.ndarray) -> torch.Tensor: |
| return self.global_enc_processor(Image.fromarray(image_np)) |
|
|
| def _prepare_encodings(self, image_np: np.ndarray, image_identifier: str): |
| original_size = image_np.shape[:2] |
|
|
| global_enc_image = self._global_enc_from_array(image_np) |
| global_enc_image = global_enc_image.to( |
| device=self.runtime_device, dtype=self.runtime_dtype |
| ) |
|
|
| grounding_image = self.transform.apply_image(image_np) |
| resize = grounding_image.shape[:2] |
| grounding_enc_image = self.grounding_enc_processor( |
| torch.from_numpy(grounding_image).permute(2, 0, 1).contiguous() |
| ).to(device=self.grounding_device, dtype=self.grounding_dtype) |
|
|
| return ( |
| global_enc_image, |
| grounding_enc_image, |
| original_size, |
| resize, |
| image_np, |
| image_identifier, |
| ) |
|
|
|
|
|
|
|
|
| def grounding_enc_processor(self, x: torch.Tensor) -> torch.Tensor: |
| x = (x - self.IMG_MEAN) / self.IMG_STD |
| h, w = x.shape[-2:] |
| x = F.pad(x, (0, self.IMG_SIZE - w, 0, self.IMG_SIZE - h)) |
| return x |
|
|
| @property |
| def base_model(self): |
| if isinstance(self.model, nn.DataParallel): |
| return self.model.module |
| else: |
| return self.model |
|
|
| def ask(self, text, conv, num_images=2): |
| text = text.replace("<", "<").replace(">", ">") |
| prompt = text |
| if len(conv.messages) == 0: |
| if num_images > 1: |
| img_tokens = [ |
| f"{DEFAULT_IMAGE_TOKEN} (IMAGE{i + 1})" for i in range(num_images) |
| ] |
| img_tokens_list = _write_list(img_tokens) |
| begin_sentence = ( |
| f"The {img_tokens_list} provide an overview of the pictures.\n" |
| ) |
| elif num_images == 1: |
| begin_sentence = ( |
| f"The {DEFAULT_IMAGE_TOKEN} (IMAGE1) provides an overview of the picture.\n" |
| ) |
| else: |
| begin_sentence = "" |
| prompt = begin_sentence + prompt |
| if self.use_mm_start_end: |
| replace_token = ( |
| DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN |
| ) |
| prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token) |
| conv.append_message(conv.roles[0], prompt) |
|
|
| def answer( |
| self, |
| conv, |
| img_list, |
| question, |
| answer_prefix: str = "", |
| max_new_tokens=512, |
| max_length=2300, |
| ): |
| conv.append_message(conv.roles[1], "") |
| prompt = conv.get_prompt() |
| input_ids = tokenizer_image_token(prompt, self.tokenizer, return_tensors="pt") |
| if answer_prefix: |
| if DEFAULT_IMAGE_TOKEN in answer_prefix: |
| prefix_ids = tokenizer_image_token(answer_prefix, self.tokenizer) |
| else: |
| prefix_ids = self.tokenizer(answer_prefix).input_ids |
| if len(prefix_ids) > 0 and prefix_ids[0] == self.tokenizer.bos_token_id: |
| prefix_ids = prefix_ids[1:] |
| if len(prefix_ids) > 0 and prefix_ids[-1] == self.tokenizer.eos_token_id: |
| prefix_ids = prefix_ids[:-1] |
| if len(prefix_ids) > 0: |
| prefix_ids = torch.tensor(prefix_ids, dtype=input_ids.dtype) |
| input_ids = torch.cat([input_ids, prefix_ids], dim=0) |
| input_ids = input_ids.unsqueeze(0).to(self.device) |
|
|
| current_max_len = input_ids.shape[1] + max_new_tokens |
| begin_idx = max(0, current_max_len - max_length) |
| input_ids = input_ids[:, begin_idx:] |
| attention_masks = input_ids.ne(self.tokenizer.pad_token_id) |
| attention_masks = attention_masks.to(self.device) |
| conv_offset = torch.tensor([0, 1], device=self.device) |
|
|
| questions_list = question if isinstance(question, list) else [question] |
| if self.qformer_tokenizer is not None: |
| qformer_inputs = self.qformer_tokenizer( |
| questions_list, padding=True, truncation=True, return_tensors="pt" |
| ) |
| qformer_input_ids = qformer_inputs.input_ids.to(self.device) |
| qformer_attention_masks = qformer_inputs.attention_mask.to(self.device) |
| else: |
| qformer_input_ids = None |
| qformer_attention_masks = None |
|
|
| if len(img_list) > 0: |
| if len(img_list) > 1: |
| ( |
| global_enc_images, |
| grounding_enc_images, |
| original_sizes, |
| resizes, |
| _, |
| image_paths, |
| ) = list(zip(*img_list)) |
| global_enc_images = torch.stack(global_enc_images, dim=0) |
| grounding_enc_images = torch.stack(grounding_enc_images, dim=0) |
| original_size_list = list(zip_longest([None, *original_sizes])) |
| resize_list = list(zip_longest([None, *resizes])) |
| else: |
| ( |
| global_enc_image, |
| grounding_enc_image, |
| original_size, |
| resize, |
| _, |
| image_path, |
| ) = img_list[0] |
| global_enc_images = global_enc_image.unsqueeze(0) |
| grounding_enc_images = grounding_enc_image.unsqueeze(0) |
| image_paths = [image_path] |
| resizes = [resize] |
| original_sizes = [original_size] |
| slots = [None, *original_sizes] |
| resize_slots = [None, *resizes] |
| original_size_list = list(zip_longest(slots)) |
| resize_list = list(zip_longest(resize_slots)) |
| image_paths = [image_paths] |
|
|
| image_offset = torch.tensor( |
| [0, global_enc_images.shape[0]], device=self.device |
| ) |
| grounding_offset = torch.tensor( |
| [0, grounding_enc_images.shape[0]], device=self.device |
| ) |
|
|
| output_ids, pred_masks = self.model.evaluate( |
| global_enc_images, |
| grounding_enc_images, |
| input_ids, |
| attention_masks, |
| resize_list, |
| original_size_list, |
| conv_offset, |
| image_offset, |
| grounding_offset, |
| qformer_input_ids, |
| qformer_attention_masks, |
| max_new_tokens=max_new_tokens, |
| image_paths=image_paths, |
| eos_token_id=self.tokenizer.eos_token_id, |
| ) |
| output_ids = output_ids[0][output_ids[0] != IMAGE_TOKEN_INDEX] |
| else: |
| image_offset = torch.tensor([0, 0], device=self.device) |
| grounding_offset = torch.tensor([0, 0], device=self.device) |
| output_ids, _ = self.model.evaluate( |
| None, |
| None, |
| input_ids, |
| attention_masks, |
| None, |
| None, |
| conv_offset, |
| image_offset, |
| grounding_offset, |
| qformer_input_ids, |
| qformer_attention_masks, |
| max_new_tokens=max_new_tokens, |
| eos_token_id=self.tokenizer.eos_token_id, |
| ) |
| output_ids = output_ids[0] |
| pred_masks = None |
|
|
| text_output = self.tokenizer.decode(output_ids, skip_special_tokens=False) |
| text_output = " ".join(text_output.split()) |
| chunks = text_output.split("</s>") |
| llm_message = chunks[0] if len(chunks) <= 2 else chunks[1] |
| llm_message = llm_message.split("ASSISTANT: ")[-1].strip() |
| conv.messages.pop() |
| conv.append_message(conv.roles[1], llm_message) |
|
|
| return llm_message, pred_masks, text_output |
|
|
| def upload_img(self, image, conv, img_list): |
| image_np = self._load_image_to_rgb(image) |
| img_list.append(self._prepare_encodings(image_np, image)) |
| msg = "Received." |
| return msg |
|
|
|
|
