forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
"""
|
|
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 ImportError(
|
|
"`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 ImportError(
|
|
"`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
|