flan/winnow
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Implement advanced diversity selection algorithms, comprehensive quality filtering, and caching for improved image curation.

Browse Source
This commit is contained in:
Sebastian G
2026-03-02 20:09:45 -05:00
parent 2c4f93bfb9
commit cc7293b66a
15 changed files with 1836 additions and 781 deletions
Download Patch File Download Diff File Expand all files Collapse all files
if_curator/diversity.py
+332 -32
View File
@@ -1,9 +1,12 @@
"""
Diversity selection for training data curation.
Uses Farthest Point Sampling (FPS) algorithm with embeddings:
- Faces: InsightFace embeddings
- Objects: SigLIP embeddings
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
@@ -15,6 +18,7 @@ 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__)
@@ -28,7 +32,7 @@ def select_diverse_assets(
progress_callback=None,
) -> list:
"""
Select diverse assets using Farthest Point Sampling or time spread.
Select diverse assets using cluster-aware FPS or time spread.
Args:
assets: List of asset dicts from Immich API
@@ -61,6 +65,11 @@ def select_diverse_assets(
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:
@@ -71,13 +80,106 @@ def _fetch_thumbnail(asset_id: str, timeout: int = 10) -> Image.Image | None:
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 Farthest Point Sampling."""
"""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)))
@@ -89,24 +191,68 @@ def _select_by_embedding(
else:
candidates = assets
# Compute embeddings
embeddings, valid_candidates = [], []
for i, asset in enumerate(candidates):
if progress_callback:
progress_callback(i, len(candidates))
# --- Phase 1: Concurrent thumbnail download ---
from concurrent.futures import ThreadPoolExecutor, as_completed
img = _fetch_thumbnail(asset["id"])
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
emb = get_embedding(img, entity_type)
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)
@@ -115,58 +261,212 @@ def _select_by_embedding(
logger.warning(f"Only {len(valid_candidates)} valid embeddings. Returning all.")
return valid_candidates
# Farthest Point Sampling with vectorized distance computation
return _farthest_point_sampling(
embeddings, valid_candidates, limit, auto_threshold=0.15
# --- Phase 5: Cluster-aware selection ---
return _cluster_aware_selection(
embeddings,
valid_candidates,
limit,
entity_type=entity_type,
confidence_scores=confidence_scores,
)
def _farthest_point_sampling(
# =============================================================================
# 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,
auto_threshold: float = 0.15,
entity_type: str = "face",
confidence_scores: list | None = None,
) -> list:
"""Vectorized Farthest Point Sampling."""
"""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 (cosine_dist = 1 - cosine_sim)
# Normalize for cosine distance
norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
emb_normed = emb_matrix / np.maximum(norms, 1e-8)
# Start from median-time sample
selected = [n // 2]
# 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)
target = 500 if limit == "auto" else limit
# 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:
# Update min distances with last selected point
last_emb = emb_normed[selected[-1]]
dists_to_last = 1 - emb_normed @ last_emb # Cosine distance
min_dists = np.minimum(min_dists, dists_to_last)
min_dists[selected[-1]] = -np.inf # Exclude already selected
# 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
# Find farthest point
best_idx = np.argmax(min_dists)
best_dist = min_dists[best_idx]
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 (threshold {auto_threshold})."
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.")
logger.info(f"Smart selection complete. Picked {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":