upstream-archive/fastembed
1
0
Fork 0
mirror of https://github.com/qdrant/fastembed.git synced 2026-07-23 11:20:51 -05:00
Code Issues Packages Projects Releases Wiki Activity

Opened images support (#315)

Browse Source 
* Opened image support

---------

Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>
This commit is contained in:
Dmitrii Ogn
2024-07-31 13:23:17 +03:00
committed by GitHub
parent 0e258ab875
commit 9c72d2f59f
5 changed files with 49 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
fastembed/common/__init__.py
View File

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

4
fastembed/common/types.py
View File

@@ -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]]]

13
fastembed/image/onnx_embedding.py
View File

@@ -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)

19
fastembed/image/onnx_image_model.py
View File

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

42
tests/test_image_onnx_embeddings.py
View File

@@ -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)