| """ |
| 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 |
|
|
