fastembed-bio 0.1.0__py3-none-any.whl

fastembed/image/transform/operators.py

@@ -0,0 +1,499 @@
+from typing import Any
+import math
+
+from PIL import Image
+
+from fastembed.common.types import NumpyArray
+from fastembed.image.transform.functional import (
+    center_crop,
+    convert_to_rgb,
+    crop_ndarray,
+    normalize,
+    pil2ndarray,
+    rescale,
+    resize,
+    resize_longest_edge,
+    resize_ndarray,
+    pad2square,
+)
+
+
+class Transform:
+    def __call__(self, images: list[Any]) -> list[Image.Image] | list[NumpyArray]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class ConvertToRGB(Transform):
+    def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
+        return [convert_to_rgb(image=image) for image in images]
+
+
+class CenterCrop(Transform):
+    def __init__(self, size: tuple[int, int]):
+        self.size = size
+
+    def __call__(self, images: list[Image.Image]) -> list[NumpyArray]:
+        return [center_crop(image=image, size=self.size) for image in images]
+
+
+class Normalize(Transform):
+    def __init__(self, mean: float | list[float], std: float | list[float]):
+        self.mean = mean
+        self.std = std
+
+    def __call__(  # type: ignore[override]
+        self, images: list[NumpyArray] | list[list[NumpyArray]]
+    ) -> list[NumpyArray] | list[list[NumpyArray]]:
+        if images and isinstance(images[0], list):
+            # Nested structure from ImageSplitter
+            return [
+                [normalize(image, mean=self.mean, std=self.std) for image in img_patches]  # type: ignore[arg-type]
+                for img_patches in images
+            ]
+        else:
+            # Flat structure (backward compatibility)
+            return [normalize(image, mean=self.mean, std=self.std) for image in images]  # type: ignore[arg-type]
+
+
+class Resize(Transform):
+    def __init__(
+        self,
+        size: int | tuple[int, int],
+        resample: Image.Resampling = Image.Resampling.BICUBIC,
+    ):
+        self.size = size
+        self.resample = resample
+
+    def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
+        return [resize(image, size=self.size, resample=self.resample) for image in images]
+
+
+class Rescale(Transform):
+    def __init__(self, scale: float = 1 / 255):
+        self.scale = scale
+
+    def __call__(  # type: ignore[override]
+        self, images: list[NumpyArray] | list[list[NumpyArray]]
+    ) -> list[NumpyArray] | list[list[NumpyArray]]:
+        if images and isinstance(images[0], list):
+            # Nested structure from ImageSplitter
+            return [
+                [rescale(image, scale=self.scale) for image in img_patches]  # type: ignore[arg-type]
+                for img_patches in images
+            ]
+        else:
+            # Flat structure (backward compatibility)
+            return [rescale(image, scale=self.scale) for image in images]  # type: ignore[arg-type]
+
+
+class PILtoNDarray(Transform):
+    def __call__(self, images: list[Image.Image | NumpyArray]) -> list[NumpyArray]:
+        return [pil2ndarray(image) for image in images]
+
+
+class PadtoSquare(Transform):
+    def __init__(
+        self,
+        size: int,
+        fill_color: str | int | tuple[int, ...],
+    ):
+        self.size = size
+        self.fill_color = fill_color
+
+    def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
+        return [
+            pad2square(image=image, size=self.size, fill_color=self.fill_color) for image in images
+        ]
+
+
+class ResizeLongestEdge(Transform):
+    """Resize images so the longest edge equals target size, preserving aspect ratio."""
+
+    def __init__(
+        self,
+        size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.size = size
+        self.resample = resample
+
+    def __call__(self, images: list[Image.Image]) -> list[Image.Image]:
+        return [resize_longest_edge(image, self.size, self.resample) for image in images]
+
+
+class ResizeForVisionEncoder(Transform):
+    """
+    Resize both dimensions to be multiples of vision_encoder_max_size.
+    Preserves aspect ratio approximately.
+    Works on numpy arrays in (C, H, W) format.
+    """
+
+    def __init__(
+        self,
+        max_size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.max_size = max_size
+        self.resample = resample
+
+    def __call__(self, images: list[NumpyArray]) -> list[NumpyArray]:
+        result = []
+        for image in images:
+            # Assume (C, H, W) format
+            _, height, width = image.shape
+
+            aspect_ratio = width / height
+
+            if width >= height:
+                # Calculate new width as multiple of max_size
+                new_width = math.ceil(width / self.max_size) * self.max_size
+                new_height = int(new_width / aspect_ratio)
+                new_height = math.ceil(new_height / self.max_size) * self.max_size
+            else:
+                # Calculate new height as multiple of max_size
+                new_height = math.ceil(height / self.max_size) * self.max_size
+                new_width = int(new_height * aspect_ratio)
+                new_width = math.ceil(new_width / self.max_size) * self.max_size
+
+            # Resize using the ndarray resize function
+            resized = resize_ndarray(
+                image,
+                size=(new_width, new_height),  # PIL expects (width, height)
+                resample=self.resample,
+                channel_first=True,
+            )
+            result.append(resized)
+
+        return result
+
+
+class ImageSplitter(Transform):
+    """
+    Split images into grid of patches plus a global view.
+
+    If image dimensions exceed max_size:
+    - Divide into ceil(H/max_size) x ceil(W/max_size) patches
+    - Each patch is cropped from the image
+    - Add a global view (original resized to max_size x max_size)
+
+    If image is smaller than max_size:
+    - Return single image unchanged
+
+    Works on numpy arrays in (C, H, W) format.
+    """
+
+    def __init__(
+        self,
+        max_size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.max_size = max_size
+        self.resample = resample
+
+    def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]:  # type: ignore[override]
+        result = []
+
+        for image in images:
+            # Assume (C, H, W) format
+            _, height, width = image.shape
+            max_height = max_width = self.max_size
+
+            frames = []
+
+            if height > max_height or width > max_width:
+                # Calculate the number of splits needed
+                num_splits_h = math.ceil(height / max_height)
+                num_splits_w = math.ceil(width / max_width)
+
+                # Calculate optimal patch dimensions
+                optimal_height = math.ceil(height / num_splits_h)
+                optimal_width = math.ceil(width / num_splits_w)
+
+                # Generate patches in grid order (row by row)
+                for r in range(num_splits_h):
+                    for c in range(num_splits_w):
+                        # Calculate crop coordinates
+                        start_x = c * optimal_width
+                        start_y = r * optimal_height
+                        end_x = min(start_x + optimal_width, width)
+                        end_y = min(start_y + optimal_height, height)
+
+                        # Crop the patch
+                        cropped = crop_ndarray(
+                            image, x1=start_x, y1=start_y, x2=end_x, y2=end_y, channel_first=True
+                        )
+                        frames.append(cropped)
+
+                # Add global view (resized to max_size x max_size)
+                global_view = resize_ndarray(
+                    image,
+                    size=(max_width, max_height),  # PIL expects (width, height)
+                    resample=self.resample,
+                    channel_first=True,
+                )
+                frames.append(global_view)
+            else:
+                # Image is small enough, no splitting needed
+                frames.append(image)
+
+            # Append (not extend) to preserve per-image grouping
+            result.append(frames)
+
+        return result
+
+
+class SquareResize(Transform):
+    """
+    Resize images to square dimensions (max_size x max_size).
+    Works on numpy arrays in (C, H, W) format.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ):
+        self.size = size
+        self.resample = resample
+
+    def __call__(self, images: list[NumpyArray]) -> list[list[NumpyArray]]:  # type: ignore[override]
+        return [
+            [
+                resize_ndarray(
+                    image, size=(self.size, self.size), resample=self.resample, channel_first=True
+                )
+            ]
+            for image in images
+        ]
+
+
+class Compose:
+    def __init__(self, transforms: list[Transform]):
+        self.transforms = transforms
+
+    def __call__(
+        self, images: list[Image.Image] | list[NumpyArray]
+    ) -> list[NumpyArray] | list[Image.Image]:
+        for transform in self.transforms:
+            images = transform(images)
+        return images
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "Compose":
+        """Creates processor from a config dict.
+        Args:
+            config (dict[str, Any]): Configuration dictionary.
+
+                Valid keys:
+                    - do_resize
+                    - resize_mode
+                    - size
+                    - fill_color
+                    - do_center_crop
+                    - crop_size
+                    - do_rescale
+                    - rescale_factor
+                    - do_normalize
+                    - image_mean
+                    - mean
+                    - image_std
+                    - std
+                    - resample
+                    - interpolation
+                Valid size keys (nested):
+                    - {"height", "width"}
+                    - {"shortest_edge"}
+                    - {"longest_edge"}
+
+        Returns:
+            Compose: Image processor.
+        """
+        transforms: list[Transform] = []
+        cls._get_convert_to_rgb(transforms, config)
+        cls._get_resize(transforms, config)
+        cls._get_pad2square(transforms, config)
+        cls._get_center_crop(transforms, config)
+        cls._get_pil2ndarray(transforms, config)
+        cls._get_image_splitting(transforms, config)
+        cls._get_rescale(transforms, config)
+        cls._get_normalize(transforms, config)
+        return cls(transforms=transforms)
+
+    @staticmethod
+    def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]) -> None:
+        transforms.append(ConvertToRGB())
+
+    @classmethod
+    def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]) -> None:
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+        if mode in ("CLIPImageProcessor", "SiglipImageProcessor"):
+            if config.get("do_resize", False):
+                size = config["size"]
+                if "shortest_edge" in size:
+                    size = size["shortest_edge"]
+                elif "height" in size and "width" in size:
+                    size = (size["height"], size["width"])
+                else:
+                    raise ValueError(
+                        "Size must contain either 'shortest_edge' or 'height' and 'width'."
+                    )
+                transforms.append(
+                    Resize(
+                        size=size,
+                        resample=config.get("resample", Image.Resampling.BICUBIC),
+                    )
+                )
+        elif mode == "ConvNextFeatureExtractor":
+            if "size" in config and "shortest_edge" not in config["size"]:
+                raise ValueError(
+                    f"Size dictionary must contain 'shortest_edge' key. Got {config['size'].keys()}"
+                )
+            shortest_edge = config["size"]["shortest_edge"]
+            crop_pct = config.get("crop_pct", 0.875)
+            if shortest_edge < 384:
+                # maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
+                resize_shortest_edge = int(shortest_edge / crop_pct)
+                transforms.append(
+                    Resize(
+                        size=resize_shortest_edge,
+                        resample=config.get("resample", Image.Resampling.BICUBIC),
+                    )
+                )
+                transforms.append(CenterCrop(size=(shortest_edge, shortest_edge)))
+            else:
+                transforms.append(
+                    Resize(
+                        size=(shortest_edge, shortest_edge),
+                        resample=config.get("resample", Image.Resampling.BICUBIC),
+                    )
+                )
+        elif mode == "JinaCLIPImageProcessor":
+            interpolation = config.get("interpolation")
+            if isinstance(interpolation, str):
+                resample = cls._interpolation_resolver(interpolation)
+            else:
+                resample = interpolation or Image.Resampling.BICUBIC
+
+            if "size" in config:
+                resize_mode = config.get("resize_mode", "shortest")
+                if resize_mode == "shortest":
+                    transforms.append(
+                        Resize(
+                            size=config["size"],
+                            resample=resample,
+                        )
+                    )
+        elif mode == "Idefics3ImageProcessor":
+            if config.get("do_resize", False):
+                size = config.get("size", {})
+                if "longest_edge" not in size:
+                    raise ValueError(
+                        "Size dictionary must contain 'longest_edge' key for Idefics3ImageProcessor"
+                    )
+
+                # Handle resample parameter - can be int enum or PIL.Image.Resampling
+                resample = config.get("resample", Image.Resampling.LANCZOS)
+                if isinstance(resample, int):
+                    resample = Image.Resampling(resample)
+
+                transforms.append(
+                    ResizeLongestEdge(
+                        size=size["longest_edge"],
+                        resample=resample,
+                    )
+                )
+        else:
+            raise ValueError(f"Preprocessor {mode} is not supported")
+
+    @staticmethod
+    def _get_center_crop(transforms: list[Transform], config: dict[str, Any]) -> None:
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+        if mode in ("CLIPImageProcessor", "SiglipImageProcessor"):
+            if config.get("do_center_crop", False):
+                crop_size_raw = config["crop_size"]
+                crop_size: tuple[int, int]
+                if isinstance(crop_size_raw, int):
+                    crop_size = (crop_size_raw, crop_size_raw)
+                elif isinstance(crop_size_raw, dict):
+                    crop_size = (crop_size_raw["height"], crop_size_raw["width"])
+                else:
+                    raise ValueError(f"Invalid crop size: {crop_size_raw}")
+                transforms.append(CenterCrop(size=crop_size))
+        elif mode == "ConvNextFeatureExtractor":
+            pass
+        elif mode == "JinaCLIPImageProcessor":
+            pass
+        elif mode == "Idefics3ImageProcessor":
+            pass
+        else:
+            raise ValueError(f"Preprocessor {mode} is not supported")
+
+    @staticmethod
+    def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]) -> None:
+        transforms.append(PILtoNDarray())
+
+    @classmethod
+    def _get_image_splitting(cls, transforms: list[Transform], config: dict[str, Any]) -> None:
+        """
+        Add image splitting transforms for Idefics3.
+        Handles conditional logic: splitting vs square resize.
+        Must be called AFTER PILtoNDarray.
+        """
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+
+        if mode == "Idefics3ImageProcessor":
+            do_splitting = config.get("do_image_splitting", False)
+            max_size = config.get("max_image_size", {}).get("longest_edge", 512)
+            resample = config.get("resample", Image.Resampling.LANCZOS)
+            if isinstance(resample, int):
+                resample = Image.Resampling(resample)
+
+            if do_splitting:
+                transforms.append(ResizeForVisionEncoder(max_size, resample))
+                transforms.append(ImageSplitter(max_size, resample))
+            else:
+                transforms.append(SquareResize(max_size, resample))
+
+    @staticmethod
+    def _get_rescale(transforms: list[Transform], config: dict[str, Any]) -> None:
+        if config.get("do_rescale", True):
+            rescale_factor = config.get("rescale_factor", 1 / 255)
+            transforms.append(Rescale(scale=rescale_factor))
+
+    @staticmethod
+    def _get_normalize(transforms: list[Transform], config: dict[str, Any]) -> None:
+        if config.get("do_normalize", False):
+            transforms.append(Normalize(mean=config["image_mean"], std=config["image_std"]))
+        elif "mean" in config and "std" in config:
+            transforms.append(Normalize(mean=config["mean"], std=config["std"]))
+
+    @staticmethod
+    def _get_pad2square(transforms: list[Transform], config: dict[str, Any]) -> None:
+        mode = config.get("image_processor_type", "CLIPImageProcessor")
+        if mode == "CLIPImageProcessor":
+            pass
+        elif mode == "ConvNextFeatureExtractor":
+            pass
+        elif mode == "JinaCLIPImageProcessor":
+            transforms.append(
+                PadtoSquare(
+                    size=config["size"],
+                    fill_color=config.get("fill_color", 0),
+                )
+            )
+
+    @staticmethod
+    def _interpolation_resolver(resample: str | None = None) -> Image.Resampling:
+        interpolation_map = {
+            "nearest": Image.Resampling.NEAREST,
+            "lanczos": Image.Resampling.LANCZOS,
+            "bilinear": Image.Resampling.BILINEAR,
+            "bicubic": Image.Resampling.BICUBIC,
+            "box": Image.Resampling.BOX,
+            "hamming": Image.Resampling.HAMMING,
+        }
+
+        if resample and (method := interpolation_map.get(resample.lower())):
+            return method
+
+        raise ValueError(f"Unknown interpolation method: {resample}")

fastembed/late_interaction/__init__.py

@@ -0,0 +1,5 @@
+from fastembed.late_interaction.late_interaction_text_embedding import (
+    LateInteractionTextEmbedding,
+)
+
+__all__ = ["LateInteractionTextEmbedding"]