| |
| |
|
|
| import numpy as np |
| from PIL import Image |
| from typing import List, Optional, Union, Dict, Any |
| import torch |
| from torchvision import transforms as T |
| import json |
|
|
| from transformers import BatchFeature, ProcessorMixin, BaseImageProcessor, ImageProcessingMixin |
| from transformers.image_utils import ChannelDimension, ImageInput, PILImageResampling, infer_channel_dimension_format |
| from transformers.utils import TensorType |
|
|
| OPENAI_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073) |
| OPENAI_CLIP_STD = (0.26862954, 0.26130258, 0.27577711) |
|
|
|
|
| from transformers.processing_utils import ImagesKwargs |
|
|
| class NemotronParseImagesKwargs(ImagesKwargs, total=False): |
| content_hw: Optional[List[int]] |
|
|
| class NemotronParseImageProcessor(BaseImageProcessor, ImageProcessingMixin): |
| """ |
| Image processor for NemotronParse model. |
| |
| This processor inherits from BaseImageProcessor to be compatible with transformers AutoImageProcessor. |
| """ |
| |
| valid_kwargs = NemotronParseImagesKwargs |
| model_input_names = ["pixel_values", "image_sizes", "content_sizes"] |
| |
| def __init__( |
| self, |
| final_size: tuple = (2048, 1664), |
| do_normalize: bool = False, |
| image_mean: Optional[Union[List[float], tuple]] = None, |
| image_std: Optional[Union[List[float], tuple]] = None, |
| **kwargs, |
| ): |
| self.navit_packing = bool(kwargs.pop('navit_packing', False)) |
| self.navit_padding_mode = kwargs.pop('navit_padding_mode', 'top_left_white') |
| self.navit_coordinate_frame = kwargs.pop('navit_coordinate_frame', 'content_normalized') |
| clean_kwargs = {} |
| for k, v in kwargs.items(): |
| if not k.startswith('_') and k not in ['transform', 'torch_transform']: |
| clean_kwargs[k] = v |
| |
| if 'size' in clean_kwargs: |
| size_config = clean_kwargs.pop('size') |
| if isinstance(size_config, dict): |
| if 'longest_edge' in size_config: |
| longest_edge = size_config['longest_edge'] |
| if isinstance(longest_edge, (list, tuple)): |
| final_size = tuple(int(x) for x in longest_edge) |
| else: |
| final_size = (int(longest_edge), int(longest_edge)) |
| elif 'height' in size_config and 'width' in size_config: |
| final_size = (int(size_config['height']), int(size_config['width'])) |
| |
| super().__init__(**clean_kwargs) |
| self.do_normalize = bool(do_normalize) |
| self.image_mean = list(image_mean or OPENAI_CLIP_MEAN) |
| self.image_std = list(image_std or OPENAI_CLIP_STD) |
| |
| if isinstance(final_size, (list, tuple)) and len(final_size) >= 2: |
| self.final_size = (int(final_size[0]), int(final_size[1])) |
| elif isinstance(final_size, (int, float)): |
| self.final_size = (int(final_size), int(final_size)) |
| else: |
| self.final_size = (2048, 1664) |
| |
| self._create_transforms() |
| |
| def _create_transforms(self): |
| """Create transform objects (not serialized to JSON).""" |
| if isinstance(self.final_size, (list, tuple)): |
| self.target_height, self.target_width = int(self.final_size[0]), int(self.final_size[1]) |
| else: |
| self.target_height = self.target_width = int(self.final_size) |
| |
| self.transform = None |
| |
| self.torch_transform = T.Compose([ |
| T.ToTensor(), |
| ]) |
|
|
| def to_dict(self): |
| """Override to exclude non-serializable transforms.""" |
| output = super().to_dict() |
| output.pop('transform', None) |
| output.pop('torch_transform', None) |
| return output |
| |
| @classmethod |
| def from_dict(cls, config_dict: dict, **kwargs): |
| """Override to recreate transforms after loading.""" |
| config_dict = config_dict.copy() |
| config_dict.pop('transform', None) |
| config_dict.pop('torch_transform', None) |
| |
| |
| for key in list(config_dict.keys()): |
| if key.startswith('_') or config_dict[key] is None: |
| config_dict.pop(key, None) |
| |
| |
| if 'final_size' in config_dict: |
| final_size = config_dict['final_size'] |
| if isinstance(final_size, (list, tuple)): |
| config_dict['final_size'] = tuple(int(x) for x in final_size) |
| if 'image_mean' in config_dict: |
| config_dict['image_mean'] = [float(x) for x in config_dict['image_mean']] |
| if 'image_std' in config_dict: |
| config_dict['image_std'] = [float(x) for x in config_dict['image_std']] |
| |
| try: |
| return cls(**config_dict, **kwargs) |
| except Exception as e: |
| print(f"Warning: Error in from_dict: {e}") |
| print("Using default parameters...") |
| return cls(**kwargs) |
| |
| def save_pretrained(self, save_directory, **kwargs): |
| """Save image processor configuration.""" |
| import os |
| import json |
| |
| os.makedirs(save_directory, exist_ok=True) |
| |
| |
| config = { |
| "feature_extractor_type": "NemotronParseImageProcessor", |
| "image_processor_type": "NemotronParseImageProcessor", |
| "processor_class": "NemotronParseProcessor", |
| "size": { |
| "height": self.final_size[0], |
| "width": self.final_size[1], |
| "longest_edge": self.final_size |
| }, |
| "final_size": self.final_size, |
| "do_normalize": self.do_normalize, |
| "navit_packing": self.navit_packing, |
| "navit_padding_mode": self.navit_padding_mode, |
| "navit_coordinate_frame": self.navit_coordinate_frame, |
| "image_mean": self.image_mean, |
| "image_std": self.image_std, |
| } |
| |
| config_path = os.path.join(save_directory, "preprocessor_config.json") |
| with open(config_path, 'w') as f: |
| json.dump(config, f, indent=2) |
|
|
| def _resize_with_aspect_ratio(self, image: np.ndarray) -> np.ndarray: |
| """Resize image maintaining aspect ratio (exact replica of original LongestMaxSizeHW).""" |
| height, width = image.shape[:2] |
| max_size_height = self.target_height |
| max_size_width = self.target_width |
| |
| |
| aspect_ratio = width / height |
| new_height = height |
| new_width = width |
|
|
| if height > max_size_height: |
| new_height = max_size_height |
| new_width = int(new_height * aspect_ratio) |
|
|
| if new_width > max_size_width: |
| new_width = max_size_width |
| new_height = int(new_width / aspect_ratio) |
| |
| if new_height == height and new_width == width: |
| return image |
| resampling = getattr(Image, "Resampling", Image).BILINEAR |
| return np.asarray(Image.fromarray(image).resize((new_width, new_height), resampling)) |
| |
| def _pad_to_size(self, image: np.ndarray) -> np.ndarray: |
| """Center-pad image to target size with white padding.""" |
| h, w = image.shape[:2] |
| min_height, min_width = self.target_height, self.target_width |
| |
| pad_h = max(0, min_height - h) |
| pad_w = max(0, min_width - w) |
| |
| if pad_h == 0 and pad_w == 0: |
| return image |
| |
| pad_top = pad_h // 2 |
| pad_bottom = pad_h - pad_top |
| pad_left = pad_w // 2 |
| pad_right = pad_w - pad_left |
| |
| if len(image.shape) == 3: |
| return np.pad( |
| image, |
| ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), |
| mode='constant', |
| constant_values=255, |
| ) |
| return np.pad( |
| image, |
| ((pad_top, pad_bottom), (pad_left, pad_right)), |
| mode='constant', |
| constant_values=255, |
| ) |
| |
| def _pad_to_navit_grid(self, image: np.ndarray) -> np.ndarray: |
| height, width = image.shape[:2] |
| target_height = ((height + 15) // 16) * 16 |
| target_width = ((width + 63) // 64) * 64 |
| pad_h, pad_w = target_height - height, target_width - width |
| if not pad_h and not pad_w: |
| return image |
| padding = ((0, pad_h), (0, pad_w), (0, 0)) if image.ndim == 3 else ((0, pad_h), (0, pad_w)) |
| return np.pad(image, padding, mode="constant", constant_values=255) |
|
|
| def preprocess( |
| self, |
| images: ImageInput, |
| return_tensors: Optional[Union[str, TensorType]] = None, |
| **kwargs, |
| ) -> Dict[str, torch.Tensor]: |
| """ |
| Preprocess an image or batch of images for the NemotronParse model. |
| |
| Args: |
| images: Input image(s) |
| return_tensors: Type of tensors to return |
| """ |
| |
| captured_hw = kwargs.pop("content_hw", None) |
| image_options = kwargs.pop("images_kwargs", {}) or {} |
| nested_hw = image_options.get("content_hw") |
| if captured_hw is not None and nested_hw is not None and list(captured_hw) != list(nested_hw): |
| raise ValueError("conflicting captured content dimensions") |
| if captured_hw is None: |
| captured_hw = nested_hw |
| if captured_hw is not None: |
| if not self.navit_packing or self.navit_coordinate_frame != "content_normalized": |
| raise ValueError("loader capture requires content-normalized NaViT") |
| if self.navit_padding_mode != "top_left_white": |
| raise ValueError("unsupported captured loader pixel contract") |
| pages = images if isinstance(images, list) else [images] |
| if len(pages) != 1 or len(captured_hw) != 2: |
| raise ValueError("one captured page and [content_height, content_width] required") |
| pixels = np.array(pages[0].convert("RGB"), copy=True) |
| height, width = pixels.shape[:2] |
| ch, cw = (int(value) for value in captured_hw) |
| if list(captured_hw) != [ch, cw] or not (0 < ch <= height and 0 < cw <= width): |
| raise ValueError("invalid captured content dimensions") |
| if height % 16 or width % 64: |
| raise ValueError("captured canvas is not aligned to the NaViT grid") |
| page = torch.from_numpy(pixels).permute(2, 0, 1).contiguous().float().div_(255.0) |
| |
| |
| if self.do_normalize: |
| mean = page.new_tensor(self.image_mean).view(-1, 1, 1) |
| std = page.new_tensor(self.image_std).view(-1, 1, 1) |
| page = (page - mean) / std |
| patches = page.unfold(1, 16, 16).unfold(2, 16, 16) |
| patches = patches.permute(1, 2, 0, 3, 4).reshape(-1, 3 * 16 * 16) |
| return BatchFeature(data={ |
| "pixel_values": patches, |
| "image_sizes": torch.tensor([[height, width]], dtype=torch.long), |
| "content_sizes": torch.tensor([[ch, cw]], dtype=torch.long), |
| }, tensor_type=return_tensors) |
|
|
| |
| if not isinstance(images, list): |
| images = [images] |
|
|
| |
| for i in range(len(images)): |
| images[i] = images[i].convert('RGB') |
| |
| |
| processed_images = [] |
| for image in images: |
| if isinstance(image, Image.Image): |
| image = np.asarray(image) |
| processed_images.append(image) |
| |
| |
| pixel_values = [] |
| content_sizes = [] |
| for image in processed_images: |
| processed_image = self._resize_with_aspect_ratio(image) |
| content_sizes.append( |
| [int(processed_image.shape[0]), int(processed_image.shape[1])] |
| ) |
| if self.navit_packing: |
| processed_image = self._pad_to_navit_grid(processed_image) |
| |
| if self.transform is not None: |
| transformed = self.transform(image=processed_image) |
| processed_image = transformed["image"] |
| else: |
| |
| if not self.navit_packing: |
| processed_image = self._pad_to_size(processed_image) |
| |
| pixel_values_tensor = self.torch_transform(processed_image) |
| |
| if pixel_values_tensor.shape[0] == 1: |
| pixel_values_tensor = pixel_values_tensor.expand(3, -1, -1) |
| |
| pixel_values.append(pixel_values_tensor) |
| |
| if self.navit_packing: |
| content_sizes = torch.tensor(content_sizes, dtype=torch.long) |
| image_sizes = torch.tensor( |
| [[int(page.shape[-2]), int(page.shape[-1])] for page in pixel_values], |
| dtype=torch.long, |
| ) |
| raw_patches = [] |
| for page in pixel_values: |
| if self.do_normalize: |
| mean = page.new_tensor(self.image_mean).view(-1, 1, 1) |
| std = page.new_tensor(self.image_std).view(-1, 1, 1) |
| page = (page - mean) / std |
| patches = page.unfold(1, 16, 16).unfold(2, 16, 16) |
| patches = patches.permute(1, 2, 0, 3, 4).reshape(-1, 3 * 16 * 16) |
| raw_patches.append(patches) |
| pixel_values = torch.cat(raw_patches, dim=0) |
| data = { |
| "pixel_values": pixel_values, |
| "image_sizes": image_sizes, |
| "content_sizes": content_sizes, |
| } |
| else: |
| pixel_values = torch.stack(pixel_values) |
| if self.do_normalize: |
| mean = pixel_values.new_tensor(self.image_mean).view(1, -1, 1, 1) |
| std = pixel_values.new_tensor(self.image_std).view(1, -1, 1, 1) |
| pixel_values = (pixel_values - mean) / std |
| data = {"pixel_values": pixel_values} |
| |
| if return_tensors is not None: |
| data = self._convert_output_format(data, return_tensors) |
| |
| return data |
| |
| def _convert_output_format(self, data: Dict[str, torch.Tensor], return_tensors: Union[str, TensorType]) -> Dict: |
| """Convert output format based on return_tensors parameter.""" |
| if return_tensors == "pt" or return_tensors == TensorType.PYTORCH: |
| return data |
| elif return_tensors == "np" or return_tensors == TensorType.NUMPY: |
| return {k: v.numpy() for k, v in data.items()} |
| else: |
| return data |
| |
| def __call__(self, images: Union[Image.Image, List[Image.Image]], **kwargs) -> Dict[str, torch.Tensor]: |
| """Process images for the model (backward compatibility).""" |
| return self.preprocess(images, **kwargs) |
|
|
|
|
| class NemotronParseProcessor(ProcessorMixin): |
| |
| attributes = ["image_processor", "tokenizer"] |
| image_processor_class = "AutoImageProcessor" |
| tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast") |
| |
| def __init__(self, image_processor=None, tokenizer=None, **kwargs): |
| if image_processor is None: |
| image_processor = NemotronParseImageProcessor(**kwargs) |
| |
| super().__init__(image_processor, tokenizer) |
| |
| |
| def __call__( |
| self, |
| images: Union[Image.Image, List[Image.Image]] = None, |
| text: Union[str, List[str]] = None, |
| add_special_tokens: bool = True, |
| padding: Union[bool, str] = False, |
| truncation: Union[bool, str] = False, |
| max_length: Optional[int] = None, |
| stride: int = 0, |
| pad_to_multiple_of: Optional[int] = None, |
| return_attention_mask: Optional[bool] = None, |
| return_overflowing_tokens: bool = False, |
| return_special_tokens_mask: bool = False, |
| return_offsets_mapping: bool = False, |
| return_token_type_ids: bool = False, |
| return_length: bool = False, |
| verbose: bool = True, |
| return_tensors: Optional[Union[str, "TensorType"]] = None, |
| content_hw: Optional[List[int]] = None, |
| **kwargs |
| ) -> BatchFeature: |
| """ |
| Main method to prepare for the model one or several text(s) and image(s). |
| """ |
| |
| |
| if images is not None: |
| image_inputs = self.image_processor( |
| images, content_hw=content_hw, **kwargs) |
| else: |
| image_inputs = {} |
| |
| |
| if text is not None: |
| text_inputs = self.tokenizer( |
| text, |
| add_special_tokens=add_special_tokens, |
| padding=padding, |
| truncation=truncation, |
| max_length=max_length, |
| stride=stride, |
| pad_to_multiple_of=pad_to_multiple_of, |
| return_attention_mask=return_attention_mask, |
| return_overflowing_tokens=return_overflowing_tokens, |
| return_special_tokens_mask=return_special_tokens_mask, |
| return_offsets_mapping=return_offsets_mapping, |
| return_token_type_ids=return_token_type_ids, |
| return_length=return_length, |
| verbose=verbose, |
| return_tensors=return_tensors, |
| **kwargs, |
| ) |
| else: |
| text_inputs = {} |
| |
| |
| return BatchFeature(data={**image_inputs, **text_inputs}) |
| |
| def decode(self, *args, **kwargs): |
| """Decode token ids to strings.""" |
| return self.tokenizer.decode(*args, **kwargs) |
| |
| def batch_decode(self, *args, **kwargs): |
| """Batch decode token ids to strings.""" |
| return self.tokenizer.batch_decode(*args, **kwargs) |
| |
| def post_process_generation(self, sequences, fix_markdown=False): |
| """Post-process generated sequences.""" |
| if hasattr(self.tokenizer, 'post_process_generation'): |
| return self.tokenizer.post_process_generation(sequences, fix_markdown=fix_markdown) |
| else: |
| |
| if isinstance(sequences, str): |
| sequences = [sequences] |
| |
| processed = [] |
| for seq in sequences: |
| |
| seq = seq.replace('<s>', '').replace('</s>', '').strip() |
| processed.append(seq) |
| |
| return processed[0] if len(processed) == 1 else processed |
| |
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): |
| """ |
| Load processor from pretrained model. |
| |
| This method is compatible with AutoProcessor.from_pretrained(). |
| """ |
| |
| from transformers import AutoImageProcessor, AutoTokenizer |
| trust_remote_code = kwargs.get("trust_remote_code", None) |
| revision = kwargs.get("revision", None) |
| token = kwargs.get("token", None) |
| image_processor = AutoImageProcessor.from_pretrained( |
| pretrained_model_name_or_path, |
| trust_remote_code=trust_remote_code, |
| revision=revision, |
| token=token, |
| ) |
| tokenizer = AutoTokenizer.from_pretrained( |
| pretrained_model_name_or_path, |
| trust_remote_code=trust_remote_code, |
| revision=revision, |
| token=token, |
| ) |
| return cls(image_processor=image_processor, tokenizer=tokenizer) |
| |
| def save_pretrained(self, save_directory, **kwargs): |
| """ |
| Save processor to directory. |
| |
| This method is compatible with AutoProcessor/AutoImageProcessor loading. |
| """ |
| import os |
| os.makedirs(save_directory, exist_ok=True) |
| |
| |
| print("Saving tokenizer for AutoTokenizer compatibility...") |
| self.tokenizer.save_pretrained(save_directory, **kwargs) |
| |
| |
| print("Saving image processor...") |
| self.image_processor.save_pretrained(save_directory, **kwargs) |
| |
| |
| super().save_pretrained(save_directory, **kwargs) |
| print(f"NemotronParseProcessor saved to {save_directory}") |
| print(f"AutoTokenizer.from_pretrained('{save_directory}') should now work!") |
|
|
