# Copyright 2024 The Feast Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Image processing utilities for Feast image search capabilities.
Provides image embedding generation and combination functions for multi-modal search.
"""
import io
from typing import List
try:
import timm
import torch
from PIL import Image
from sklearn.preprocessing import normalize
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform
_image_dependencies_available = True
except ImportError:
_image_dependencies_available = False
COMBINATION_STRATEGIES = ["weighted_sum", "concatenate", "average"]
def _check_image_dependencies():
"""Check if image processing dependencies are available."""
if not _image_dependencies_available:
raise ImportError(
"Image processing dependencies are not installed. "
"Please install with: pip install feast[image]"
)
class ImageFeatureExtractor:
"""
Extract image embeddings using pre-trained vision models.
This class uses timm (PyTorch Image Models) to generate embeddings
from images using pre-trained vision models like ResNet, ViT, etc.
Examples:
Basic usage::
extractor = ImageFeatureExtractor()
with open("image.jpg", "rb") as f:
image_bytes = f.read()
embedding = extractor.extract_embedding(image_bytes)
Using different models::
# ResNet-50
extractor = ImageFeatureExtractor("resnet50")
embedding = extractor.extract_embedding(image_bytes)
# ViT model
extractor = ImageFeatureExtractor("vit_base_patch16_224")
embedding = extractor.extract_embedding(image_bytes)
"""
def __init__(self, model_name: str = "resnet34"):
"""
Initialize with a pre-trained model.
Args:
model_name: Model name from timm library. Popular choices:
- "resnet34": Fast, good for general use (default)
- "resnet50": Better accuracy than ResNet-34
- "vit_base_patch16_224": Vision Transformer, high accuracy
- "efficientnet_b0": Good balance of speed and accuracy
- "mobilenetv3_large_100": Fast inference for mobile/edge
Raises:
ImportError: If image processing dependencies are not installed
RuntimeError: If the specified model cannot be loaded
"""
_check_image_dependencies()
try:
self.model_name = model_name
self.model = timm.create_model(
model_name, pretrained=True, num_classes=0, global_pool="avg"
)
self.model.eval()
config = resolve_data_config({}, model=model_name)
self.preprocess = create_transform(**config)
except Exception as e:
raise RuntimeError(f"Failed to load model '{model_name}': {e}")
def extract_embedding(self, image_bytes: bytes) -> List[float]:
"""
Extract embedding from image bytes.
Args:
image_bytes: Image data as bytes (JPEG, PNG, WebP, etc.)
Returns:
Normalized embedding vector as list of floats
Raises:
ValueError: If image cannot be processed or is invalid
"""
try:
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
input_tensor = self.preprocess(image).unsqueeze(0)
with torch.no_grad():
output = self.model(input_tensor)
feature_vector = output.squeeze().numpy()
normalized = normalize(feature_vector.reshape(1, -1), norm="l2")
return normalized.flatten().tolist()
except Exception as e:
raise ValueError(f"Failed to extract embedding from image: {e}")
def batch_extract_embeddings(
self, image_bytes_list: List[bytes]
) -> List[List[float]]:
"""
Extract embeddings from multiple images in batch for efficiency.
Args:
image_bytes_list: List of image data as bytes
Returns:
List of normalized embedding vectors
Raises:
ValueError: If any image cannot be processed
"""
embeddings = []
images = []
for image_bytes in image_bytes_list:
try:
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
preprocessed = self.preprocess(image)
images.append(preprocessed)
except Exception as e:
raise ValueError(f"Failed to preprocess image: {e}")
batch_tensor = torch.stack(images)
with torch.no_grad():
outputs = self.model(batch_tensor)
for output in outputs:
feature_vector = output.numpy()
normalized = normalize(feature_vector.reshape(1, -1), norm="l2")
embeddings.append(normalized.flatten().tolist())
return embeddings
def combine_embeddings(
text_embedding: List[float],
image_embedding: List[float],
strategy: str = "weighted_sum",
text_weight: float = 0.5,
image_weight: float = 0.5,
) -> List[float]:
"""
Combine text and image embeddings search.
This function provides several strategies for combining embeddings from
different modalities (text and image) into a single vector for search.
Args:
text_embedding: Text embedding vector
image_embedding: Image embedding vector
strategy: Combination strategy (default: "weighted_sum")
text_weight: Weight for text embedding (for weighted strategies)
image_weight: Weight for image embedding (for weighted strategies)
Returns:
Combined embedding vector as list of floats
Raises:
ValueError: If weights don't sum to 1.0 for weighted_sum strategy
Examples:
Weighted combination (emphasize image)::
combined = combine_embeddings(
[0.1, 0.2], [0.8, 0.9], # text_emb, image_emb
strategy="weighted_sum",
text_weight=0.3, image_weight=0.7
)
Concatenation for full information::
combined = combine_embeddings(
[0.1, 0.2], [0.8, 0.9], # text_emb, image_emb
strategy="concatenate"
)
"""
if strategy == "weighted_sum":
if abs(text_weight + image_weight - 1.0) > 1e-6:
raise ValueError(
"text_weight + image_weight must equal 1.0 for weighted_sum"
)
max_dim = max(len(text_embedding), len(image_embedding))
text_padded = text_embedding + [0.0] * (max_dim - len(text_embedding))
image_padded = image_embedding + [0.0] * (max_dim - len(image_embedding))
combined = [
text_weight * t + image_weight * i
for t, i in zip(text_padded, image_padded)
]
return combined
elif strategy == "concatenate":
return text_embedding + image_embedding
elif strategy == "average":
max_dim = max(len(text_embedding), len(image_embedding))
text_padded = text_embedding + [0.0] * (max_dim - len(text_embedding))
image_padded = image_embedding + [0.0] * (max_dim - len(image_embedding))
combined = [(t + i) / 2.0 for t, i in zip(text_padded, image_padded)]
return combined
else:
raise ValueError(
f"Unknown combination strategy: {strategy}. "
f"Supported strategies: {', '.join(COMBINATION_STRATEGIES)}"
)
def validate_image_format(image_bytes: bytes) -> bool:
"""
Validate that the provided bytes represent a valid image.
Args:
image_bytes: Image data as bytes
Returns:
True if valid image, False otherwise
"""
try:
with Image.open(io.BytesIO(image_bytes)) as img:
img.verify()
return True
except Exception:
return False
def get_image_metadata(image_bytes: bytes) -> dict:
"""
Extract metadata from image bytes.
Args:
image_bytes: Image data as bytes
Returns:
Dictionary with image metadata (format, size, mode, etc.)
Raises:
ValueError: If image cannot be processed
"""
try:
with Image.open(io.BytesIO(image_bytes)) as img:
return {
"format": img.format,
"mode": img.mode,
"width": img.width,
"height": img.height,
"size_bytes": len(image_bytes),
}
except Exception as e:
raise ValueError(f"Failed to extract image metadata: {e}")