dgc-image-embedding 1.0.0__tar.gz

dgc_image_embedding-1.0.0/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: dgc_image_embedding
+Version: 1.0.0
+Author-email: velopace <velopace7@gmail.com>
+License: MIT
+Description-Content-Type: text/markdown
+Requires-Dist: pillow==12.2.0
+Requires-Dist: torch==2.12.0
+Requires-Dist: torchvision==0.27.0
+Requires-Dist: open-clip-torch==3.3.0
+Provides-Extra: dev
+Requires-Dist: build>=1.3.0; extra == "dev"
+Requires-Dist: coverage>=7.10.2; extra == "dev"
+Requires-Dist: pytest>=8.4.1; extra == "dev"
+Requires-Dist: pytest-cov>=6.2.1; extra == "dev"
+Requires-Dist: twine>=6.1.0; extra == "dev"
+Requires-Dist: toml>=0.10.2; extra == "dev"

dgc_image_embedding-1.0.0/dgc_image_embedding/__init__.py

@@ -0,0 +1,13 @@
+from .embedding import embedding_metadata, generate_embedding, generate_embeddings_batch, get_device, get_model
+from .image import load_image_from_bytes, prepare_image_from_bytes, resize_and_pad
+
+__all__ = [
+    "embedding_metadata",
+    "generate_embedding",
+    "generate_embeddings_batch",
+    "get_device",
+    "get_model",
+    "load_image_from_bytes",
+    "prepare_image_from_bytes",
+    "resize_and_pad",
+]

dgc_image_embedding-1.0.0/dgc_image_embedding/config.py

@@ -0,0 +1,13 @@
+TARGET_IMAGE_SIZE = 224
+
+MAX_FILE_SIZE_BYTES = 5 * 1024 * 1024
+MIN_SOURCE_SIDE = 100
+MAX_SOURCE_SIDE = 3000
+
+PAD_COLOR = (255, 255, 255)
+
+MODEL_NAME = "ViT-B-32"
+PRETRAINED = "laion2b_s34b_b79k"
+
+CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
+CLIP_STD = (0.26862954, 0.26130258, 0.27577711)

dgc_image_embedding-1.0.0/dgc_image_embedding/embedding.py

@@ -0,0 +1,89 @@
+import threading
+
+import open_clip
+import torch
+from PIL import Image
+from torchvision import transforms
+
+from .config import (
+    CLIP_MEAN,
+    CLIP_STD,
+    MODEL_NAME,
+    PRETRAINED,
+)
+
+
+_model = None
+_device = "cuda" if torch.cuda.is_available() else "cpu"
+_model_lock = threading.Lock()
+
+
+def get_device() -> str:
+    return _device
+
+
+def get_model():
+    global _model
+
+    if _model is not None:
+        return _model
+
+    with _model_lock:
+        if _model is None:
+            model, _, _ = open_clip.create_model_and_transforms(
+                MODEL_NAME,
+                pretrained=PRETRAINED,
+                device=_device,
+            )
+            model.eval()
+            _model = model
+
+    return _model
+
+
+def image_to_tensor(image: Image.Image) -> torch.Tensor:
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD),
+        ]
+    )
+
+    return transform(image).unsqueeze(0)
+
+
+def generate_embedding(processed_image: Image.Image) -> list[float]:
+    model = get_model()
+    tensor = image_to_tensor(processed_image).to(_device)
+
+    with torch.no_grad():
+        embedding = model.encode_image(tensor)
+        embedding = embedding / embedding.norm(dim=-1, keepdim=True)
+
+    return embedding.squeeze(0).cpu().tolist()
+
+
+def generate_embeddings_batch(processed_images: list[Image.Image]) -> list[list[float]]:
+    if not processed_images:
+        return []
+
+    model = get_model()
+
+    tensors = torch.cat(
+        [image_to_tensor(image) for image in processed_images],
+        dim=0,
+    ).to(_device)
+
+    with torch.no_grad():
+        embeddings = model.encode_image(tensors)
+        embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)
+
+    return embeddings.cpu().tolist()
+
+
+def embedding_metadata() -> dict:
+    return {
+        "modelName": MODEL_NAME,
+        "pretrained": PRETRAINED,
+        "device": _device,
+    }

dgc_image_embedding-1.0.0/dgc_image_embedding/image.py

@@ -0,0 +1,73 @@
+from io import BytesIO
+
+from PIL import Image, ImageOps, UnidentifiedImageError
+
+from .config import (
+    MAX_FILE_SIZE_BYTES,
+    MAX_SOURCE_SIDE,
+    MIN_SOURCE_SIDE,
+    PAD_COLOR,
+    TARGET_IMAGE_SIZE,
+)
+
+
+def load_image_from_bytes(image_bytes: bytes) -> Image.Image:
+    if not image_bytes:
+        raise ValueError("Image file is empty")
+
+    if len(image_bytes) > MAX_FILE_SIZE_BYTES:
+        raise ValueError("Image file is too large")
+
+    try:
+        image = Image.open(BytesIO(image_bytes))
+        image = ImageOps.exif_transpose(image)
+    except UnidentifiedImageError as err:
+        raise ValueError("Invalid image file") from err
+
+    validate_image_size(image)
+
+    return image
+
+
+def validate_image_size(image: Image.Image) -> None:
+    width, height = image.size
+
+    if width < MIN_SOURCE_SIDE or height < MIN_SOURCE_SIDE:
+        raise ValueError(f"Image is too small: {width}x{height}")
+
+    if width > MAX_SOURCE_SIDE or height > MAX_SOURCE_SIDE:
+        raise ValueError(f"Image is too large: {width}x{height}")
+
+
+def resize_and_pad(
+    image: Image.Image,
+    target_size: int = TARGET_IMAGE_SIZE,
+    pad_color: tuple[int, int, int] = PAD_COLOR,
+) -> Image.Image:
+    image = ImageOps.exif_transpose(image)
+    image = image.convert("RGB")
+
+    width, height = image.size
+    scale = target_size / max(width, height)
+
+    new_width = max(1, round(width * scale))
+    new_height = max(1, round(height * scale))
+
+    resized = image.resize(
+        (new_width, new_height),
+        Image.Resampling.LANCZOS,
+    )
+
+    padded = Image.new("RGB", (target_size, target_size), pad_color)
+
+    paste_x = (target_size - new_width) // 2
+    paste_y = (target_size - new_height) // 2
+
+    padded.paste(resized, (paste_x, paste_y))
+
+    return padded
+
+
+def prepare_image_from_bytes(image_bytes: bytes) -> Image.Image:
+    image = load_image_from_bytes(image_bytes)
+    return resize_and_pad(image)

dgc_image_embedding-1.0.0/dgc_image_embedding.egg-info/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: dgc_image_embedding
+Version: 1.0.0
+Author-email: velopace <velopace7@gmail.com>
+License: MIT
+Description-Content-Type: text/markdown
+Requires-Dist: pillow==12.2.0
+Requires-Dist: torch==2.12.0
+Requires-Dist: torchvision==0.27.0
+Requires-Dist: open-clip-torch==3.3.0
+Provides-Extra: dev
+Requires-Dist: build>=1.3.0; extra == "dev"
+Requires-Dist: coverage>=7.10.2; extra == "dev"
+Requires-Dist: pytest>=8.4.1; extra == "dev"
+Requires-Dist: pytest-cov>=6.2.1; extra == "dev"
+Requires-Dist: twine>=6.1.0; extra == "dev"
+Requires-Dist: toml>=0.10.2; extra == "dev"

dgc_image_embedding-1.0.0/dgc_image_embedding.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+pyproject.toml
+dgc_image_embedding/__init__.py
+dgc_image_embedding/config.py
+dgc_image_embedding/embedding.py
+dgc_image_embedding/image.py
+dgc_image_embedding.egg-info/PKG-INFO
+dgc_image_embedding.egg-info/SOURCES.txt
+dgc_image_embedding.egg-info/dependency_links.txt
+dgc_image_embedding.egg-info/requires.txt
+dgc_image_embedding.egg-info/top_level.txt

dgc_image_embedding-1.0.0/dgc_image_embedding.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

dgc_image_embedding-1.0.0/dgc_image_embedding.egg-info/requires.txt

@@ -0,0 +1,12 @@
+pillow==12.2.0
+torch==2.12.0
+torchvision==0.27.0
+open-clip-torch==3.3.0
+
+[dev]
+build>=1.3.0
+coverage>=7.10.2
+pytest>=8.4.1
+pytest-cov>=6.2.1
+twine>=6.1.0
+toml>=0.10.2

dgc_image_embedding-1.0.0/dgc_image_embedding.egg-info/top_level.txt

@@ -0,0 +1 @@
+dgc_image_embedding

dgc_image_embedding-1.0.0/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = [ "setuptools", "wheel",]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "dgc_image_embedding"
+version = "1.0.0"
+description = ""
+readme = "README.md"
+dependencies = [ "pillow==12.2.0", "torch==2.12.0", "torchvision==0.27.0", "open-clip-torch==3.3.0",]
+[[project.authors]]
+name = "velopace"
+email = "velopace7@gmail.com"
+
+[project.optional-dependencies]
+dev = [ "build>=1.3.0", "coverage>=7.10.2", "pytest>=8.4.1", "pytest-cov>=6.2.1", "twine>=6.1.0", "toml>=0.10.2",]
+
+[project.license]
+text = "MIT"
+
+[tool.setuptools.package-data]
+dgc_image_embedding = [ "py.typed",]
+
+[tool.pytest.ini_options]
+testpaths = [ "tests",]
+
+[tool.coverage.run]
+branch = true
+source = [ "dgc_image_embedding",]
+
+[tool.coverage.report]
+show_missing = true
+skip_covered = true
+
+[tool.setuptools.packages.find]
+where = [ ".",]
+include = [ "dgc_image_embedding*",]

dgc_image_embedding-1.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+