| """ | |
| Based upon ImageCaptionLoader in LangChain version: langchain/document_loaders/image_captions.py | |
| But accepts preloaded model to avoid slowness in use and CUDA forking issues | |
| Loader that uses Pix2Struct models to image caption | |
| """ | |
| from typing import List, Union, Any, Tuple | |
| from langchain.docstore.document import Document | |
| from langchain.document_loaders import ImageCaptionLoader | |
| from utils import get_device, clear_torch_cache | |
| from PIL import Image | |
| class H2OPix2StructLoader(ImageCaptionLoader): | |
| """Loader that extracts text from images""" | |
| def __init__(self, path_images: Union[str, List[str]] = None, model_type="google/pix2struct-textcaps-base", | |
| max_new_tokens=50): | |
| super().__init__(path_images) | |
| self._pix2struct_model = None | |
| self._model_type = model_type | |
| self._max_new_tokens = max_new_tokens | |
| def set_context(self): | |
| if get_device() == 'cuda': | |
| import torch | |
| n_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0 | |
| if n_gpus > 0: | |
| self.context_class = torch.device | |
| self.device = 'cuda' | |
| else: | |
| self.device = 'cpu' | |
| else: | |
| self.device = 'cpu' | |
| def load_model(self): | |
| try: | |
| from transformers import AutoProcessor, Pix2StructForConditionalGeneration | |
| except ImportError: | |
| raise ValueError( | |
| "`transformers` package not found, please install with " | |
| "`pip install transformers`." | |
| ) | |
| if self._pix2struct_model: | |
| self._pix2struct_model = self._pix2struct_model.to(self.device) | |
| return self | |
| self.set_context() | |
| self._pix2struct_processor = AutoProcessor.from_pretrained(self._model_type) | |
| self._pix2struct_model = Pix2StructForConditionalGeneration.from_pretrained(self._model_type).to(self.device) | |
| return self | |
| def unload_model(self): | |
| if hasattr(self._pix2struct_model, 'cpu'): | |
| self._pix2struct_model.cpu() | |
| clear_torch_cache() | |
| def set_image_paths(self, path_images: Union[str, List[str]]): | |
| """ | |
| Load from a list of image files | |
| """ | |
| if isinstance(path_images, str): | |
| self.image_paths = [path_images] | |
| else: | |
| self.image_paths = path_images | |
| def load(self, prompt=None) -> List[Document]: | |
| if self._pix2struct_model is None: | |
| self.load_model() | |
| results = [] | |
| for path_image in self.image_paths: | |
| caption, metadata = self._get_captions_and_metadata( | |
| processor=self._pix2struct_processor, model=self._pix2struct_model, path_image=path_image | |
| ) | |
| doc = Document(page_content=caption, metadata=metadata) | |
| results.append(doc) | |
| return results | |
| def _get_captions_and_metadata( | |
| self, processor: Any, model: Any, path_image: str) -> Tuple[str, dict]: | |
| """ | |
| Helper function for getting the captions and metadata of an image | |
| """ | |
| try: | |
| image = Image.open(path_image) | |
| except Exception: | |
| raise ValueError(f"Could not get image data for {path_image}") | |
| inputs = self._pix2struct_processor(images=image, return_tensors="pt") | |
| inputs = inputs.to(self.device) | |
| generated_ids = self._pix2struct_model.generate(**inputs, max_new_tokens=self._max_new_tokens) | |
| generated_text = self._pix2struct_processor.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
| metadata: dict = {"image_path": path_image} | |
| return generated_text, metadata | |