codekingpro/portable-devtools

1import importlib
2import multiprocessing
3from typing import Optional, Sequence, List, Tuple
4import numpy as np
5from chromadb.api.types import URI, DataLoader, Image, URIs
6from concurrent.futures import ThreadPoolExecutor
7
8
9class ImageLoader(DataLoader[List[Optional[Image]]]):
10    def __init__(self, max_workers: int = multiprocessing.cpu_count()) -> None:
11        try:
12            self._PILImage = importlib.import_module("PIL.Image")
13            self._max_workers = max_workers
14        except ImportError:
15            raise ValueError(
16                "The PIL python package is not installed. Please install it with `pip install pillow`"
17            )
18
19    def _load_image(self, uri: Optional[URI]) -> Optional[Image]:
20        return np.array(self._PILImage.open(uri)) if uri is not None else None
21
22    def __call__(self, uris: Sequence[Optional[URI]]) -> List[Optional[Image]]:
23        with ThreadPoolExecutor(max_workers=self._max_workers) as executor:
24            return list(executor.map(self._load_image, uris))
25
26
27class ChromaLangchainPassthroughDataLoader(DataLoader[List[Optional[Image]]]):
28    # This is a simple pass through data loader that just returns the input data with "images"
29    # flag which lets the langchain embedding function know that the data is image uris
30    def __call__(self, uris: URIs) -> Tuple[str, URIs]:  # type: ignore
31        return ("images", uris)
32