482 lines
16 KiB
Python
482 lines
16 KiB
Python
"""
|
||
Diversity selection for training data curation.
|
||
|
||
Selection pipeline:
|
||
1. Concurrent thumbnail download
|
||
2. Quality filtering (blur, IR, exposure, confidence, face size)
|
||
3. Face crop extraction (embed person's face, not full image)
|
||
4. Embedding computation (InsightFace or SigLIP)
|
||
5. Cluster-aware selection (K-Medoids + FPS with hard example weighting)
|
||
"""
|
||
|
||
import logging
|
||
from io import BytesIO
|
||
|
||
import numpy as np
|
||
import requests
|
||
from PIL import Image
|
||
|
||
from .config import Config, get_headers
|
||
from .embeddings import get_embedding, is_embedding_available
|
||
from .quality import assess_quality
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def select_diverse_assets(
|
||
assets: list,
|
||
limit: int | str,
|
||
entity_name: str,
|
||
selection_mode: str = "smart",
|
||
entity_type: str = "face",
|
||
progress_callback=None,
|
||
) -> list:
|
||
"""
|
||
Select diverse assets using cluster-aware FPS or time spread.
|
||
|
||
Args:
|
||
assets: List of asset dicts from Immich API
|
||
limit: Number to select, or "auto" for dynamic selection
|
||
entity_name: Name of the person/object for logging
|
||
selection_mode: 'smart' (embedding-based) or 'time' (time spread)
|
||
entity_type: 'face' or 'object' - determines embedding model
|
||
progress_callback: Optional callback(current, total) for progress
|
||
|
||
Returns:
|
||
List of selected assets
|
||
"""
|
||
# Fast path: fewer assets than limit
|
||
if limit != "auto" and len(assets) <= limit:
|
||
return assets
|
||
|
||
# Sort by creation time
|
||
assets = sorted(assets, key=lambda x: x.get("fileCreatedAt", ""))
|
||
|
||
if selection_mode != "smart" or not is_embedding_available(entity_type):
|
||
if selection_mode == "smart":
|
||
model_name = "InsightFace" if entity_type == "face" else "SigLIP"
|
||
logger.warning(f"{model_name} unavailable. Falling back to time spread.")
|
||
return _select_time_spread(assets, limit)
|
||
|
||
try:
|
||
return _select_by_embedding(assets, limit, entity_type, progress_callback)
|
||
except Exception as e:
|
||
logger.error(f"Smart Diversity failed: {e}. Falling back to time spread.")
|
||
return _select_time_spread(assets, limit)
|
||
|
||
|
||
# =============================================================================
|
||
# Thumbnail & Metadata Helpers
|
||
# =============================================================================
|
||
|
||
|
||
def _fetch_thumbnail(asset_id: str, timeout: int = 10) -> Image.Image | None:
|
||
"""Fetch thumbnail from Immich API."""
|
||
try:
|
||
url = f"{Config.IMMICH_URL}/api/assets/{asset_id}/thumbnail?size=preview&format=JPEG"
|
||
resp = requests.get(url, headers=get_headers(), timeout=timeout)
|
||
return Image.open(BytesIO(resp.content)) if resp.ok else None
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _get_face_bbox(asset: dict) -> tuple[float, float, float, float] | None:
|
||
"""Extract face bounding box from asset metadata if available."""
|
||
for person in asset.get("people", []):
|
||
faces = person.get("faces", [])
|
||
if faces:
|
||
f = faces[0]
|
||
return (
|
||
f.get("boundingBoxX1", 0),
|
||
f.get("boundingBoxY1", 0),
|
||
f.get("boundingBoxX2", 0),
|
||
f.get("boundingBoxY2", 0),
|
||
)
|
||
return None
|
||
|
||
|
||
def _get_face_confidence(asset: dict) -> float | None:
|
||
"""Extract face detection confidence from asset metadata if available."""
|
||
for person in asset.get("people", []):
|
||
faces = person.get("faces", [])
|
||
if faces:
|
||
return faces[0].get("score") or faces[0].get("confidence")
|
||
return None
|
||
|
||
|
||
def _crop_face_from_thumbnail(
|
||
img: Image.Image,
|
||
asset: dict,
|
||
margin: float = 0.25,
|
||
) -> Image.Image | None:
|
||
"""Crop the face region from a thumbnail using Immich bbox metadata.
|
||
|
||
By cropping before embedding, we guarantee InsightFace embeds the
|
||
correct person's face (not the largest face in a group photo).
|
||
|
||
Args:
|
||
img: Full preview thumbnail
|
||
asset: Asset dict with people/faces metadata
|
||
margin: Extra margin around the bbox (fraction, default 25%)
|
||
|
||
Returns:
|
||
Cropped face PIL image, or None if no face metadata available
|
||
"""
|
||
bbox = _get_face_bbox(asset)
|
||
if bbox is None:
|
||
return None
|
||
|
||
x1, y1, x2, y2 = bbox
|
||
img_w, img_h = img.size
|
||
|
||
# Get metadata dimensions to scale bbox
|
||
for person in asset.get("people", []):
|
||
faces = person.get("faces", [])
|
||
if faces:
|
||
meta_w = faces[0].get("imageWidth") or img_w
|
||
meta_h = faces[0].get("imageHeight") or img_h
|
||
scale_x, scale_y = img_w / meta_w, img_h / meta_h
|
||
x1, y1 = x1 * scale_x, y1 * scale_y
|
||
x2, y2 = x2 * scale_x, y2 * scale_y
|
||
break
|
||
|
||
face_w, face_h = x2 - x1, y2 - y1
|
||
|
||
# Add margin so InsightFace's internal alignment has context
|
||
mx, my = face_w * margin, face_h * margin
|
||
crop = img.crop(
|
||
(
|
||
max(0, x1 - mx),
|
||
max(0, y1 - my),
|
||
min(img_w, x2 + mx),
|
||
min(img_h, y2 + my),
|
||
)
|
||
)
|
||
|
||
# Skip if too small for meaningful embedding
|
||
if crop.width < 30 or crop.height < 30:
|
||
return None
|
||
|
||
return crop
|
||
|
||
|
||
# =============================================================================
|
||
# Embedding Collection
|
||
# =============================================================================
|
||
|
||
|
||
def _select_by_embedding(
|
||
assets: list,
|
||
limit: int | str,
|
||
entity_type: str,
|
||
progress_callback=None,
|
||
) -> list:
|
||
"""Select assets using embedding-based cluster-aware FPS.
|
||
|
||
Pipeline:
|
||
1. Concurrent thumbnail download
|
||
2. Quality filtering
|
||
3. Face crop extraction (face mode only)
|
||
4. Embedding computation
|
||
5. Cluster-aware selection with hard example weighting
|
||
"""
|
||
# Determine candidate pool (cap at 3000 for performance)
|
||
effective_limit = 30 if limit == "auto" else limit
|
||
pool_size = min(3000, max(effective_limit * 20, len(assets)))
|
||
|
||
# Subsample if needed (evenly distributed in time)
|
||
if len(assets) > pool_size:
|
||
indices = np.linspace(0, len(assets) - 1, pool_size, dtype=int)
|
||
candidates = [assets[i] for i in indices]
|
||
else:
|
||
candidates = assets
|
||
|
||
# --- Phase 1: Concurrent thumbnail download ---
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
|
||
thumbnail_map: dict[str, Image.Image] = {}
|
||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||
futures = {pool.submit(_fetch_thumbnail, a["id"]): a for a in candidates}
|
||
for i, future in enumerate(as_completed(futures)):
|
||
if progress_callback:
|
||
progress_callback(i, len(candidates))
|
||
asset = futures[future]
|
||
try:
|
||
img = future.result()
|
||
if img is not None:
|
||
thumbnail_map[asset["id"]] = img
|
||
except Exception:
|
||
continue
|
||
|
||
# --- Phase 2-4: Quality filter → Crop → Embed ---
|
||
embeddings, valid_candidates, confidence_scores = [], [], []
|
||
quality_filtered = 0
|
||
|
||
for asset in candidates:
|
||
img = thumbnail_map.get(asset["id"])
|
||
if img is None:
|
||
continue
|
||
|
||
confidence = _get_face_confidence(asset)
|
||
|
||
# Quality gate: filter before expensive embedding computation
|
||
if entity_type == "face":
|
||
face_bbox = _get_face_bbox(asset)
|
||
quality = assess_quality(
|
||
img,
|
||
face_bbox=face_bbox,
|
||
confidence=confidence,
|
||
blur_threshold=Config.BLUR_THRESHOLD,
|
||
min_face_px=Config.MIN_FACE_WIDTH,
|
||
min_confidence=Config.MIN_CONFIDENCE,
|
||
)
|
||
if not quality.passed:
|
||
quality_filtered += 1
|
||
logger.debug(f"Quality filtered {asset['id']}: {quality.reason}")
|
||
continue
|
||
|
||
# Crop the target person's face before embedding
|
||
face_crop = _crop_face_from_thumbnail(img, asset)
|
||
embed_img = face_crop if face_crop is not None else img
|
||
else:
|
||
embed_img = img
|
||
|
||
emb = get_embedding(embed_img, entity_type, asset_id=asset["id"])
|
||
if emb is not None:
|
||
embeddings.append(emb)
|
||
valid_candidates.append(asset)
|
||
confidence_scores.append(confidence)
|
||
|
||
if progress_callback:
|
||
progress_callback(len(candidates), len(candidates))
|
||
|
||
if quality_filtered > 0:
|
||
logger.info(f"Quality filtering removed {quality_filtered} images.")
|
||
|
||
if not embeddings:
|
||
logger.warning("No valid embeddings found. Falling back to time spread.")
|
||
return _select_time_spread(assets, limit)
|
||
|
||
if limit != "auto" and len(valid_candidates) < limit:
|
||
logger.warning(f"Only {len(valid_candidates)} valid embeddings. Returning all.")
|
||
return valid_candidates
|
||
|
||
# --- Phase 5: Cluster-aware selection ---
|
||
return _cluster_aware_selection(
|
||
embeddings,
|
||
valid_candidates,
|
||
limit,
|
||
entity_type=entity_type,
|
||
confidence_scores=confidence_scores,
|
||
)
|
||
|
||
|
||
# =============================================================================
|
||
# K-Medoids (Lightweight Implementation)
|
||
# =============================================================================
|
||
|
||
|
||
def _kmedoids(dist_matrix: np.ndarray, k: int, max_iter: int = 50) -> tuple[list[int], np.ndarray]:
|
||
"""Lightweight K-Medoids clustering using cosine distance matrix.
|
||
|
||
Args:
|
||
dist_matrix: (N, N) pairwise distance matrix
|
||
k: Number of clusters
|
||
max_iter: Maximum iterations for swap step
|
||
|
||
Returns:
|
||
(medoid_indices, cluster_labels) tuple
|
||
"""
|
||
n = dist_matrix.shape[0]
|
||
rng = np.random.default_rng(42)
|
||
|
||
# Initialize medoids: first = most central point, rest = farthest from chosen
|
||
total_dist = dist_matrix.sum(axis=1)
|
||
medoids = [int(np.argmin(total_dist))]
|
||
|
||
for _ in range(k - 1):
|
||
dists_to_chosen = dist_matrix[:, medoids].min(axis=1)
|
||
dists_to_chosen[medoids] = -np.inf
|
||
medoids.append(int(np.argmax(dists_to_chosen)))
|
||
|
||
# Iterative swap step
|
||
medoids = list(medoids)
|
||
labels = np.argmin(dist_matrix[:, medoids], axis=1)
|
||
cost = sum(dist_matrix[i, medoids[labels[i]]] for i in range(n))
|
||
|
||
for _ in range(max_iter):
|
||
improved = False
|
||
# Try swapping each medoid with a random non-medoid
|
||
non_medoids = [i for i in range(n) if i not in medoids]
|
||
if not non_medoids:
|
||
break
|
||
|
||
for m_idx in range(k):
|
||
candidates = rng.choice(non_medoids, size=min(10, len(non_medoids)), replace=False)
|
||
for cand in candidates:
|
||
new_medoids = medoids.copy()
|
||
new_medoids[m_idx] = cand
|
||
new_labels = np.argmin(dist_matrix[:, new_medoids], axis=1)
|
||
new_cost = sum(dist_matrix[i, new_medoids[new_labels[i]]] for i in range(n))
|
||
if new_cost < cost:
|
||
medoids = new_medoids
|
||
labels = new_labels
|
||
cost = new_cost
|
||
improved = True
|
||
break
|
||
if improved:
|
||
break
|
||
|
||
if not improved:
|
||
break
|
||
|
||
return medoids, labels
|
||
|
||
|
||
# =============================================================================
|
||
# Cluster-Aware Selection (K-Medoids + FPS Hybrid)
|
||
# =============================================================================
|
||
|
||
|
||
def _compute_adaptive_threshold(emb_normed: np.ndarray, entity_type: str) -> float:
|
||
"""Compute adaptive FPS stop threshold based on actual embedding distribution.
|
||
|
||
Instead of a hardcoded threshold, samples pairwise distances and sets
|
||
the threshold as a fraction of the median pairwise distance.
|
||
"""
|
||
n = len(emb_normed)
|
||
sample_size = min(200, n)
|
||
rng = np.random.default_rng(42)
|
||
indices = rng.choice(n, sample_size, replace=False) if n > sample_size else np.arange(n)
|
||
sample = emb_normed[indices]
|
||
|
||
# Compute pairwise cosine distances for the sample
|
||
pairwise = 1 - sample @ sample.T
|
||
upper_tri = pairwise[np.triu_indices(len(sample), k=1)]
|
||
median_dist = float(np.median(upper_tri))
|
||
|
||
# Faces: 20% of median (tighter — want fewer, more distinct images)
|
||
# Objects: 10% of median (wider — want more diversity)
|
||
fraction = 0.20 if entity_type == "face" else 0.10
|
||
threshold = max(0.05, median_dist * fraction)
|
||
|
||
logger.info(
|
||
f"Adaptive threshold: {threshold:.4f} "
|
||
f"(median_dist={median_dist:.4f}, fraction={fraction}, type={entity_type})"
|
||
)
|
||
return threshold
|
||
|
||
|
||
def _cluster_aware_selection(
|
||
embeddings: list,
|
||
candidates: list,
|
||
limit: int | str,
|
||
entity_type: str = "face",
|
||
confidence_scores: list | None = None,
|
||
) -> list:
|
||
"""Two-stage selection: K-Medoids clustering → FPS with hard example weighting.
|
||
|
||
Stage 1: Cluster embeddings into k groups, select medoids as initial picks.
|
||
This guarantees at least one representative from every distinct "look".
|
||
|
||
Stage 2: Fill remaining budget with FPS across cluster boundaries,
|
||
biasing toward hard examples (low-confidence candidates).
|
||
"""
|
||
emb_matrix = np.vstack(embeddings) # (N, D)
|
||
n = len(emb_matrix)
|
||
|
||
# Normalize for cosine distance
|
||
norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
|
||
emb_normed = emb_matrix / np.maximum(norms, 1e-8)
|
||
|
||
# Build confidence weight array for hard example boosting
|
||
conf_array = np.ones(n)
|
||
if confidence_scores and entity_type == "face":
|
||
for i, c in enumerate(confidence_scores):
|
||
if c is not None:
|
||
conf_array[i] = c
|
||
|
||
# Compute adaptive threshold for auto mode
|
||
auto_threshold = _compute_adaptive_threshold(emb_normed, entity_type) if limit == "auto" else 0.0
|
||
target = Config.MAX_AUTO_IMAGES if limit == "auto" else limit
|
||
|
||
# --- Stage 1: K-Medoids clustering ---
|
||
k = min(max(5, target // 4), n // 3, n) # e.g., 5-20 clusters
|
||
logger.info(f"Clustering {n} embeddings into {k} groups (K-Medoids)...")
|
||
|
||
# Compute full cosine distance matrix
|
||
dist_matrix = 1 - emb_normed @ emb_normed.T
|
||
|
||
medoid_indices, cluster_labels = _kmedoids(dist_matrix, k)
|
||
selected = list(medoid_indices)
|
||
selected_set = set(selected)
|
||
|
||
logger.info(f"Selected {len(selected)} cluster medoids as initial picks.")
|
||
|
||
# --- Stage 2: FPS with hard example weighting ---
|
||
min_dists = np.full(n, np.inf)
|
||
|
||
# Initialize min distances from all medoids
|
||
for idx in selected:
|
||
dists = dist_matrix[idx]
|
||
min_dists = np.minimum(min_dists, dists)
|
||
for idx in selected:
|
||
min_dists[idx] = -np.inf
|
||
|
||
while len(selected) < target:
|
||
# Hard example weighting: boost distance for low-confidence candidates
|
||
# Confidence < 0.85 gets up to 1.5× distance boost
|
||
hard_weight = np.where(conf_array < 0.85, 1.0 + (0.85 - conf_array) * 2.0, 1.0)
|
||
weighted_dists = min_dists * hard_weight
|
||
|
||
best_idx = int(np.argmax(weighted_dists))
|
||
best_dist = min_dists[best_idx] # Use unweighted for threshold comparison
|
||
|
||
if best_dist == -np.inf:
|
||
break # All points selected
|
||
|
||
if limit == "auto" and best_dist < auto_threshold:
|
||
logger.info(
|
||
f"Auto-stop: Next best image {best_dist:.3f} away " f"(adaptive threshold {auto_threshold:.4f})."
|
||
)
|
||
break
|
||
|
||
selected.append(best_idx)
|
||
selected_set.add(best_idx)
|
||
|
||
# Update min distances
|
||
dists_to_new = dist_matrix[best_idx]
|
||
min_dists = np.minimum(min_dists, dists_to_new)
|
||
min_dists[best_idx] = -np.inf
|
||
|
||
# Log hard example stats
|
||
if entity_type == "face":
|
||
selected_conf = [conf_array[i] for i in selected if conf_array[i] < 1.0]
|
||
hard_count = sum(1 for c in selected_conf if c < 0.85)
|
||
logger.info(
|
||
f"Selection complete: {len(selected)} images " f"({hard_count} hard examples with confidence < 0.85)."
|
||
)
|
||
else:
|
||
logger.info(f"Selection complete: {len(selected)} diverse images.")
|
||
|
||
return [candidates[i] for i in selected]
|
||
|
||
|
||
# =============================================================================
|
||
# Time Spread Fallback
|
||
# =============================================================================
|
||
|
||
|
||
def _select_time_spread(assets: list, limit: int | str) -> list:
|
||
"""Select N assets evenly distributed in time."""
|
||
if limit == "auto":
|
||
limit = 30
|
||
|
||
logger.info(f"Selecting {limit} images using time spread.")
|
||
|
||
if len(assets) <= limit:
|
||
return assets
|
||
|
||
indices = np.linspace(0, len(assets) - 1, limit, dtype=int)
|
||
return [assets[i] for i in np.unique(indices)]
|