neuralabs/deepseek_ocr_de
526
1# MIT License modified from prithivMLmods/DeepSeek-OCR-Latest-BF16.I642import os3import math4import re5from tqdm import tqdm6from abc import ABC7from typing import List, Optional, Tuple, Union8 9from addict import Dict10from PIL import Image, ImageOps, ImageDraw, ImageFont11import numpy as np12 13import torch14import torch.nn as nn15from torch.nn import CrossEntropyLoss16from torchvision import transforms17 18from transformers.cache_utils import Cache19from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast20from transformers import DeepseekV2Model, DeepseekV2ForCausalLM21from transformers import DeepseekV2Config22from transformers.models.deepseek_v2.modeling_deepseek_v2 import (23 DeepseekV2Attention, DeepseekV2MLP, DeepseekV2MoE, DeepseekV2RMSNorm, DeepseekV2DecoderLayer)24from transformers.models.llama.modeling_llama import LlamaAttention, LlamaRotaryEmbedding25from transformers import TextStreamer26from .deepencoder import build_sam_vit_b, build_clip_l, MlpProjector27from .conversation import get_conv_template28 29torch_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float1630 31def load_image(image_path):32 33 try:34 image = Image.open(image_path)35 36 corrected_image = ImageOps.exif_transpose(image)37 38 return corrected_image39 40 except Exception as e:41 print(f"error: {e}")42 try:43 return Image.open(image_path)44 except:45 return None46 47 48def re_match(text):49 pattern = r'(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)'50 matches = re.findall(pattern, text, re.DOTALL)51 52 # pattern1 = r'<\|ref\|>.*?<\|/ref\|>\n'53 # new_text1 = re.sub(pattern1, '', text, flags=re.DOTALL)54 55 mathes_image = []56 mathes_other = []57 for a_match in matches:58 if '<|ref|>image<|/ref|>' in a_match[0]:59 mathes_image.append(a_match[0])60 else:61 mathes_other.append(a_match[0])62 return matches, mathes_image, mathes_other63 64 65def extract_coordinates_and_label(ref_text, image_width, image_height):66 67 try:68 label_type = ref_text[1]69 cor_list = eval(ref_text[2])70 except Exception as e:71 print(e)72 return None73 74 return (label_type, cor_list)75 76 77def draw_bounding_boxes(image, refs, ouput_path):78 79 image_width, image_height = image.size80 81 img_draw = image.copy()82 draw = ImageDraw.Draw(img_draw)83 84 overlay = Image.new('RGBA', img_draw.size, (0, 0, 0, 0))85 draw2 = ImageDraw.Draw(overlay)86 87 # try:88 # except IOError:89 # try:90 # font = ImageFont.truetype("DejaVuSans.ttf", 20) 91 # except IOError:92 font = ImageFont.load_default()93 94 img_idx = 095 96 for i, ref in enumerate(refs):97 try:98 result = extract_coordinates_and_label(ref, image_width, image_height)99 if result:100 label_type, points_list = result101 102 color = (np.random.randint(0, 200), np.random.randint(0, 200), np.random.randint(0, 255))103 104 color_a = color + (20, )105 for points in points_list:106 x1, y1, x2, y2 = points107 108 x1 = int(x1 / 999 * image_width)109 y1 = int(y1 / 999 * image_height)110 111 x2 = int(x2 / 999 * image_width)112 y2 = int(y2 / 999 * image_height)113 114 if label_type == 'image':115 try:116 cropped = image.crop((x1, y1, x2, y2))117 cropped.save(f"{ouput_path}/images/{img_idx}.jpg")118 except Exception as e:119 print(e)120 pass121 img_idx += 1122 123 try:124 if label_type == 'title':125 draw.rectangle([x1, y1, x2, y2], outline=color, width=4)126 draw2.rectangle([x1, y1, x2, y2], fill=color_a, outline=(0, 0, 0, 0), width=1)127 else:128 draw.rectangle([x1, y1, x2, y2], outline=color, width=2)129 draw2.rectangle([x1, y1, x2, y2], fill=color_a, outline=(0, 0, 0, 0), width=1)130 text_x = x1131 text_y = max(0, y1 - 15)132 133 134 text_bbox = draw.textbbox((0, 0), label_type, font=font)135 text_width = text_bbox[2] - text_bbox[0]136 text_height = text_bbox[3] - text_bbox[1]137 draw.rectangle([text_x, text_y, text_x + text_width, text_y + text_height], 138 fill=(255, 255, 255, 30))139 140 draw.text((text_x, text_y), label_type, font=font, fill=color)141 except:142 pass143 except:144 continue145 img_draw.paste(overlay, (0, 0), overlay)146 return img_draw147 148 149def process_image_with_refs(image, ref_texts, output_path):150 151 result_image = draw_bounding_boxes(image, ref_texts, output_path)152 153 return result_image154 155 156 157 158 159def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):160 best_ratio_diff = float('inf')161 best_ratio = (1, 1)162 area = width * height163 for ratio in target_ratios:164 target_aspect_ratio = ratio[0] / ratio[1]165 ratio_diff = abs(aspect_ratio - target_aspect_ratio)166 if ratio_diff < best_ratio_diff:167 best_ratio_diff = ratio_diff168 best_ratio = ratio169 elif ratio_diff == best_ratio_diff:170 if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:171 best_ratio = ratio172 # print(f'width: {width}, height: {height}, best_ratio: {best_ratio}')173 return best_ratio174 175 176def dynamic_preprocess(image, min_num=2, max_num=9, image_size=640, use_thumbnail=False):177 orig_width, orig_height = image.size178 aspect_ratio = orig_width / orig_height179 180 # calculate the existing image aspect ratio181 target_ratios = set(182 (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if183 i * j <= max_num and i * j >= min_num)184 # print(target_ratios)185 target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])186 187 # find the closest aspect ratio to the target188 target_aspect_ratio = find_closest_aspect_ratio(189 aspect_ratio, target_ratios, orig_width, orig_height, image_size)190 191 # print(target_aspect_ratio)192 # calculate the target width and height193 target_width = image_size * target_aspect_ratio[0]194 target_height = image_size * target_aspect_ratio[1]195 blocks = target_aspect_ratio[0] * target_aspect_ratio[1]196 197 # resize the image198 resized_img = image.resize((target_width, target_height))199 processed_images = []200 for i in range(blocks):201 box = (202 (i % (target_width // image_size)) * image_size,203 (i // (target_width // image_size)) * image_size,204 ((i % (target_width // image_size)) + 1) * image_size,205 ((i // (target_width // image_size)) + 1) * image_size206 )207 # split the image208 split_img = resized_img.crop(box)209 processed_images.append(split_img)210 assert len(processed_images) == blocks211 if use_thumbnail and len(processed_images) != 1:212 thumbnail_img = image.resize((image_size, image_size))213 processed_images.append(thumbnail_img)214 return processed_images, target_aspect_ratio215 216 217 218def normalize_transform(mean, std):219 if mean is None and std is None:220 transform = None221 elif mean is None and std is not None:222 mean = [0.] * len(std)223 transform = transforms.Normalize(mean=mean, std=std)224 elif mean is not None and std is None:225 std = [1.] * len(mean)226 transform = transforms.Normalize(mean=mean, std=std)227 else:228 transform = transforms.Normalize(mean=mean, std=std)229 230 return transform231 232 233 234def format_messages(235 conversations: List[Dict[str, str]],236 sft_format: str = "deepseek",237 system_prompt: str = "",238):239 """240 Applies the SFT template to conversation.241 242 Args:243 conversations (List[Dict]): A List of messages.244 sft_format (str, optional): The format of the SFT template to use. Defaults to "deepseek".245 system_prompt (str, optional): The system prompt to use in the SFT template. Defaults to "".246 247 Returns:248 sft_prompt (str): The formatted text.249 """250 251 conv = get_conv_template(sft_format)252 conv.set_system_message(system_prompt)253 for message in conversations:254 conv.append_message(message["role"], message["content"].strip())255 sft_prompt = conv.get_prompt().strip()256 257 return sft_prompt258 259 260def text_encode(tokenizer, text: str, bos: bool = True, eos: bool = False):261 t = tokenizer.encode(text, add_special_tokens=False)262 bos_id = 0263 eos_id = 1264 if bos:265 t = [bos_id] + t266 if eos:267 t = t + [eos_id]268 269 return t270 271def load_pil_images(conversations: List[Dict[str, str]]) -> List[Image.Image]:272 """273 274 Args:275 conversations (List[Dict[str, str]]): the conversations with a list of messages. An example is :276 [277 {278 "role": "User",279 "content": "<image_placeholder>\nExtract all information from this image and convert them into markdown format.",280 "images": ["./examples/table_datasets.png"]281 },282 {"role": "Assistant", "content": ""},283 ]284 285 Returns:286 pil_images (List[PIL.Image.Image]): the list of PIL images.287 288 """289 290 pil_images = []291 292 for message in conversations:293 if "images" not in message:294 continue295 296 for image_path in message["images"]:297 # print('----------------')298 # print(image_path)299 # print('----------------')300 # exit()301 302 # pil_img = Image.open(image_path)303 pil_img = load_image(image_path)304 pil_img = pil_img.convert("RGB")305 pil_images.append(pil_img)306 307 return pil_images308 309 310class BaseTransform(ABC):311 312 def set_rng(self, *args, **kwargs):313 pass314 315 def __call__(self, *args, **kwargs) -> torch.Tensor:316 pass317 318 @property319 def default_shape(self):320 raise NotImplementedError321 322 323class BasicImageTransform(BaseTransform):324 def __init__(325 self, 326 mean: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5),327 std: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5),328 normalize: bool = True329 ):330 self.mean = mean331 self.std = std332 333 transform_pipelines = [334 transforms.ToTensor()335 ]336 337 normalize = normalize_transform(mean, std) if normalize else nn.Identity()338 if normalize is not None:339 transform_pipelines.append(normalize)340 341 self.transform = transforms.Compose(transform_pipelines)342 343 def __call__(self, x):344 x = self.transform(x)345 return x346 347class NoEOSTextStreamer(TextStreamer):348 def on_finalized_text(self, text: str, stream_end: bool = False):349 350 eos_text = self.tokenizer.decode([self.tokenizer.eos_token_id], skip_special_tokens=False)351 text = text.replace(eos_text, "\n")352 print(text, flush=True, end="")353 354 355def decoder_layer_init(self, config: DeepseekV2Config, layer_idx: int):356 nn.Module.__init__(self)357 self.hidden_size = config.hidden_size358 359 if config.use_mla:360 self.self_attn = DeepseekV2Attention(config=config, layer_idx=layer_idx)361 else:362 config.head_dim = config.hidden_size // config.num_attention_heads363 self.self_attn = LlamaAttention(config, layer_idx)364 self.mlp = DeepseekV2MoE(config) if layer_idx >= config.first_k_dense_replace else DeepseekV2MLP(config)365 366 self.input_layernorm = DeepseekV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)367 self.post_attention_layernorm = DeepseekV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)368 369 370DeepseekV2DecoderLayer.__init__ = decoder_layer_init371 372class DeepseekOCRConfig(DeepseekV2Config):373 model_type = "DeepseekOCR"374 375class DeepseekOCRModel(DeepseekV2Model):376 config_class = DeepseekOCRConfig377 378 def __init__(self, config: DeepseekV2Config):379 super(DeepseekOCRModel, self).__init__(config)380 381 self.sam_model = build_sam_vit_b()382 self.vision_model = build_clip_l()383 # self.conv_2 = nn.Conv2d(in_channels=1024, out_channels=2048, kernel_size=2, stride=2)384 n_embed = 1280385 self.projector = MlpProjector(Dict(projector_type="linear", input_dim=2048, n_embed=n_embed))386 embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32))387 self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std)388 self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std)389 390 self.rotary_emb = LlamaRotaryEmbedding(config=config)391 392 def forward(393 self,394 input_ids: torch.LongTensor = None,395 attention_mask: Optional[torch.Tensor] = None,396 position_ids: Optional[torch.LongTensor] = None,397 past_key_values: Optional[List[torch.FloatTensor]] = None,398 inputs_embeds: Optional[torch.FloatTensor] = None,399 use_cache: Optional[bool] = None,400 output_attentions: Optional[bool] = None,401 output_hidden_states: Optional[bool] = None,402 images: Optional[torch.FloatTensor] = None,403 images_seq_mask: Optional[torch.FloatTensor] = None,404 images_spatial_crop: Optional[torch.FloatTensor] = None,405 return_dict: Optional[bool] = None,406 ) -> Union[Tuple, BaseModelOutputWithPast]:407 408 409 410 if inputs_embeds is None:411 # inputs_embeds = self.embed_tokens(input_ids)412 inputs_embeds = self.get_input_embeddings()(input_ids)413 414 inputs_embeds = inputs_embeds.clone()415 416 sam_model = getattr(self, 'sam_model', None)417 # sam_model = self.sam_model418 vision_model = getattr(self, 'vision_model', None)419 420 421 422 if sam_model is not None and (input_ids.shape[1] != 1 or self.training) and torch.sum(images[0][1]).item() != 0:423 424 idx = 0425 426 # sam_model = torch.jit.script(sam_model)427 428 # start_time = time.time()429 for image, crop_shape in zip(images, images_spatial_crop):430 images_in_this_batch = []431 432 patches = image[0]433 image_ori = image[1]434 435 with torch.no_grad():436 # with torch.inference_mode(): 437 438 if torch.sum(patches).item() != 0:439 # P, C, H, W = patches.shape440 crop_flag = 1441 local_features_1 = sam_model(patches)442 443 local_features_2 = vision_model(patches, local_features_1) 444 # vit_time = time.time()445 local_features = torch.cat((local_features_2[:, 1:], local_features_1.flatten(2).permute(0, 2, 1)), dim=-1) 446 local_features = self.projector(local_features)447 448 449 global_features_1 = sam_model(image_ori)450 global_features_2 = vision_model(image_ori, global_features_1) 451 global_features = torch.cat((global_features_2[:, 1:], global_features_1.flatten(2).permute(0, 2, 1)), dim=-1) 452 global_features = self.projector(global_features)453 454 print('=====================')455 print('BASE: ', global_features.shape)456 print('PATCHES: ', local_features.shape)457 print('=====================')458 459 _, hw, n_dim = global_features.shape460 h = w = int(hw ** 0.5)461 462 _2, hw2, n_dim2 = local_features.shape463 h2 = w2 = int(hw2 ** 0.5)464 465 width_crop_num, height_crop_num = crop_shape[0], crop_shape[1]466 467 global_features = global_features.view(h, w, n_dim)468 469 global_features = torch.cat(470 [global_features, self.image_newline[None, None, :].expand(h, 1, n_dim)], dim=1471 )472 473 global_features = global_features.view(-1, n_dim)474 475 476 local_features = local_features.view(height_crop_num, width_crop_num, h2, w2, n_dim2).permute(0, 2, 1, 3, 4).reshape(height_crop_num*h2, width_crop_num*w2, n_dim2)477 local_features = torch.cat(478 [local_features, self.image_newline[None, None, :].expand(height_crop_num * h2, 1, n_dim2)], dim=1479 )480 local_features = local_features.view(-1, n_dim2)481 482 global_local_features = torch.cat([local_features, global_features, self.view_seperator[None, :]], dim=0)483 484 # end_time = time.time()485 486 # print('sam: ', sam_time - start_time)487 # print('vit: ', vit_time - sam_time)488 # print('all: ', end_time - start_time)489 490 # exit()491 492 else:493 global_features_1 = sam_model(image_ori)494 global_features_2 = vision_model(image_ori, global_features_1) 495 global_features = torch.cat((global_features_2[:, 1:], global_features_1.flatten(2).permute(0, 2, 1)), dim=-1) 496 global_features = self.projector(global_features)497 _, hw, n_dim = global_features.shape498 h = w = int(hw ** 0.5)499 500 501 global_features = global_features.view(h, w, n_dim)502 503 global_features = torch.cat(504 [global_features, self.image_newline[None, None, :].expand(h, 1, n_dim)], dim=1505 )506 507 global_features = global_features.view(-1, n_dim)508 509 global_local_features = torch.cat([global_features, self.view_seperator[None, :]], dim=0)510 511 images_in_this_batch.append(global_local_features)512 513 514 if images_in_this_batch:515 images_in_this_batch = torch.cat(images_in_this_batch, dim=0)516 images_in_this_batch = images_in_this_batch.to(517 device=inputs_embeds.device, dtype=inputs_embeds.dtype518 )519 mask = images_seq_mask[idx].unsqueeze(-1).to(inputs_embeds.device) # bool [T, 1]520 updated_row = inputs_embeds[idx].masked_scatter(mask, images_in_this_batch)521 inputs_embeds[idx] = updated_row522 523 idx += 1524 525 return super(DeepseekOCRModel, self).forward(526 input_ids=None, attention_mask=attention_mask, past_key_values=past_key_values,527 inputs_embeds=inputs_embeds, use_cache=use_cache, position_ids = position_ids,528 output_attentions=output_attentions, output_hidden_states=output_hidden_states,529 return_dict=return_dict530 )531 532 533class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):534 535 config_class = DeepseekOCRConfig536 # supports_gradient_checkpointing = True537 538 def __init__(self, config):539 super(DeepseekV2ForCausalLM, self).__init__(config)540 self.model = DeepseekOCRModel(config)541 542 self.vocab_size = config.vocab_size543 544 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)545 546 # Initialize weights and apply final processing547 self.post_init()548 549 def get_model(self):550 return self.model551 552 553 def forward(554 self,555 input_ids: torch.LongTensor = None,556 attention_mask: Optional[torch.Tensor] = None,557 position_ids: Optional[torch.LongTensor] = None,558 past_key_values: Optional[List[torch.FloatTensor]] = None,559 inputs_embeds: Optional[torch.FloatTensor] = None,560 labels: Optional[torch.LongTensor] = None,561 use_cache: Optional[bool] = None,562 output_attentions: Optional[bool] = None,563 output_hidden_states: Optional[bool] = None,564 images: Optional[torch.FloatTensor] = None,565 images_seq_mask: Optional[torch.FloatTensor] = None,566 images_spatial_crop: Optional[torch.FloatTensor] = None,567 return_dict: Optional[bool] = None,568 569 ) -> Union[Tuple, CausalLMOutputWithPast]:570 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions571 output_hidden_states = (572 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states573 )574 return_dict = return_dict if return_dict is not None else self.config.use_return_dict575 576 577 578 outputs = self.model(579 input_ids=input_ids,580 past_key_values=past_key_values,581 attention_mask=attention_mask,582 position_ids=position_ids,583 inputs_embeds=inputs_embeds,584 use_cache=use_cache,585 output_attentions=output_attentions,586 output_hidden_states=output_hidden_states,587 images=images,588 images_seq_mask = images_seq_mask,589 images_spatial_crop = images_spatial_crop,590 return_dict=return_dict591 592 )593 594 hidden_states = outputs[0]595 logits = self.lm_head(hidden_states)596 logits = logits.float()597 598 # logits599 600 loss = None601 if labels is not None:602 # Shift so that tokens < n predict n603 shift_logits = logits[..., :-1, :].contiguous()604 shift_labels = labels[..., 1:].contiguous()605 # Flatten the tokens606 loss_fct = CrossEntropyLoss()607 shift_logits = shift_logits.view(-1, self.config.vocab_size)608 shift_labels = shift_labels.view(-1)609 # Enable model parallelism610 shift_labels = shift_labels.to(shift_logits.device)611 loss = loss_fct(shift_logits, shift_labels)612 613 if not return_dict:614 output = (logits,) + outputs[1:]615 return (loss,) + output if loss is not None else output616 617 return CausalLMOutputWithPast(618 loss=loss,619 logits=logits,620 past_key_values=outputs.past_key_values,621 hidden_states=outputs.hidden_states,622 attentions=outputs.attentions,623 )624 625 626 def prepare_inputs_for_generation(627 self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs628 ):629 # Omit tokens covered by past_key_values630 past_length = 0631 if past_key_values is not None:632 if isinstance(past_key_values, Cache):633 cache_length = past_key_values.get_seq_length()634 past_length = past_key_values.get_seq_length()635 max_cache_length = None636 else:637 cache_length = past_length = past_key_values[0][0].shape[2]638 max_cache_length = None639 640 # Keep only the unprocessed tokens:641 # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where642 # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as643 # input)644 if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:645 input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]646 # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard647 # input_ids based on the past_length.648 elif past_length < input_ids.shape[1]:649 input_ids = input_ids[:, past_length:]650 # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.651 652 # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.653 if (654 max_cache_length is not None655 and attention_mask is not None656 and cache_length + input_ids.shape[1] > max_cache_length657 ):658 attention_mask = attention_mask[:, -max_cache_length:]659 660 position_ids = kwargs.get("position_ids", None)661 if attention_mask is not None and position_ids is None:662 # create position_ids on the fly for batch generation663 position_ids = attention_mask.long().cumsum(-1) - 1664 position_ids.masked_fill_(attention_mask == 0, 1)665 if past_key_values:666 position_ids = position_ids[:, -input_ids.shape[1] :]667 668 # if self.generation_config.cache_implementation == "static":669 # # generation with static cache670 # cache_position = kwargs.get("cache_position", None)671 # if cache_position is None:672 # past_length = 0673 # else:674 # past_length = cache_position[-1] + 1675 # input_ids = input_ids[:, past_length:]676 # position_ids = position_ids[:, past_length:]677 678 # TODO @gante we should only keep a `cache_position` in generate, and do +=1.679 # same goes for position ids. Could also help with continued generation.680 cache_position = torch.arange(past_length, past_length + position_ids.shape[-1], device=position_ids.device)681 682 # if `inputs_embeds` are passed, we only want to use them in the 1st generation step683 if inputs_embeds is not None and past_key_values is None:684 model_inputs = {"inputs_embeds": inputs_embeds}685 else:686 model_inputs = {"input_ids": input_ids}687 688 model_inputs.update(689 {690 "position_ids": position_ids,691 "past_key_values": past_key_values,692 "use_cache": kwargs.get("use_cache"),693 "attention_mask": attention_mask,694 "images": kwargs.get("images", None),695 "images_seq_mask": kwargs.get("images_seq_mask", None),696 "images_spatial_crop": kwargs.get("images_spatial_crop", None),697 }698 )699 return model_inputs700 701 702 def disable_torch_init(self):703 """704 Disable the redundant torch default initialization to accelerate model creation.705 """706 import torch707 setattr(torch.nn.Linear, "reset_parameters", lambda self: None)708 setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)709 710 711 712 def infer(self, tokenizer, prompt='', image_file='', output_path = '', base_size=1024, image_size=640, crop_mode=True, test_compress=False, save_results=False, eval_mode=False):713 self.disable_torch_init()714 715 os.makedirs(output_path, exist_ok=True)716 os.makedirs(f'{output_path}/images', exist_ok=True)717 718 if prompt and image_file:719 conversation = [720 {721 "role": "<|User|>",722 # "content": "<image>\n<|grounding|>Given the layout of the image. ",723 "content": f'{prompt}',724 # "content": "君不见黄河之水天上来的下一句是什么？",725 # "content": "<image>\nFree OCR. ",726 # "content": "<image>\nParse the figure. ",727 # "content": "<image>\nExtract the text in the image. ",728 "images": [f'{image_file}'],729 },730 {"role": "<|Assistant|>", "content": ""},731 ]732 733 elif prompt:734 conversation = [735 {736 "role": "<|User|>",737 # "content": "<image>\n<|grounding|>Given the layout of the image. ",738 "content": f'{prompt}',739 # "content": "君不见黄河之水天上来的下一句是什么？",740 # "content": "<image>\nFree OCR. ",741 # "content": "<image>\nParse the figure. ",742 # "content": "<image>\nExtract the text in the image. ",743 # "images": [f'{image_file}'],744 },745 {"role": "<|Assistant|>", "content": ""},746 ]747 else:748 assert False, f'prompt is none!'749 750 prompt = format_messages(conversations=conversation, sft_format='plain', system_prompt='')751 752 patch_size = 16753 downsample_ratio = 4754 images = load_pil_images(conversation)755 756 valid_img_tokens = 0757 ratio = 1758 759 image_draw = images[0].copy()760 761 w,h = image_draw.size762 # print(w, h)763 ratio = 1 - ((max(w, h) - min(w, h)) / (max(w, h)))764 765 766 image_transform=BasicImageTransform(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), normalize=True)767 images_seq_mask = []768 769 image_token = '<image>'770 image_token_id = 128815771 text_splits = prompt.split(image_token)772 773 images_list, images_crop_list, images_seq_mask = [], [], []774 tokenized_str = []775 images_spatial_crop = []776 for text_sep, image in zip(text_splits, images):777 778 tokenized_sep = text_encode(tokenizer, text_sep, bos=False, eos=False)779 tokenized_str += tokenized_sep780 images_seq_mask += [False] * len(tokenized_sep)781 782 if crop_mode:783 784 if image.size[0] <= 640 and image.size[1] <= 640:785 crop_ratio = [1, 1]786 787 else:788 if crop_mode:789 # best_width, best_height = select_best_resolution(image.size, self.candidate_resolutions)790 images_crop_raw, crop_ratio = dynamic_preprocess(image)791 else:792 # best_width, best_height = self.image_size, self.image_size793 crop_ratio = [1, 1]794 795 """process the global view"""796 # image = image.resize((base_size, base_size))797 global_view = ImageOps.pad(image, (base_size, base_size),798 color=tuple(int(x * 255) for x in image_transform.mean))799 800 if base_size == 1024:801 valid_img_tokens += int(256 * ratio)802 elif base_size == 1280:803 valid_img_tokens += int(400 * ratio)804 # elif base_size == 640:805 # valid_img_tokens += int(100 * ratio)806 807 808 809 810 811 images_list.append(image_transform(global_view).to(torch_dtype))812 813 # global_view_tensor = image_transform(global_view).to(torch_dtype)814 815 width_crop_num, height_crop_num = crop_ratio816 817 images_spatial_crop.append([width_crop_num, height_crop_num])818 819 820 if width_crop_num > 1 or height_crop_num > 1:821 """process the local views"""822 823 for i in range(len(images_crop_raw)):824 images_crop_list.append(image_transform(images_crop_raw[i]).to(torch_dtype))825 826 if image_size == 640:827 valid_img_tokens += len(images_crop_list) * 100828 829 num_queries = math.ceil((image_size // patch_size) / downsample_ratio)830 num_queries_base = math.ceil((base_size // patch_size) / downsample_ratio)831 832 833 834 """add image tokens"""835 836 837 838 tokenized_image = ([image_token_id] * num_queries_base + [image_token_id]) * num_queries_base839 tokenized_image += [image_token_id]840 if width_crop_num > 1 or height_crop_num > 1:841 tokenized_image += ([image_token_id] * (num_queries * width_crop_num) + [image_token_id]) * (842 num_queries * height_crop_num)843 tokenized_str += tokenized_image844 images_seq_mask += [True] * len(tokenized_image)845 # num_image_tokens.append(len(tokenized_image))846 847 else:848 # best_width, best_height = self.image_size, self.image_size849 # print(image.size, (best_width, best_height)) # check the select_best_resolutions func850 851 """process the global view"""852 if image_size <= 640:853 print('directly resize')854 image = image.resize((image_size, image_size))855 # else:856 global_view = ImageOps.pad(image, (image_size, image_size),857 color=tuple(int(x * 255) for x in image_transform.mean))858 images_list.append(image_transform(global_view).to(torch_dtype))859 860 if base_size == 1024:861 valid_img_tokens += int(256 * ratio)862 elif base_size == 1280:863 valid_img_tokens += int(400 * ratio)864 elif base_size == 640:865 valid_img_tokens += int(100 * 1)866 elif base_size == 512:867 valid_img_tokens += int(64 * 1)868 869 width_crop_num, height_crop_num = 1, 1870 871 images_spatial_crop.append([width_crop_num, height_crop_num])872 873 874 """add image tokens"""875 num_queries = math.ceil((image_size // patch_size) / downsample_ratio)876 877 tokenized_image = ([image_token_id] * num_queries + [image_token_id]) * num_queries878 tokenized_image += [image_token_id]879 # tokenized_image += ([self.image_token_id] * (num_queries * width_crop_num) + [self.image_token_id]) * (880 # num_queries * height_crop_num)881 tokenized_str += tokenized_image882 images_seq_mask += [True] * len(tokenized_image)883 # num_image_tokens.append(len(tokenized_image))884 885 886 """process the last text split"""887 tokenized_sep = text_encode(tokenizer, text_splits[-1], bos=False, eos=False)888 tokenized_str += tokenized_sep889 images_seq_mask += [False] * len(tokenized_sep)890 891 """add the bos tokens"""892 bos_id = 0893 tokenized_str = [bos_id] + tokenized_str 894 images_seq_mask = [False] + images_seq_mask895 896 897 898 input_ids = torch.LongTensor(tokenized_str)899 900 images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool)901 902 903 if len(images_list) == 0:904 images_ori = torch.zeros((1, 3, image_size, image_size))905 images_spatial_crop = torch.zeros((1, 2), dtype=torch.long)906 images_crop = torch.zeros((1, 3, base_size, base_size))907 908 else:909 images_ori = torch.stack(images_list, dim=0)910 images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long)911 if images_crop_list:912 images_crop = torch.stack(images_crop_list, dim=0)913 else:914 images_crop = torch.zeros((1, 3, base_size, base_size))915 916 917 918 if not eval_mode:919 streamer = NoEOSTextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)920 with torch.autocast("cuda", dtype=torch_dtype):921 with torch.no_grad():922 output_ids = self.generate(923 input_ids.unsqueeze(0).cuda(),924 images=[(images_crop.cuda(), images_ori.cuda())],925 images_seq_mask = images_seq_mask.unsqueeze(0).cuda(),926 images_spatial_crop = images_spatial_crop,927 # do_sample=False,928 # num_beams = 1,929 temperature=0.0,930 eos_token_id=tokenizer.eos_token_id,931 streamer=streamer,932 max_new_tokens=8192,933 no_repeat_ngram_size = 20,934 use_cache = True935 )936 937 else:938 with torch.autocast("cuda", dtype=torch_dtype):939 with torch.no_grad():940 output_ids = self.generate(941 input_ids.unsqueeze(0).cuda(),942 images=[(images_crop.cuda(), images_ori.cuda())],943 images_seq_mask = images_seq_mask.unsqueeze(0).cuda(),944 images_spatial_crop = images_spatial_crop,945 # do_sample=False,946 # num_beams = 1,947 temperature=0.0,948 eos_token_id=tokenizer.eos_token_id,949 max_new_tokens=8192,950 no_repeat_ngram_size = 35,951 use_cache = True952 )953 954 955 if '<image>' in conversation[0]['content'] and eval_mode:956 outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])957 stop_str = '<｜end▁of▁sentence｜>'958 if outputs.endswith(stop_str):959 outputs = outputs[:-len(stop_str)]960 # re_match961 outputs = outputs.strip()962 963 return outputs964 965 if '<image>' in conversation[0]['content'] and test_compress:966 outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])967 pure_texts_outputs_token_length = len(text_encode(tokenizer, outputs, bos=False, eos=False))968 print('='*50)969 print('image size: ', (w, h))970 print('valid image tokens: ', int(valid_img_tokens))971 print('output texts tokens (valid): ', pure_texts_outputs_token_length)972 print('compression ratio: ', round(pure_texts_outputs_token_length/valid_img_tokens, 2))973 print('='*50)974 975 976 if '<image>' in conversation[0]['content'] and save_results:977 outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])978 stop_str = '<｜end▁of▁sentence｜>'979 980 print('='*15 + 'save results:' + '='*15)981 982 # # # # conv.messages[-1][-1] = outputs983 if outputs.endswith(stop_str):984 outputs = outputs[:-len(stop_str)]985 outputs = outputs.strip()986 987 matches_ref, matches_images, mathes_other = re_match(outputs)988 # print(matches_ref)989 result = process_image_with_refs(image_draw, matches_ref, output_path)990 991 992 for idx, a_match_image in enumerate(tqdm(matches_images, desc="image")):993 outputs = outputs.replace(a_match_image, ' + '.jpg)\n')994 995 for idx, a_match_other in enumerate(tqdm(mathes_other, desc="other")):996 outputs = outputs.replace(a_match_other, '').replace('\\coloneqq', ':=').replace('\\eqqcolon', '=:')997 998 999 # if 'structural formula' in conversation[0]['content']:1000 # outputs = '<smiles>' + outputs + '</smiles>'1001 with open(f'{output_path}/result.mmd', 'w', encoding = 'utf-8') as afile:1002 afile.write(outputs)1003 1004 if 'line_type' in outputs:1005 import matplotlib.pyplot as plt1006 lines = eval(outputs)['Line']['line']1007 1008 line_type = eval(outputs)['Line']['line_type']1009 # print(lines)1010 1011 endpoints = eval(outputs)['Line']['line_endpoint']1012 1013 fig, ax = plt.subplots(figsize=(3,3), dpi=200)1014 ax.set_xlim(-15, 15)1015 ax.set_ylim(-15, 15)1016 1017 for idx, line in enumerate(lines):1018 try:1019 p0 = eval(line.split(' -- ')[0])1020 p1 = eval(line.split(' -- ')[-1])1021 1022 if line_type[idx] == '--':1023 ax.plot([p0[0], p1[0]], [p0[1], p1[1]], linewidth=0.8, color='k')1024 else:1025 ax.plot([p0[0], p1[0]], [p0[1], p1[1]], linewidth = 0.8, color = 'k')1026 1027 ax.scatter(p0[0], p0[1], s=5, color = 'k')1028 ax.scatter(p1[0], p1[1], s=5, color = 'k')1029 except:1030 pass1031 1032 for endpoint in endpoints:1033 1034 label = endpoint.split(': ')[0]1035 (x, y) = eval(endpoint.split(': ')[1])1036 ax.annotate(label, (x, y), xytext=(1, 1), textcoords='offset points', 1037 fontsize=5, fontweight='light')1038 1039 1040 plt.savefig(f'{output_path}/geo.jpg')1041 plt.close()1042 1043 result.save(f"{output_path}/result_with_boxes.jpg")