| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
| from typing import Dict, List, Optional |
| import numpy as np |
|
|
|
|
| @dataclass |
| class TaskConfig: |
| task_id: str |
| name: str |
| classes: List[str] |
| pool_size: int |
| budget: int |
| success_threshold: float |
| rare_classes: List[str] |
| modality: str = "xray" |
| body_part: str = "chest" |
| variable_cost: bool = False |
| base_annotation_cost: int = 1 |
| rare_annotation_cost: int = 2 |
|
|
|
|
| @dataclass |
| class TaskState: |
| config: TaskConfig |
| pool_embeddings: np.ndarray |
| pool_labels: np.ndarray |
| pool_image_ids: List[str] |
| val_embeddings: np.ndarray |
| val_labels: np.ndarray |
| labeled_indices: List[int] = field(default_factory=list) |
| labeled_labels: List[str] = field(default_factory=list) |
| unlabeled_indices: List[int] = field(default_factory=list) |
| budget_used: int = 0 |
| annotation_cost_spent: int = 0 |
| current_auc: float = 0.5 |
| step_count: int = 0 |
| cumulative_reward: float = 0.0 |
| auc_trajectory: List[float] = field(default_factory=list) |
| classes_discovered: List[str] = field(default_factory=list) |
|
|
| def get_labeled_embeddings(self) -> np.ndarray: |
| if not self.labeled_indices: |
| return np.zeros((0, self.pool_embeddings.shape[1]), dtype=np.float32) |
| return self.pool_embeddings[self.labeled_indices] |
|
|
| def get_unlabeled_embeddings(self) -> np.ndarray: |
| if not self.unlabeled_indices: |
| return np.zeros((0, self.pool_embeddings.shape[1]), dtype=np.float32) |
| return self.pool_embeddings[self.unlabeled_indices] |
|
|
| def get_unlabeled_image_ids(self) -> List[str]: |
| return [self.pool_image_ids[i] for i in self.unlabeled_indices] |
|
|
| def find_index_by_image_id(self, image_id: str) -> Optional[int]: |
| try: |
| pool_pos = self.pool_image_ids.index(image_id) |
| if pool_pos in self.unlabeled_indices: |
| return pool_pos |
| return None |
| except ValueError: |
| return None |
|
|
| def get_budget_phase(self) -> str: |
| if self.config.budget == 0: |
| return "late" |
| progress = self.budget_used / self.config.budget |
| if progress < 0.3: |
| return "early" |
| if progress < 0.7: |
| return "mid" |
| return "late" |
|
|
| def get_class_distribution(self) -> Dict[str, int]: |
| distribution: Dict[str, int] = {cls: 0 for cls in self.config.classes} |
| for lbl in self.labeled_labels: |
| if lbl in distribution: |
| distribution[lbl] += 1 |
| return distribution |
|
|
| def get_class_coverage_ratio(self) -> float: |
| discovered = sum(1 for cls in self.config.classes if cls in self.classes_discovered) |
| return discovered / len(self.config.classes) if self.config.classes else 0.0 |
|
|
| def compute_annotation_cost(self, label: str) -> int: |
| if not self.config.variable_cost: |
| return self.config.base_annotation_cost |
| if label in self.config.rare_classes: |
| return self.config.rare_annotation_cost |
| return self.config.base_annotation_cost |
|
|
