Skip to main content
`Ctrl`+`K`

#

# Source code for langchain.document_loaders.image_captions

    """
    Loader that loads image captions
    By default, the loader utilizes the pre-trained BLIP image captioning model.
    https://huggingface.co/Salesforce/blip-image-captioning-base

    """
    from typing import Any, List, Tuple, Union

    import requests

    from langchain.docstore.document import Document
    from langchain.document_loaders.base import BaseLoader

    class ImageCaptionLoader(BaseLoader):
        """Loader that loads the captions of an image"""

        def __init__(
            self,
            path_images: Union[str, List[str]],
            blip_processor: str = "Salesforce/blip-image-captioning-base",
            blip_model: str = "Salesforce/blip-image-captioning-base",
        ):
            """
            Initialize with a list of image paths
            """
            if isinstance(path_images, str):
                self.image_paths = [path_images]
            else:
                self.image_paths = path_images

            self.blip_processor = blip_processor
            self.blip_model = blip_model

        def load(self) -> List[Document]:
            """
            Load from a list of image files
            """
            try:
                from transformers import BlipForConditionalGeneration, BlipProcessor
            except ImportError:
                raise ValueError(
                    "transformers package not found, please install with"
                    "`pip install transformers`"
                )

            processor = BlipProcessor.from_pretrained(self.blip_processor)
            model = BlipForConditionalGeneration.from_pretrained(self.blip_model)

            results = []
            for path_image in self.image_paths:
                caption, metadata = self._get_captions_and_metadata(
                    model=model, processor=processor, path_image=path_image
                )
                doc = Document(page_content=caption, metadata=metadata)
                results.append(doc)

            return results

        def _get_captions_and_metadata(
            self, model: Any, processor: Any, path_image: str
        ) -> Tuple[str, dict]:
            """
            Helper function for getting the captions and metadata of an image
            """
            try:
                from PIL import Image
            except ImportError:
                raise ValueError(
                    "PIL package not found, please install with `pip install pillow`"
                )

            try:
                if path_image.startswith("http://") or path_image.startswith("https://"):
                    image = Image.open(requests.get(path_image, stream=True).raw).convert(
                        "RGB"
                    )
                else:
                    image = Image.open(path_image).convert("RGB")
            except Exception:
                raise ValueError(f"Could not get image data for {path_image}")

            inputs = processor(image, "an image of", return_tensors="pt")
            output = model.generate(**inputs)

            caption: str = processor.decode(output[0])
            metadata: dict = {"image_path": path_image}

            return caption, metadata