Blablablab/audio-classification
0
1"""2Embedding Visualization Module3 4Provides 2D visualization of text/image embeddings for the admin dashboard,5enabling interactive exploration and prioritization of annotation items.6 7Key Components:8- EmbeddingVisualizationManager: Main class for embedding visualization9- UMAP dimensionality reduction for 2D projection10- Label coloring via MACE or majority vote11- Interactive selection and queue reordering12 13The visualization allows admins to:14- See clustering patterns in the data15- Identify annotated vs unannotated items16- Select regions to prioritize for annotation17- Interleave multiple selections for diverse sampling18"""19 20import logging21import threading22from dataclasses import dataclass, field23from typing import Any, Dict, List, Optional, Set, Tuple24import hashlib25import json26 27logger = logging.getLogger(__name__)28 29# Guarded imports for optional dependencies30try:31 import numpy as np32 _NUMPY_AVAILABLE = True33except ImportError:34 _NUMPY_AVAILABLE = False35 np = None36 37try:38 import umap39 _UMAP_AVAILABLE = True40except ImportError:41 _UMAP_AVAILABLE = False42 umap = None43 44# Singleton45_EMBEDDING_VIZ_MANAGER: Optional['EmbeddingVisualizationManager'] = None46_EMBEDDING_VIZ_LOCK = threading.Lock()47 48 49@dataclass50class EmbeddingVizConfig:51 """Configuration for embedding visualization."""52 enabled: bool = True53 sample_size: int = 100054 include_all_annotated: bool = True55 embedding_model: str = "all-MiniLM-L6-v2"56 image_embedding_model: str = "clip-ViT-B-32"57 umap_n_neighbors: int = 1558 umap_min_dist: float = 0.159 umap_metric: str = "cosine"60 label_source: str = "mace" # "mace" or "majority"61 62 63@dataclass64class VisualizationPoint:65 """A single point in the visualization."""66 instance_id: str67 x: float68 y: float69 label: Optional[str] = None70 label_source: Optional[str] = None71 preview: str = ""72 preview_type: str = "text" # "text" or "image"73 annotated: bool = False74 annotation_count: int = 075 76 77@dataclass78class VisualizationData:79 """Complete visualization data for the scatter plot."""80 points: List[VisualizationPoint] = field(default_factory=list)81 labels: List[Optional[str]] = field(default_factory=list)82 label_colors: Dict[Optional[str], str] = field(default_factory=dict)83 stats: Dict[str, Any] = field(default_factory=dict)84 85 86# Default color palette for labels87DEFAULT_COLORS = [88 "#22c55e", # green89 "#ef4444", # red90 "#3b82f6", # blue91 "#eab308", # yellow92 "#8b5cf6", # purple93 "#f97316", # orange94 "#06b6d4", # cyan95 "#ec4899", # pink96 "#14b8a6", # teal97 "#f59e0b", # amber98]99 100UNANNOTATED_COLOR = "#94a3b8" # slate gray101 102 103class EmbeddingVisualizationManager:104 """105 Manages embedding visualization for the admin dashboard.106 107 This class provides:108 - 2D UMAP projections of text/image embeddings109 - Label coloring via MACE or majority vote110 - Interactive selection and queue reordering111 - Caching with invalidation on new annotations112 """113 114 def __init__(self, config: EmbeddingVizConfig, app_config: Dict[str, Any]):115 """116 Initialize the embedding visualization manager.117 118 Args:119 config: EmbeddingVizConfig instance120 app_config: Full application configuration dictionary121 """122 self.config = config123 self.app_config = app_config124 self.logger = logging.getLogger(__name__)125 self._lock = threading.RLock()126 127 # State128 self.enabled = False129 self._projection_cache: Optional[Dict[str, Tuple[float, float]]] = None130 self._cache_hash: Optional[str] = None131 self._label_cache: Dict[str, Optional[str]] = {}132 133 # Check dependencies134 if not _NUMPY_AVAILABLE:135 self.logger.warning(136 "numpy not available. Embedding visualization disabled."137 )138 return139 140 if not _UMAP_AVAILABLE:141 self.logger.warning(142 "umap-learn not installed. Embedding visualization disabled. "143 "Install with: pip install umap-learn"144 )145 return146 147 if not config.enabled:148 self.logger.info("Embedding visualization disabled in config")149 return150 151 self.enabled = True152 self.logger.info("Embedding visualization manager initialized")153 154 def _get_diversity_manager(self):155 """Get the DiversityManager singleton."""156 from potato.diversity_manager import get_diversity_manager157 return get_diversity_manager()158 159 def _get_item_state_manager(self):160 """Get the ItemStateManager singleton."""161 from potato.item_state_management import get_item_state_manager162 return get_item_state_manager()163 164 def _get_user_state_manager(self):165 """Get the UserStateManager singleton."""166 from potato.user_state_management import get_user_state_manager167 return get_user_state_manager()168 169 def _compute_embedding_hash(self, embeddings: Dict[str, Any]) -> str:170 """Compute a hash of embedding IDs for cache invalidation."""171 sorted_ids = sorted(embeddings.keys())172 return hashlib.md5(",".join(sorted_ids).encode()).hexdigest()173 174 def compute_umap_projection(175 self,176 embeddings: Dict[str, Any],177 force: bool = False178 ) -> Dict[str, Tuple[float, float]]:179 """180 Compute UMAP 2D projection of embeddings.181 182 Args:183 embeddings: Dict mapping instance_id to embedding vector184 force: Force recomputation even if cached185 186 Returns:187 Dict mapping instance_id to (x, y) coordinates188 """189 if not self.enabled or not embeddings:190 return {}191 192 with self._lock:193 # Check cache194 current_hash = self._compute_embedding_hash(embeddings)195 if not force and self._projection_cache and self._cache_hash == current_hash:196 self.logger.debug("Using cached UMAP projection")197 return self._projection_cache198 199 try:200 self.logger.info(f"Computing UMAP projection for {len(embeddings)} embeddings")201 202 # Convert to numpy array203 instance_ids = list(embeddings.keys())204 vectors = np.array([embeddings[iid] for iid in instance_ids])205 206 # Ensure we have enough samples for UMAP207 n_samples = len(vectors)208 n_neighbors = min(self.config.umap_n_neighbors, n_samples - 1)209 if n_neighbors < 2:210 self.logger.warning(f"Not enough samples ({n_samples}) for UMAP")211 return {}212 213 # Run UMAP214 reducer = umap.UMAP(215 n_neighbors=n_neighbors,216 min_dist=self.config.umap_min_dist,217 metric=self.config.umap_metric,218 n_components=2,219 random_state=42220 )221 projection = reducer.fit_transform(vectors)222 223 # Build result dict224 result = {}225 for i, instance_id in enumerate(instance_ids):226 result[instance_id] = (float(projection[i, 0]), float(projection[i, 1]))227 228 # Cache result229 self._projection_cache = result230 self._cache_hash = current_hash231 232 self.logger.info(f"UMAP projection complete: {len(result)} points")233 return result234 235 except Exception as e:236 self.logger.error(f"UMAP projection failed: {e}")237 return {}238 239 def get_labels_for_instances(240 self,241 instance_ids: List[str],242 source: str = "mace"243 ) -> Dict[str, Optional[str]]:244 """245 Get predicted labels for instances.246 247 Args:248 instance_ids: List of instance IDs249 source: Label source - "mace" or "majority"250 251 Returns:252 Dict mapping instance_id to label (or None if unannotated)253 """254 result = {}255 256 if source == "mace":257 result = self._get_mace_labels(instance_ids)258 else:259 result = self._get_majority_labels(instance_ids)260 261 return result262 263 def _get_mace_labels(self, instance_ids: List[str]) -> Dict[str, Optional[str]]:264 """Get MACE predicted labels for instances."""265 result = {iid: None for iid in instance_ids}266 267 try:268 from potato.mace_manager import get_mace_manager269 270 mace_mgr = get_mace_manager()271 if not mace_mgr or not mace_mgr.mace_config.enabled:272 self.logger.debug("MACE not available, falling back to majority")273 return self._get_majority_labels(instance_ids)274 275 # Get predictions from all schemas276 summary = mace_mgr.get_results_summary()277 if "error" in summary or not summary.get("enabled"):278 return self._get_majority_labels(instance_ids)279 280 # Use first schema's predictions (most common case)281 schemas = summary.get("schemas", {})282 if not schemas:283 return self._get_majority_labels(instance_ids)284 285 # Get first schema with predictions286 for schema_name, schema_data in schemas.items():287 predictions = schema_data.get("predictions", {})288 label_names = schema_data.get("label_names", [])289 290 for instance_id in instance_ids:291 if instance_id in predictions:292 pred_idx = predictions[instance_id]293 if isinstance(pred_idx, int) and pred_idx < len(label_names):294 result[instance_id] = label_names[pred_idx]295 break # Use first schema only296 297 except ImportError:298 self.logger.debug("MACE manager not available")299 except Exception as e:300 self.logger.error(f"Error getting MACE labels: {e}")301 302 return result303 304 def _get_majority_labels(self, instance_ids: List[str]) -> Dict[str, Optional[str]]:305 """Get majority vote labels for instances."""306 from collections import Counter307 308 result = {iid: None for iid in instance_ids}309 310 try:311 usm = self._get_user_state_manager()312 if not usm:313 return result314 315 # Get annotation schemes316 annotation_schemes = self.app_config.get("annotation_schemes", [])317 if not annotation_schemes:318 return result319 320 # Use first categorical schema321 target_schema = None322 for scheme in annotation_schemes:323 if scheme.get("annotation_type") in ["radio", "select", "multiselect"]:324 target_schema = scheme.get("name")325 break326 327 if not target_schema:328 return result329 330 # Count labels per instance331 from potato.flask_server import get_users332 users = get_users()333 334 for instance_id in instance_ids:335 labels = []336 for username in users:337 user_state = usm.get_user_state(username)338 if not user_state:339 continue340 341 annotations = user_state.get_all_annotations()342 if instance_id not in annotations:343 continue344 345 instance_annot = annotations[instance_id]346 label_annotations = instance_annot.get("labels", {})347 348 for label, value in label_annotations.items():349 label_schema = None350 label_name = None351 352 if hasattr(label, 'schema'):353 label_schema = label.schema354 label_name = getattr(label, 'name', None)355 elif hasattr(label, 'get_schema'):356 label_schema = label.get_schema()357 label_name = label.get_name() if hasattr(label, 'get_name') else None358 359 if label_schema == target_schema and label_name:360 labels.append(label_name)361 362 if labels:363 counter = Counter(labels)364 result[instance_id] = counter.most_common(1)[0][0]365 366 except Exception as e:367 self.logger.error(f"Error getting majority labels: {e}")368 369 return result370 371 def _assign_label_colors(self, unique_labels: List[Optional[str]]) -> Dict[Optional[str], str]:372 """Assign consistent colors to labels."""373 colors = {}374 color_idx = 0375 376 for label in unique_labels:377 if label is None:378 colors[None] = UNANNOTATED_COLOR379 else:380 colors[label] = DEFAULT_COLORS[color_idx % len(DEFAULT_COLORS)]381 color_idx += 1382 383 return colors384 385 def get_visualization_data(self, force_refresh: bool = False) -> VisualizationData:386 """387 Get complete visualization data for the scatter plot.388 389 Args:390 force_refresh: Force recomputation of projections391 392 Returns:393 VisualizationData with points, labels, and colors394 """395 if not self.enabled:396 return VisualizationData(397 stats={"error": "Embedding visualization not enabled"}398 )399 400 with self._lock:401 dm = self._get_diversity_manager()402 ism = self._get_item_state_manager()403 404 if not dm or not dm.enabled:405 return VisualizationData(406 stats={"error": "Diversity manager not available. Enable diversity_ordering in config."}407 )408 409 if not dm.embeddings:410 return VisualizationData(411 stats={"error": "No embeddings available. Ensure items have been loaded."}412 )413 414 # Get embeddings (possibly sampled)415 all_embedding_ids = set(dm.embeddings.keys())416 annotated_ids = set()417 418 # Find annotated instances419 if ism:420 for instance_id in all_embedding_ids:421 annotators = ism.get_annotators_for_item(instance_id)422 if annotators:423 annotated_ids.add(instance_id)424 425 # Sample if needed426 sample_ids = self._sample_instances(427 all_embedding_ids,428 annotated_ids,429 self.config.sample_size,430 self.config.include_all_annotated431 )432 433 # Get embeddings for sampled instances434 sampled_embeddings = {435 iid: dm.embeddings[iid]436 for iid in sample_ids437 if iid in dm.embeddings438 }439 440 # Compute UMAP projection441 projection = self.compute_umap_projection(sampled_embeddings, force=force_refresh)442 if not projection:443 return VisualizationData(444 stats={"error": "UMAP projection failed"}445 )446 447 # Get labels448 labels = self.get_labels_for_instances(449 list(projection.keys()),450 source=self.config.label_source451 )452 453 # Build points454 points = []455 unique_labels = set()456 457 for instance_id, (x, y) in projection.items():458 label = labels.get(instance_id)459 unique_labels.add(label)460 461 # Get preview text462 preview = ""463 preview_type = "text"464 if ism:465 item = ism.get_instance_by_id(instance_id)466 if item:467 text = item.get_text()468 if text:469 preview = text[:200] + "..." if len(text) > 200 else text470 # Check for image471 if hasattr(item, 'get_image_path'):472 img_path = item.get_image_path()473 if img_path:474 preview = img_path475 preview_type = "image"476 477 annotation_count = 0478 if ism:479 annotators = ism.get_annotators_for_item(instance_id)480 annotation_count = len(annotators) if annotators else 0481 482 points.append(VisualizationPoint(483 instance_id=instance_id,484 x=x,485 y=y,486 label=label,487 label_source=self.config.label_source if label else None,488 preview=preview,489 preview_type=preview_type,490 annotated=instance_id in annotated_ids,491 annotation_count=annotation_count492 ))493 494 # Assign colors495 label_colors = self._assign_label_colors(list(unique_labels))496 497 # Build stats498 stats = {499 "total_instances": len(all_embedding_ids),500 "visualized_instances": len(points),501 "annotated_instances": len(annotated_ids),502 "unannotated_instances": len(all_embedding_ids) - len(annotated_ids),503 "label_source": self.config.label_source,504 "unique_labels": len([l for l in unique_labels if l is not None])505 }506 507 return VisualizationData(508 points=points,509 labels=sorted([l for l in unique_labels if l is not None]) + [None],510 label_colors=label_colors,511 stats=stats512 )513 514 def _sample_instances(515 self,516 all_ids: Set[str],517 annotated_ids: Set[str],518 sample_size: int,519 include_all_annotated: bool520 ) -> Set[str]:521 """522 Sample instances for visualization.523 524 Args:525 all_ids: All available instance IDs526 annotated_ids: IDs that have been annotated527 sample_size: Maximum number of instances to include528 include_all_annotated: Always include all annotated instances529 530 Returns:531 Set of instance IDs to visualize532 """533 if len(all_ids) <= sample_size:534 return all_ids535 536 result = set()537 538 if include_all_annotated:539 result.update(annotated_ids)540 541 # Sample remaining from unannotated542 remaining_needed = sample_size - len(result)543 if remaining_needed > 0:544 unannotated = all_ids - annotated_ids545 if len(unannotated) <= remaining_needed:546 result.update(unannotated)547 else:548 # Random sample549 import random550 sampled = random.sample(list(unannotated), remaining_needed)551 result.update(sampled)552 553 return result554 555 def reorder_instances(556 self,557 selections: List[Dict[str, Any]],558 interleave: bool = True559 ) -> Dict[str, Any]:560 """561 Reorder the annotation queue based on selections.562 563 Args:564 selections: List of selection groups, each with:565 - instance_ids: List of selected instance IDs566 - priority: Priority number (lower = higher priority)567 interleave: Whether to interleave selections (default True)568 569 Returns:570 Dict with success status and reordering info571 """572 if not selections:573 return {"success": False, "error": "No selections provided"}574 575 ism = self._get_item_state_manager()576 if not ism:577 return {"success": False, "error": "ItemStateManager not available"}578 579 try:580 # Build new order581 if interleave:582 new_order = self._interleave_selections(selections)583 else:584 # Concatenate by priority585 sorted_selections = sorted(selections, key=lambda s: s.get("priority", 999))586 new_order = []587 for sel in sorted_selections:588 new_order.extend(sel.get("instance_ids", []))589 590 # Deduplicate while preserving order591 seen = set()592 deduped_order = []593 for iid in new_order:594 if iid not in seen:595 seen.add(iid)596 deduped_order.append(iid)597 598 # Apply reordering599 ism.reorder_instances(deduped_order)600 601 # Build preview of new order (first 10)602 preview = deduped_order[:10]603 604 return {605 "success": True,606 "reordered_count": len(deduped_order),607 "new_order_preview": preview608 }609 610 except Exception as e:611 self.logger.error(f"Error reordering instances: {e}")612 return {"success": False, "error": str(e)}613 614 def _interleave_selections(self, selections: List[Dict[str, Any]]) -> List[str]:615 """616 Interleave instances from multiple selections by priority.617 618 Example:619 selections = [620 {"instance_ids": ["a", "b", "c"], "priority": 1},621 {"instance_ids": ["x", "y"], "priority": 2}622 ]623 Result: ["a", "x", "b", "y", "c"]624 625 Lower priority number = higher priority (comes first in each round)626 627 Args:628 selections: List of selection dicts with instance_ids and priority629 630 Returns:631 List of interleaved instance IDs632 """633 # Sort by priority634 sorted_selections = sorted(selections, key=lambda s: s.get("priority", 999))635 636 # Create iterators637 iterators = [iter(s.get("instance_ids", [])) for s in sorted_selections]638 639 result = []640 while iterators:641 exhausted = []642 for i, it in enumerate(iterators):643 try:644 result.append(next(it))645 except StopIteration:646 exhausted.append(i)647 648 # Remove exhausted iterators (in reverse to maintain indices)649 for i in reversed(exhausted):650 iterators.pop(i)651 652 return result653 654 def invalidate_cache(self) -> None:655 """Invalidate the projection cache."""656 with self._lock:657 self._projection_cache = None658 self._cache_hash = None659 self._label_cache = {}660 self.logger.info("Embedding visualization cache invalidated")661 662 def get_stats(self) -> Dict[str, Any]:663 """Get visualization manager statistics."""664 dm = self._get_diversity_manager()665 666 return {667 "enabled": self.enabled,668 "umap_available": _UMAP_AVAILABLE,669 "numpy_available": _NUMPY_AVAILABLE,670 "embeddings_available": dm.enabled if dm else False,671 "embedding_count": len(dm.embeddings) if dm and dm.embeddings else 0,672 "cache_valid": self._projection_cache is not None,673 "config": {674 "sample_size": self.config.sample_size,675 "include_all_annotated": self.config.include_all_annotated,676 "label_source": self.config.label_source,677 "umap_n_neighbors": self.config.umap_n_neighbors,678 "umap_min_dist": self.config.umap_min_dist,679 }680 }681 682 def to_json(self) -> Dict[str, Any]:683 """Convert visualization data to JSON-serializable format."""684 data = self.get_visualization_data()685 686 points_json = []687 for p in data.points:688 points_json.append({689 "instance_id": p.instance_id,690 "x": p.x,691 "y": p.y,692 "label": p.label,693 "label_source": p.label_source,694 "preview": p.preview,695 "preview_type": p.preview_type,696 "annotated": p.annotated,697 "annotation_count": p.annotation_count698 })699 700 return {701 "points": points_json,702 "labels": data.labels,703 "label_colors": data.label_colors,704 "stats": data.stats705 }706 707 708def parse_embedding_viz_config(config_data: Dict[str, Any]) -> EmbeddingVizConfig:709 """710 Parse embedding_visualization section from config.711 712 Args:713 config_data: Full application configuration714 715 Returns:716 EmbeddingVizConfig instance717 """718 ev = config_data.get("embedding_visualization", {})719 720 return EmbeddingVizConfig(721 enabled=ev.get("enabled", True),722 sample_size=ev.get("sample_size", 1000),723 include_all_annotated=ev.get("include_all_annotated", True),724 embedding_model=ev.get("embedding_model", "all-MiniLM-L6-v2"),725 image_embedding_model=ev.get("image_embedding_model", "clip-ViT-B-32"),726 umap_n_neighbors=ev.get("umap", {}).get("n_neighbors", 15),727 umap_min_dist=ev.get("umap", {}).get("min_dist", 0.1),728 umap_metric=ev.get("umap", {}).get("metric", "cosine"),729 label_source=ev.get("label_source", "mace"),730 )731 732 733def init_embedding_viz_manager(734 config_data: Dict[str, Any]735) -> Optional[EmbeddingVisualizationManager]:736 """737 Initialize the singleton EmbeddingVisualizationManager.738 739 Args:740 config_data: Full application configuration741 742 Returns:743 EmbeddingVisualizationManager instance, or None if disabled744 """745 global _EMBEDDING_VIZ_MANAGER746 747 with _EMBEDDING_VIZ_LOCK:748 if _EMBEDDING_VIZ_MANAGER is None:749 viz_config = parse_embedding_viz_config(config_data)750 _EMBEDDING_VIZ_MANAGER = EmbeddingVisualizationManager(viz_config, config_data)751 752 return _EMBEDDING_VIZ_MANAGER753 754 755def get_embedding_viz_manager() -> Optional[EmbeddingVisualizationManager]:756 """Get the singleton EmbeddingVisualizationManager instance."""757 return _EMBEDDING_VIZ_MANAGER758 759 760def clear_embedding_viz_manager() -> None:761 """Clear the singleton (for testing)."""762 global _EMBEDDING_VIZ_MANAGER763 with _EMBEDDING_VIZ_LOCK:764 _EMBEDDING_VIZ_MANAGER = None765