diff --git a/fastembed/common/__init__.py b/fastembed/common/__init__.py index 5caf45f..b749189 100644 --- a/fastembed/common/__init__.py +++ b/fastembed/common/__init__.py @@ -1,3 +1,3 @@ -from fastembed.common.types import ImageInput, OnnxProvider, PathInput +from fastembed.common.types import ImageInput, OnnxProvider, PathInput, PilInput -__all__ = ["OnnxProvider", "ImageInput", "PathInput"] +__all__ = ["OnnxProvider", "ImageInput", "PathInput", "PilInput"] diff --git a/fastembed/common/types.py b/fastembed/common/types.py index c428a3f..1c6f39a 100644 --- a/fastembed/common/types.py +++ b/fastembed/common/types.py @@ -1,5 +1,6 @@ import os import sys +from PIL import Image from typing import Any, Dict, Iterable, Tuple, Union if sys.version_info >= (3, 10): @@ -9,6 +10,7 @@ else: PathInput: TypeAlias = Union[str, os.PathLike] -ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput]] +PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]] +ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput] OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]] diff --git a/fastembed/image/onnx_embedding.py b/fastembed/image/onnx_embedding.py index c39e9ca..b6e082d 100644 --- a/fastembed/image/onnx_embedding.py +++ b/fastembed/image/onnx_embedding.py @@ -51,6 +51,7 @@ supported_onnx_models = [ }, ] + class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]): def __init__( self, @@ -141,16 +142,10 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]): return onnx_input - def _post_process_onnx_output( - self, output: OnnxOutputContext - ) -> Iterable[np.ndarray]: + def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]: return normalize(output.model_output).astype(np.float32) class OnnxImageEmbeddingWorker(ImageEmbeddingWorker): - def init_embedding( - self, model_name: str, cache_dir: str, **kwargs - ) -> OnnxImageEmbedding: - return OnnxImageEmbedding( - model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs - ) + def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding: + return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs) diff --git a/fastembed/image/onnx_image_model.py b/fastembed/image/onnx_image_model.py index cfaa152..05dc36c 100644 --- a/fastembed/image/onnx_image_model.py +++ b/fastembed/image/onnx_image_model.py @@ -7,7 +7,7 @@ from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type import numpy as np from PIL import Image -from fastembed.common import ImageInput, OnnxProvider, PathInput +from fastembed.common import ImageInput, OnnxProvider from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T from fastembed.common.preprocessor_utils import load_preprocessor from fastembed.common.utils import iter_batch @@ -54,9 +54,12 @@ class OnnxImageModel(OnnxModel[T]): def _build_onnx_input(self, encoded: np.ndarray) -> Dict[str, np.ndarray]: return {node.name: encoded for node in self.model.get_inputs()} - def onnx_embed(self, images: List[PathInput], **kwargs) -> OnnxOutputContext: + def onnx_embed(self, images: List[ImageInput], **kwargs) -> OnnxOutputContext: with contextlib.ExitStack(): - image_files = [Image.open(image) for image in images] + image_files = [ + Image.open(image) if not isinstance(image, Image.Image) else image + for image in images + ] encoded = self.processor(image_files) onnx_input = self._build_onnx_input(encoded) onnx_input = self._preprocess_onnx_input(onnx_input) @@ -75,7 +78,11 @@ class OnnxImageModel(OnnxModel[T]): ) -> Iterable[T]: is_small = False - if isinstance(images, str) or isinstance(images, Path): + if ( + isinstance(images, str) + or isinstance(images, Path) + or (isinstance(images, Image.Image)) + ): images = [images] is_small = True @@ -90,9 +97,7 @@ class OnnxImageModel(OnnxModel[T]): for batch in iter_batch(images, batch_size): yield from self._post_process_onnx_output(self.onnx_embed(batch)) else: - start_method = ( - "forkserver" if "forkserver" in get_all_start_methods() else "spawn" - ) + start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn" params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs} pool = ParallelWorkerPool( parallel, self._get_worker_class(), start_method=start_method diff --git a/tests/test_image_onnx_embeddings.py b/tests/test_image_onnx_embeddings.py index bc6a4d0..be3703c 100644 --- a/tests/test_image_onnx_embeddings.py +++ b/tests/test_image_onnx_embeddings.py @@ -1,7 +1,10 @@ import os +from io import BytesIO import numpy as np import pytest +import requests +from PIL import Image from fastembed import ImageEmbedding from tests.config import TEST_MISC_DIR @@ -12,12 +15,10 @@ CANONICAL_VECTOR_VALUES = { [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.01046245, 0.01171397, 0.00705971, 0.0] ), "Qdrant/Unicom-ViT-B-16": np.array( - [ 0.0170, -0.0361, 0.0125, -0.0428, -0.0232, 0.0232, -0.0602, -0.0333, - 0.0155, 0.0497] + [0.0170, -0.0361, 0.0125, -0.0428, -0.0232, 0.0232, -0.0602, -0.0333, 0.0155, 0.0497] ), "Qdrant/Unicom-ViT-B-32": np.array( - [0.0418, 0.0550, 0.0003, 0.0253, -0.0185, 0.0016, -0.0368, -0.0402, - -0.0891, -0.0186] + [0.0418, 0.0550, 0.0003, 0.0253, -0.0185, 0.0016, -0.0368, -0.0402, -0.0891, -0.0186] ), } @@ -33,10 +34,15 @@ def test_embedding(): model = ImageEmbedding(model_name=model_desc["model"]) - images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] + images = [ + TEST_MISC_DIR / "image.jpeg", + str(TEST_MISC_DIR / "small_image.jpeg"), + Image.open((TEST_MISC_DIR / "small_image.jpeg")), + Image.open(BytesIO(requests.get("https://qdrant.tech/img/logo.png").content)), + ] embeddings = list(model.embed(images)) embeddings = np.stack(embeddings, axis=0) - assert embeddings.shape == (2, dim) + assert embeddings.shape == (len(images), dim) canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]] @@ -44,19 +50,24 @@ def test_embedding(): embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3 ), model_desc["model"] + assert np.allclose(embeddings[1], embeddings[2]), model_desc["model"] + @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")]) def test_batch_embedding(n_dims, model_name): model = ImageEmbedding(model_name=model_name) n_images = 32 - images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] * ( - n_images // 2 - ) + test_images = [ + TEST_MISC_DIR / "image.jpeg", + str(TEST_MISC_DIR / "small_image.jpeg"), + Image.open(TEST_MISC_DIR / "small_image.jpeg"), + ] + images = test_images * n_images embeddings = list(model.embed(images, batch_size=10)) embeddings = np.stack(embeddings, axis=0) - assert embeddings.shape == (n_images, n_dims) + assert embeddings.shape == (len(test_images) * n_images, n_dims) @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")]) @@ -64,9 +75,12 @@ def test_parallel_processing(n_dims, model_name): model = ImageEmbedding(model_name=model_name) n_images = 32 - images = [TEST_MISC_DIR / "image.jpeg", str(TEST_MISC_DIR / "small_image.jpeg")] * ( - n_images // 2 - ) + test_images = [ + TEST_MISC_DIR / "image.jpeg", + str(TEST_MISC_DIR / "small_image.jpeg"), + Image.open(TEST_MISC_DIR / "small_image.jpeg"), + ] + images = test_images * n_images embeddings = list(model.embed(images, batch_size=10, parallel=2)) embeddings = np.stack(embeddings, axis=0) @@ -76,6 +90,6 @@ def test_parallel_processing(n_dims, model_name): embeddings_3 = list(model.embed(images, batch_size=10, parallel=0)) embeddings_3 = np.stack(embeddings_3, axis=0) - assert embeddings.shape == (n_images, n_dims) + assert embeddings.shape == (n_images * len(test_images), n_dims) assert np.allclose(embeddings, embeddings_2, atol=1e-3) assert np.allclose(embeddings, embeddings_3, atol=1e-3)