Coverage for packages/dqm-ml-pytorch/src/dqm_ml_pytorch/image_embedding.py: 82%
286 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-07-21 08:27 +0000
1"""Image embedding processor using pre-trained deep learning models.
3This module contains the ImageEmbeddingProcessor class that extracts
4high-dimensional embeddings from images using PyTorch and torchvision
5pre-trained models.
6"""
8from __future__ import annotations
10import io
11import logging
12from pathlib import Path
13from typing import Any
14import warnings
16from dqm_ml_core import FeaturesProcessor
17from dqm_ml_core.models.columns import ColumnsConfig
18from dqm_ml_core.models.processors import FeaturesEmbeddingsProcessorConfig
19from dqm_ml_core.utils.image_loading import ImageLoadingMixin
20from dqm_ml_core.utils.matching import resolve_include_exclude
21import numpy as np
22from PIL import Image
23import pyarrow as pa
24import torch
25import torchvision
26from torchvision import transforms
27from torchvision.models.feature_extraction import create_feature_extractor
29# COMPATIBILITY : from typing import Any, override # When support of 3.10 and 3.11 will be removed
30from typing_extensions import override
32logger = logging.getLogger(__name__)
35class ImageEmbeddingProcessor(ImageLoadingMixin, FeaturesProcessor):
36 """
37 Computes high-dimensional latent vectors (embeddings) for images
38 using deep learning models.
40 This processor uses PyTorch and Torchvision to:
41 1. Load images from bytes or file paths.
42 2. Preprocess images (resize, normalize) for the selected model.
43 3. Run batch inference using a pre-trained model (e.g., ResNet, ViT).
44 4. Extract features from a specific layer (e.g., 'avgpool').
46 The resulting embeddings are stored as a `FixedSizeListArray`
47 in the features.
48 """
50 def __init__(
51 self,
52 name: str = "image_embedding",
53 config: dict[str, Any] | None = None,
54 ):
55 """
56 Initialize the image embedding processor.
58 Args:
59 name: Unique name of the processor instance.
60 config: Configuration dictionary containing:
61 - infer:
62 - width, height: Input resolution for the model (default: 224x224).
63 - batch_size: Number of images per inference pass (default: 32).
64 - norm_mean, norm_std: Preprocessing normalization stats.
65 - model:
66 - arch: Torchvision model name (default: "resnet18").
67 - n_layer_feature: Target layer for feature extraction (default: "avgpool").
68 - device: Execution device, "cpu" or "cuda" (default: "cpu").
69 """
70 super().__init__(name, config)
72 self.columns_config: ColumnsConfig | None = None
73 raw_columns = self.config.get("columns")
74 if isinstance(raw_columns, dict):
75 self.columns_config = ColumnsConfig.model_validate(raw_columns)
77 cfg = FeaturesEmbeddingsProcessorConfig.model_validate({**self.config, "name": self.name})
79 # Storage filesystem support
80 self.s3_fs = None
81 storage_cfg = self.storage_raw
82 if storage_cfg: 82 ↛ 83line 82 didn't jump to line 83 because the condition on line 82 was never true
83 from dqm_ml_core.models.global_ import StorageConfig
84 from dqm_ml_job.utils import get_s3_filesystem
86 storage_config = StorageConfig.model_validate(storage_cfg)
87 if storage_config.type == "s3":
88 self.s3_fs = get_s3_filesystem(storage_config)
90 self.size: tuple[int, int] = (cfg.infer.width, cfg.infer.height)
91 self.batch_size: int = cfg.infer.batch_size
92 self.arch: str = cfg.model.arch
93 n_layer_feature = cfg.model.n_layer_feature
95 # Multi-layer support for CMD: n_layer_feature can be a list
96 if isinstance(n_layer_feature, list):
97 self.multi_layer = True
98 self.target_layers: list[str] = n_layer_feature
99 self.target_layer: Any = n_layer_feature
100 self._embed_dims: dict[str, int] = {}
101 else:
102 self.multi_layer = False
103 self.target_layer = n_layer_feature
104 self._embed_dim: int | None = None
106 # Build transform (fast, no model needed)
107 safe_std = [s if s != 0 else 1e-12 for s in cfg.infer.norm_std]
108 self.transform = transforms.Compose(
109 [
110 transforms.Resize(self.size),
111 transforms.ToTensor(),
112 transforms.Normalize(mean=cfg.infer.norm_mean, std=safe_std),
113 ]
114 )
116 # Model and extractor — loaded lazily by _ensure_model_loaded()
117 self.model: Any = None
118 self.feature_extractor: Any = None
119 self.device = "cpu"
120 self._model_loaded = False
122 def _ensure_model_loaded(self) -> None:
123 """Load the PyTorch model and create the feature extractor.
125 This is deferred from ``__init__`` because:
126 - Model loading is expensive (download + GPU allocation).
127 - ``compute_device`` is injected by DatasetJob after __init__.
128 """
129 if self._model_loaded:
130 return
131 cfg = FeaturesEmbeddingsProcessorConfig.model_validate({**self.config, "name": self.name})
132 compute_device = getattr(self, "compute_device", None)
133 self.device = self._resolve_device(compute_device) if compute_device else self._resolve_device(cfg.model.device)
134 self.model = self._load_model(self.arch, self.device)
135 self.feature_extractor = self._make_extractor(self.model, self.target_layer)
136 self._model_loaded = True
138 def check_config(self) -> None:
139 """Validate configuration and load model.
141 Kept for backward compatibility. Delegates to ``_ensure_model_loaded``.
142 """
143 self._ensure_model_loaded()
145 @override
146 def needed_columns(self) -> list[str]:
147 """Return the list of columns required for image embedding extraction.
149 Returns:
150 List of input column names.
151 """
152 return self.input_columns or []
154 def _output_column_name(self, col: str, base: str) -> str:
155 """Generate output column name with prefix and suffix.
157 Args:
158 col: Input column name.
159 base: Base feature name (e.g., "embedding", "emb_layer1").
161 Returns:
162 Fully qualified output column name with prefix and suffix applied.
163 """
164 return super()._resolve_output_name(col, base)
166 def generated_columns(self) -> list[str]:
167 """Return the list of columns generated by this processor.
169 For multi-layer mode, returns one column per layer per input column.
170 For single-layer mode, returns one embedding column per input column.
172 Returns:
173 A list of column names.
174 """
175 if not self.input_columns:
176 return []
177 cols: list[str] = []
178 for col in self.input_columns:
179 if getattr(self, "multi_layer", False):
180 for layer in self.target_layers:
181 layer_base = f"emb_{layer.replace('.', '_')}"
182 cols.append(self._output_column_name(col, layer_base))
183 cols.append(self._output_column_name(col, f"{layer_base}_channels"))
184 else:
185 cols.append(self._output_column_name(col, "embedding"))
186 return cols
188 def _open_image(self, image_data: Any, column: str) -> Image.Image:
189 """Open a PIL Image from bytes, S3 path, or local filesystem path."""
190 if isinstance(image_data, (bytes, bytearray)):
191 return Image.open(io.BytesIO(image_data)).convert("RGB")
192 img = self._open_image_from_path(image_data, column)
193 assert img is not None
194 return img
196 @override
197 def _open_image_from_path(self, path: str, column: str | None = None) -> Image.Image | None:
198 """Open a PIL Image from an S3 or local filesystem path."""
199 prefix = self._current_image_prefix(column)
200 if prefix is not None and self.s3_fs: 200 ↛ 201line 200 didn't jump to line 201 because the condition on line 200 was never true
201 return self._open_s3_image(prefix, path)
202 full_path = Path(prefix) / path if prefix else Path(path)
203 return Image.open(full_path).convert("RGB")
205 def _handle_load_error(self, exc: Exception, idx: int) -> None:
206 """Check error config and either raise or record the failure."""
207 self._check_image_fail_fast(exc, "on_decode_failure", "on_transform_error")
208 self._failure_count += 1
209 self._total_count += 1
210 self._check_failure_rate()
211 logger.warning(f"[ImageEmbeddingProcessor] failed to load image: {exc}")
213 def _load_single_tensor(self, image_data: Any, column: str, idx: int) -> torch.Tensor | None:
214 """Load, transform, and return a single image tensor (or None on failure)."""
215 if image_data is None: 215 ↛ 216line 215 didn't jump to line 216 because the condition on line 215 was never true
216 return None
217 try:
218 pil_image = self._open_image(image_data, column)
219 return self.transform(pil_image) # type: ignore[no-any-return]
220 except Exception as e:
221 self._handle_load_error(e, idx)
222 return None
224 def _load_image_tensors(
225 self,
226 image_values: list[Any],
227 column: str = "",
228 ) -> list[torch.Tensor | None]:
229 """Load and transform images from a list of raw image values.
231 Auto-detects between bytes and path based on Python type.
233 Args:
234 image_values: List of raw image column values.
235 column: The input column name (used to resolve path prefix).
237 Returns:
238 List of preprocessed image tensors (or None for failed loads).
239 """
240 return [self._load_single_tensor(v, column, idx) for idx, v in enumerate(image_values)]
242 @override
243 def _current_image_prefix(self, column: str | None = None) -> str | None:
244 """Return the path prefix for the given column.
246 Reads from ``self.current_path_prefix``, a dict set by the job
247 mapping column names to path prefixes.
248 """
249 prefix_map: dict[str, str] = getattr(self, "current_path_prefix", {})
250 return prefix_map.get(column) # type: ignore[arg-type]
252 @override
253 def compute_features(self, batch: pa.RecordBatch, prev_features: pa.Array = None) -> dict[str, pa.Array]:
254 """
255 Extract image embeddings for all samples in the batch.
257 1. Images are loaded and transformed.
258 2. Model inference is performed in sub-batches defined by `infer.batch_size`.
259 3. Results are aggregated into a pyarrow `FixedSizeListArray`.
261 Args:
262 batch: Raw pyarrow batch.
263 prev_features: Pre-computed features (not used).
265 Returns:
266 Dictionary mapping column-prefixed embedding names to arrays.
267 """
268 self._ensure_model_loaded()
270 available = batch.schema.names
271 cols = resolve_include_exclude(
272 self.input_columns,
273 self.exclude_columns or None,
274 available,
275 )
276 if not cols: 276 ↛ 277line 276 didn't jump to line 277 because the condition on line 276 was never true
277 logger.warning(f"[{self.name}] no input columns matched in batch")
278 return {}
280 result: dict[str, pa.Array] = {}
281 for col in cols:
282 if col not in available:
283 logger.warning(f"[ImageEmbeddingProcessor] missing column '{col}'")
284 continue
286 image_values = batch.column(col).to_pylist()
287 image_tensors = self._load_image_tensors(image_values, column=col)
289 self.feature_extractor.eval()
290 with torch.no_grad():
291 if self.multi_layer:
292 raw = self._compute_features_multi_layer(image_tensors)
293 else:
294 raw = self._compute_features_single_layer(image_tensors)
296 for k, v in raw.items():
297 result[self._output_column_name(col, k)] = v
299 return result
301 @staticmethod
302 def _normalize_embedding(emb: np.ndarray | None, embed_dim: int) -> list[float]:
303 """Convert an embedding to a flat list of exactly embed_dim floats."""
304 if emb is None:
305 return [0.0] * embed_dim
306 flat_emb = emb.ravel()
307 if flat_emb.size != embed_dim:
308 if flat_emb.size > embed_dim:
309 flat_emb = flat_emb[:embed_dim]
310 else:
311 flat_emb = np.pad(flat_emb, (0, embed_dim - flat_emb.size))
312 return flat_emb.tolist()
314 def _build_fixed_array(self, embs: list[np.ndarray | None], embed_dim: int) -> pa.FixedSizeListArray:
315 """Build a FixedSizeListArray from a list of embedding vectors.
317 Args:
318 embs: List of embedding arrays or None.
319 embed_dim: Expected dimension of each embedding.
321 Returns:
322 A FixedSizeListArray of float32.
323 """
324 if embed_dim <= 0: 324 ↛ 325line 324 didn't jump to line 325 because the condition on line 324 was never true
325 raise ValueError(f"embed_dim must be positive, got {embed_dim}")
326 flat: list[float] = []
327 for emb in embs:
328 flat.extend(self._normalize_embedding(emb, embed_dim))
329 flat_array = pa.array(np.asarray(flat, dtype=np.float32))
330 return pa.FixedSizeListArray.from_arrays(flat_array, embed_dim)
332 def _compute_features_single_layer(self, image_tensors: list[torch.Tensor | None]) -> dict[str, pa.Array]:
333 """Compute embeddings for a single target layer.
335 Args:
336 image_tensors: List of preprocessed image tensors or None.
338 Returns:
339 Dictionary with 'embedding' key.
340 """
341 embs: list[np.ndarray | None] = []
342 with torch.no_grad():
343 for batch_start in range(0, len(image_tensors), self.batch_size):
344 batch_slice = image_tensors[batch_start : batch_start + self.batch_size]
345 self._process_batch_single(batch_slice, embs)
347 embed_dim = self._infer_embed_dim(embs)
348 if embed_dim is None or embed_dim <= 0:
349 return {}
350 return {"embedding": self._build_fixed_array(embs, embed_dim)}
352 def _process_batch_single(self, batch_slice: list[torch.Tensor | None], embs: list[np.ndarray | None]) -> None:
353 """Process a single batch for single-layer embedding extraction.
355 Args:
356 batch_slice: Subset of image tensors.
357 embs: Output list to append embeddings to.
358 """
359 valid = [t for t in batch_slice if t is not None]
360 if not valid:
361 embs.extend([None] * len(batch_slice))
362 return
364 batch_tensor = torch.stack(valid).to(self.device)
365 out = self.feature_extractor(batch_tensor)
366 if isinstance(out, dict): 366 ↛ 370line 366 didn't jump to line 370 because the condition on line 366 was always true
367 flat_feats = [layer_output.flatten(1) for layer_output in out.values()]
368 feats = torch.cat(flat_feats, dim=1)
369 else:
370 feats = out.flatten(1) if out.dim() > 2 else out
371 batch_embeddings_np = feats.detach().cpu().numpy().astype("float32")
373 pos = 0
374 for item_or_none in batch_slice:
375 if item_or_none is None: 375 ↛ 376line 375 didn't jump to line 376 because the condition on line 375 was never true
376 embs.append(None)
377 else:
378 embs.append(batch_embeddings_np[pos])
379 pos += 1
381 def _infer_embed_dim(self, embs: list[np.ndarray | None]) -> int | None:
382 """Infer embedding dimension from the first valid embedding.
384 Args:
385 embs: List of embeddings or None.
387 Returns:
388 Embedding dimension, or None if no valid embeddings exist.
389 """
390 if self._embed_dim is not None:
391 return self._embed_dim
392 for emb in embs:
393 if emb is not None:
394 self._embed_dim = int(emb.size)
395 return self._embed_dim
396 return None
398 def _compute_features_multi_layer(self, image_tensors: list[torch.Tensor | None]) -> dict[str, pa.Array]:
399 """Compute embeddings for multiple target layers.
401 Each layer's output is flattened and stored in a separate column
402 named ``emb_<layer_name>`` (with dots replaced by underscores).
404 Args:
405 image_tensors: List of preprocessed image tensors or None.
407 Returns:
408 Dictionary mapping layer column names to FixedSizeListArrays.
409 """
410 layer_cols = [f"emb_{layer.replace('.', '_')}" for layer in self.target_layers]
411 channel_cols = [f"{col}_channels" for col in layer_cols]
412 per_layer_embs: dict[str, list[np.ndarray | None]] = {col: [] for col in layer_cols}
413 per_layer_channels: dict[str, list[int | None]] = {col: [] for col in channel_cols}
415 with torch.no_grad():
416 for batch_start in range(0, len(image_tensors), self.batch_size):
417 batch_slice = image_tensors[batch_start : batch_start + self.batch_size]
418 self._process_batch_multi(batch_slice, layer_cols, channel_cols, per_layer_embs, per_layer_channels)
420 return self._build_multi_layer_results(layer_cols, channel_cols, per_layer_embs, per_layer_channels)
422 def _build_batch_np_dict(
423 self,
424 out_dict: dict[str, torch.Tensor],
425 valid_len: int,
426 ) -> dict[str, np.ndarray]:
427 """Build per-layer numpy arrays from a batch of forward pass outputs.
429 Args:
430 out_dict: Output dict from the feature extractor.
431 valid_len: Number of valid (non-None) samples in the batch.
433 Returns:
434 Dict mapping layer/column names to numpy arrays.
435 """
436 batch_np_dict: dict[str, np.ndarray] = {}
437 for layer_name in self.target_layers:
438 col = f"emb_{layer_name.replace('.', '_')}"
439 feats = out_dict[layer_name]
440 flat_feats = feats.flatten(1) if feats.dim() > 2 else feats
441 batch_np_dict[col] = flat_feats.detach().cpu().numpy().astype("float32")
442 batch_np_dict[f"{col}_channels"] = np.full(valid_len, feats.shape[1], dtype=np.int32)
443 return batch_np_dict
445 @staticmethod
446 def _append_none_row(
447 layer_cols: list[str],
448 channel_cols: list[str],
449 per_layer_embs: dict[str, list[np.ndarray | None]],
450 per_layer_channels: dict[str, list[int | None]],
451 ) -> None:
452 """Append None entries for all layer/channel columns."""
453 for col in layer_cols:
454 per_layer_embs[col].append(None)
455 for col in channel_cols:
456 per_layer_channels[col].append(None)
458 @staticmethod
459 def _append_valid_row(
460 pos: int,
461 layer_cols: list[str],
462 channel_cols: list[str],
463 batch_np_dict: dict[str, np.ndarray],
464 per_layer_embs: dict[str, list[np.ndarray | None]],
465 per_layer_channels: dict[str, list[int | None]],
466 ) -> None:
467 """Append embeddings for a valid (non-None) item at the given position."""
468 for col in layer_cols:
469 per_layer_embs[col].append(batch_np_dict[col][pos])
470 for col in channel_cols:
471 per_layer_channels[col].append(int(batch_np_dict[col][pos]))
473 @staticmethod
474 def _append_batch_results(
475 batch_slice: list[torch.Tensor | None],
476 layer_cols: list[str],
477 channel_cols: list[str],
478 batch_np_dict: dict[str, np.ndarray],
479 per_layer_embs: dict[str, list[np.ndarray | None]],
480 per_layer_channels: dict[str, list[int | None]],
481 ) -> None:
482 """Append per-layer results for a batch to the per-layer collections.
484 Args:
485 batch_slice: Subset of image tensors.
486 layer_cols: Layer column names.
487 channel_cols: Channel column names.
488 batch_np_dict: Numpy arrays per column.
489 per_layer_embs: Per-layer embedding lists to append to.
490 per_layer_channels: Per-layer channel lists to append to.
491 """
492 pos = 0
493 for item_or_none in batch_slice:
494 if item_or_none is None: 494 ↛ 495line 494 didn't jump to line 495 because the condition on line 494 was never true
495 ImageEmbeddingProcessor._append_none_row(layer_cols, channel_cols, per_layer_embs, per_layer_channels)
496 else:
497 ImageEmbeddingProcessor._append_valid_row(
498 pos, layer_cols, channel_cols, batch_np_dict, per_layer_embs, per_layer_channels
499 )
500 pos += 1
502 def _process_batch_multi(
503 self,
504 batch_slice: list[torch.Tensor | None],
505 layer_cols: list[str],
506 channel_cols: list[str],
507 per_layer_embs: dict[str, list[np.ndarray | None]],
508 per_layer_channels: dict[str, list[int | None]],
509 ) -> None:
510 """Process a single batch for multi-layer embedding extraction.
512 Args:
513 batch_slice: Subset of image tensors.
514 layer_cols: Layer column names.
515 channel_cols: Channel column names.
516 per_layer_embs: Per-layer embedding lists to append to.
517 per_layer_channels: Per-layer channel lists to append to.
518 """
519 valid = [t for t in batch_slice if t is not None]
520 if not valid: 520 ↛ 521line 520 didn't jump to line 521 because the condition on line 520 was never true
521 for col in layer_cols:
522 per_layer_embs[col].extend([None] * len(batch_slice))
523 for col in channel_cols:
524 per_layer_channels[col].extend([None] * len(batch_slice))
525 return
527 batch_tensor = torch.stack(valid).to(self.device)
528 out_dict = self.feature_extractor(batch_tensor)
529 batch_np_dict = self._build_batch_np_dict(out_dict, len(valid))
530 ImageEmbeddingProcessor._append_batch_results(
531 batch_slice, layer_cols, channel_cols, batch_np_dict, per_layer_embs, per_layer_channels
532 )
534 @staticmethod
535 def _find_embed_dim(embs: list[np.ndarray | None]) -> int | None:
536 """Find the embedding dimension from the first non-None embedding."""
537 for emb in embs: 537 ↛ 540line 537 didn't jump to line 540 because the loop on line 537 didn't complete
538 if emb is not None: 538 ↛ 537line 538 didn't jump to line 537 because the condition on line 538 was always true
539 return int(emb.size)
540 return None
542 def _build_multi_layer_results(
543 self,
544 layer_cols: list[str],
545 channel_cols: list[str],
546 per_layer_embs: dict[str, list[np.ndarray | None]],
547 per_layer_channels: dict[str, list[int | None]],
548 ) -> dict[str, pa.Array]:
549 """Build the result dictionary from per-layer collections.
551 Args:
552 layer_cols: Layer column names.
553 channel_cols: Channel column names.
554 per_layer_embs: Per-layer embedding lists.
555 per_layer_channels: Per-layer channel lists.
557 Returns:
558 Dictionary mapping column names to Arrow arrays.
559 """
560 result: dict[str, pa.Array] = {}
561 for col in layer_cols:
562 embs = per_layer_embs[col]
563 embed_dim = self._find_embed_dim(embs)
564 if embed_dim is None or embed_dim == 0: 564 ↛ 565line 564 didn't jump to line 565 because the condition on line 564 was never true
565 continue
566 result[col] = self._build_fixed_array(embs, embed_dim)
567 for col in channel_cols:
568 vals = [v if v is not None else 0 for v in per_layer_channels[col]]
569 if any(v is not None for v in per_layer_channels[col]): 569 ↛ 567line 569 didn't jump to line 567 because the condition on line 569 was always true
570 result[col] = pa.array(vals, type=pa.int32())
571 return result
573 # utils functions
574 @staticmethod
575 def _resolve_device(device: str) -> str:
576 """Resolve ``"auto"`` to CUDA if available, else CPU."""
577 if device == "auto":
578 return "cuda" if torch.cuda.is_available() else "cpu"
579 return device
581 def _load_model(self, arch: str, device: str) -> Any:
582 """Load a pre-trained torchvision model.
584 Args:
585 arch: Model architecture name (e.g., 'resnet18', 'resnet50').
586 device: Device to load the model on ('cpu' or 'cuda').
588 Returns:
589 The loaded PyTorch model.
590 """
591 try:
592 model = torchvision.models.get_model(arch, weights="DEFAULT")
593 except Exception:
594 # Fallback for older torchvision that lacks get_model()
595 model = getattr(torchvision.models, arch)(pretrained=True)
596 return model.to(device)
598 def _make_extractor(self, model: torch.nn.Module, target_layer: Any) -> Any:
599 """Create a feature extractor from a model.
601 Args:
602 model: The PyTorch model to extract features from.
603 target_layer: Layer name (str), index (int), or list of names to extract.
605 Returns:
606 A feature extractor that returns the requested layer outputs.
607 """
608 names = list(dict(model.named_modules()).keys())
609 if isinstance(target_layer, list):
610 nodes = {n: n for n in target_layer}
611 elif isinstance(target_layer, int): 611 ↛ 616line 611 didn't jump to line 616 because the condition on line 611 was always true
612 idx = target_layer if target_layer >= 0 else len(names) + target_layer
613 layer_name = names[idx]
614 nodes = {layer_name: "features"}
615 else:
616 nodes = {target_layer: "features"}
617 with warnings.catch_warnings():
618 warnings.simplefilter("ignore", UserWarning)
619 return create_feature_extractor(model, return_nodes=nodes)