Coverage for collectra/types/base.py: 80.31%
325 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1import uuid
2from dataclasses import dataclass, field
3from pathlib import Path
4from typing import Any
6import numpy as np
7import yaml
8from rich import print
10from collectra.cli import command
12from ..commons import BaseEntity, Node, NodeStatus
13from ..utils import (
14 change_dir,
15 error_msg,
16 load_class_from_string,
17 traceback_error,
18)
20__all__ = ["Artefact", "ArtefactNode"]
22from ..logger import get_logger
24logger = get_logger(__name__)
27class ArtefactCommands:
28 """CLI capabilities for a named pipeline artefact node.
30 Concrete artefact classes can return a subclass from ``cli_commands`` to
31 expose commands applying to every matching item in that node.
32 """
34 def __init__(self, node: "ArtefactNode", ext: str) -> None:
35 self.node = node
36 self.ext = ext
38 @command
39 def evaluate(
40 self,
41 predicted_folder: Path,
42 gold_folder: Path,
43 threshold: float = 0.5,
44 match_by_order: bool = False,
45 output: Path | None = None,
46 ) -> None:
47 """Evaluate this artefact against ground-truth Collectra results."""
48 from collectra.evaluator.base import Evaluator
50 report = Evaluator(
51 predicted_folder,
52 gold_folder,
53 ext=self.ext,
54 labels={self.node.name},
55 match_by_order=match_by_order,
56 ).evaluate(threshold=threshold)
57 for table in [*report.aggregate_tables, *report.tables]:
58 print(table)
59 if output is not None:
60 report.save_csv(output)
63@dataclass
64class Artefact(BaseEntity):
65 name: str
66 id: str = field(default="")
67 parents: list[str] = field(default_factory=list)
68 partition: str = field(default="")
69 embeddings: list[str] = field(default_factory=list)
70 orientation: str = field(default="north") # e.g., "landscape" or "portrait"
71 ensemble: list[str] = field(default_factory=list)
72 source_file: Path | None = field(
73 default=None, init=False, repr=False, compare=False
74 )
76 def set_parents(self, parents: list["Artefact"]) -> None:
77 self.parents = [parent.id for parent in parents]
79 def serialize(self) -> dict:
80 serialized = super().serialize()
81 if not self.partition:
82 serialized.pop("partition", None)
83 if "parents" in serialized:
84 if len(serialized["parents"]) == 0:
85 serialized.pop("parents")
86 elif len(serialized["parents"]) == 1:
87 serialized["parents"] = serialized["parents"][0]
88 return serialized
90 def evaluate(self, gold) -> float:
91 raise NotImplementedError(
92 "Eval method not implemented for base Artefact class."
93 )
95 def display(self, context) -> dict:
96 """Describe this artefact using the display protocol (see collectra.display).
98 Override to return an image, text, or properties view, or supply a
99 custom JavaScript module using ``context.renderer(module_path)``.
100 The context resolves related artefacts and publishes previews without
101 altering data. See docs/artefacts/display.rst for the renderer contract.
102 """
103 return {"kind": "properties", "data": self.serialize()}
105 def _extract(self, path: Path) -> None:
106 raise NotImplementedError(
107 "Extract method not implemented for base Artefact class."
108 )
110 def extract(self, path: Path | str) -> None:
111 path = Path(path)
112 path.parent.mkdir(parents=True, exist_ok=True)
113 self._extract(path)
115 @classmethod
116 def cli_commands(cls, node: "ArtefactNode", ext: str) -> ArtefactCommands:
117 """Return the CLI object representing this artefact type in ``node``."""
118 return ArtefactCommands(node, ext)
120 def _generate_id(self) -> str:
121 # Generate a unique ID based on the name and other attributes
122 return f"{self.name}-{uuid.uuid4()}"
124 def __post_init__(self) -> None:
125 super().__init__(self.name) # Initialize BaseEntity
126 if not self.id:
127 self.id = self._generate_id()
129 def attributes_to_ignore(self) -> set:
130 attributes = super().attributes_to_ignore()
131 attributes.add("name")
132 attributes.add("partition")
133 attributes.add("source_file")
134 return attributes
136 @classmethod
137 def all_attributes(cls):
138 return set(
139 [
140 attribute
141 for attribute in dir(cls)
142 if not attribute.startswith("_") and not callable(attribute)
143 ]
144 )
147@dataclass
148class ArtefactNode(Node):
150 items: dict[str, Artefact] = field(default_factory=dict)
151 ensemble_items: dict[str, Artefact] = field(default_factory=dict)
152 types: set[type] = field(default_factory=set)
153 skip_type_check: bool = False
154 ensemble: bool = False
156 def __post_init__(self) -> None:
157 super().__post_init__()
158 self.status = NodeStatus.READY if self.items else NodeStatus.NOT_READY
160 def add_item(self, item: Artefact) -> None:
161 try:
162 self.items[item.id] = item
163 self.status = NodeStatus.READY
164 except Exception as e:
165 raise RuntimeError(f"Error adding item to ArtefactNode: {str(e)}")
167 def get_item(self, item_id: str) -> Artefact | None:
168 return self.items.get(item_id, None)
170 def add_type(self, type_: type) -> None:
171 self.types.add(type_)
173 def check_type(self, type_: type) -> bool:
174 for t in self.types:
175 if issubclass(type_, t):
176 return True
177 return False
179 def __str__(self) -> str:
180 types_str = "\n".join([t.get_class_path() for t in self.types])
181 return f"{self.name}\n{types_str}"
183 def evaluate(
184 self, gold_items: "ArtefactNode", threshold: float = 0.5
185 ) -> dict[str, int | float]:
186 """
187 Evaluate predicted items against gold standard items.
189 This implements the Hungarian algorithm (optimal bipartite matching) that:
190 1. Validates type compatibility between predicted and gold items
191 2. Computes pairwise scores between all predicted and gold items
192 3. Finds optimal assignment that maximizes total matching score
193 4. Filters matches below threshold
194 5. Computes aggregate metrics (precision, recall, F1, mean score)
196 Args:
197 gold_items: ArtefactNode containing ground truth items
198 threshold: Minimum score for a valid match (default 0.5)
199 - For ImageCrop: IoU threshold
200 - For Text: similarity threshold
202 Returns:
203 dict: {
204 "precision": float, # TP / (TP + FP)
205 "recall": float, # TP / (TP + FN)
206 "f1": float, # 2 * (P * R) / (P + R)
207 "mean_score": float, # Average score of matched pairs
208 "num_predicted": int, # Total predicted items
209 "num_gold": int, # Total gold items
210 "num_matched": int, # Number of successful matches (TP)
211 "num_false_positives": int, # Predicted items not matched
212 "num_false_negatives": int, # Gold items not matched
213 "matches": [ # List of matched pairs
214 {
215 "predicted_id": str,
216 "gold_id": str,
217 "score": float
218 },
219 ...
220 ],
221 "unmatched_predicted": [str], # IDs of unmatched predictions
222 "unmatched_gold": [str] # IDs of unmatched gold items
223 }
225 Raises:
226 TypeError: If predicted and gold items are of incompatible types
227 """
229 if threshold > 1.0 or threshold < 0.0:
230 raise ValueError(f"Invalid threshold provided: {threshold}")
232 # Use gold items for type detection when self is empty, so labels with
233 # no predictions still appear in aggregated metrics (e.g. empty [] in YAML)
234 _type_source = self.items if self.items else gold_items.items
235 metrics = {
236 "types": list(set([type(item).__name__ for item in _type_source.values()])),
237 "precision": 0.0,
238 "recall": 0.0,
239 "f1": 0.0,
240 "mean_score": 0.0,
241 "num_predicted": 0,
242 "num_gold": 0,
243 "num_matched": 0,
244 "num_false_positives": 0,
245 "num_false_negatives": 0,
246 "matches": [],
247 "unmatched_predicted": [],
248 "unmatched_gold": [],
249 }
251 # Edge case: Empty inputs with explicit semantic definitions
252 if len(self.items) == 0 and len(gold_items.items) == 0:
253 # Vacuous truth: nothing to predict, nothing to find
254 metrics.update({"precision": 1.0, "recall": 1.0, "f1": 1.0})
255 return metrics
257 if len(self.items) == 0:
258 # No predictions: precision undefined (0/0), defaults to 0.0 for consistency
259 metrics.update(
260 {
261 "precision": 0.0,
262 "num_gold": len(gold_items.items),
263 "num_false_negatives": len(gold_items.items),
264 "unmatched_gold": list(gold_items.items.keys()),
265 }
266 )
267 return metrics
269 if len(gold_items.items) == 0:
270 # No gold items: zero precision (all FPs), perfect recall (nothing to miss)
271 metrics.update(
272 {
273 "recall": 1.0,
274 "num_predicted": len(self.items),
275 "num_false_positives": len(self.items),
276 "unmatched_predicted": list(self.items.keys()),
277 }
278 )
279 return metrics
281 # Step 0: Validate type compatibility upfront
282 predicted_types = set([type(value).__name__ for value in self.items.values()])
283 gold_types = set([type(value).__name__ for value in gold_items.items.values()])
284 compatible = predicted_types == gold_types
285 if not compatible:
286 raise TypeError(
287 f"Cannot evaluate {predicted_types} predictions against "
288 f"{gold_types} gold items. Types must match."
289 )
291 # Step 1: Build score matrix
292 # Shape: (num_predicted, num_gold)
293 predicted_ids = list(self.items.keys())
294 gold_ids = list(gold_items.items.keys())
295 n_pred = len(self.items.keys())
296 n_true = len(gold_items.items.keys())
297 score_matrix = np.zeros((n_pred, n_true))
299 if n_true > n_pred:
300 padding = np.full([n_true - n_pred, n_true], 1e-9)
301 score_matrix = np.vstack([score_matrix, padding])
302 elif n_true < n_pred:
303 padding = np.full([n_pred, n_pred - n_true], 1e-9)
304 score_matrix = np.hstack([score_matrix, padding])
306 for pred_idx, pred_id in enumerate(predicted_ids):
307 pred_item = self.items[pred_id]
308 for gold_idx, gold_id in enumerate(gold_ids):
309 try:
310 gold_item = gold_items.items[gold_id]
311 score = pred_item.evaluate(gold_item)
312 score_matrix[pred_idx, gold_idx] = score
313 except NotImplementedError as e:
314 logger.debug(f"Skipping evaluation for {pred_id} vs {gold_id}: {e}")
316 # Step 2: Hungarian algorithm for optimal bipartite matching
317 # Convert to cost matrix (minimize cost = maximize score)
318 # scipy.optimize.linear_sum_assignment minimizes total cost
319 # If cost_matrix has a shape of 1, skip this step
320 if score_matrix.size != 1:
321 from scipy.optimize import linear_sum_assignment
323 cost_matrix = 1.0 - score_matrix
325 # Note: When multiple predicted items have identical scores against the same
326 # gold item, the assignment is deterministic but arbitrary (based on internal
327 # algorithm traversal order). All such assignments are equally optimal.
328 try:
329 pred_indices, gold_indices = linear_sum_assignment(cost_matrix)
330 except Exception as e:
331 raise RuntimeError(f"Error during optimal assignment: {str(e)}")
332 else:
333 pred_indices = np.array([0])
334 gold_indices = np.array([0])
336 # Step 3: Filter matches by threshold and record results
337 matches = []
338 matched_predicted = set()
339 matched_gold = set()
341 for pred_idx, gold_idx in zip(pred_indices, gold_indices):
342 if pred_idx >= n_pred or gold_idx >= n_true:
343 continue
344 score = score_matrix[pred_idx, gold_idx]
345 # Only accept matches above the threshold
346 if score >= threshold:
347 predicted_id = predicted_ids[pred_idx]
348 gold_id = gold_ids[gold_idx]
349 matches.append(
350 {
351 "predicted_id": predicted_id,
352 "gold_id": gold_id,
353 "score": float(score),
354 }
355 )
356 matched_predicted.add(predicted_id)
357 matched_gold.add(gold_id)
359 # Step 4: Identify unmatched items
360 unmatched_predicted = [
361 predicted_id
362 for predicted_id in predicted_ids
363 if predicted_id not in matched_predicted
364 ]
366 unmatched_gold = [
367 gold_id for gold_id in gold_ids if gold_id not in matched_gold
368 ]
370 # Step 5: Compute metrics
371 num_true_positives = len(matches)
372 num_false_positives = len(unmatched_predicted)
373 num_false_negatives = len(unmatched_gold)
375 precision = (
376 num_true_positives / (num_true_positives + num_false_positives)
377 if (num_true_positives + num_false_positives) > 0
378 else 0.0
379 )
381 recall = (
382 num_true_positives / (num_true_positives + num_false_negatives)
383 if (num_true_positives + num_false_negatives) > 0
384 else 0.0
385 )
387 f1 = (
388 2 * (precision * recall) / (precision + recall)
389 if (precision + recall) > 0
390 else 0.0
391 )
393 mean_score = (
394 sum(match["score"] for match in matches) / len(matches) if matches else 0.0
395 )
397 # Step 6: Return comprehensive evaluation results
398 metrics.update(
399 {
400 "precision": precision,
401 "recall": recall,
402 "f1": f1,
403 "mean_score": mean_score,
404 "num_predicted": len(self.items),
405 "num_gold": len(gold_items.items),
406 "num_matched": num_true_positives,
407 "num_false_positives": num_false_positives,
408 "num_false_negatives": num_false_negatives,
409 "matches": matches,
410 "unmatched_predicted": unmatched_predicted,
411 "unmatched_gold": unmatched_gold,
412 }
413 )
414 return metrics
416 def _create_instance(
417 self, cls_: type, *, source_file: Path | None = None, **item
418 ) -> None:
419 try:
420 instance = cls_(**item)
421 if not instance:
422 raise ValueError(f"Failed to load {item} with {cls_}")
423 instance.source_file = source_file
424 self.add_item(instance)
425 except Exception as e:
426 logger.error(
427 "Failed to create instance of %s: %s", cls_.__name__, e, exc_info=True
428 )
430 def _create_ensemble_instances(self, base_path: Path, cls_: type, **item) -> None:
431 try:
432 item["id"] = f"{base_path.parent}::{item['id']}"
433 if "parents" in item:
434 for parent_idx, _ in enumerate(item["parents"]):
435 item["parents"][
436 parent_idx
437 ] = f"{base_path.parent}::{item['parents'][parent_idx]}"
438 instance = cls_(**item)
439 if not instance:
440 raise ValueError(f"Failed to load {item} with {cls_}")
441 self.ensemble_items[instance.id] = instance
442 self.status = NodeStatus.READY
443 except Exception as e:
444 raise RuntimeError(f"Error creating ensemble instances: {str(e)}")
446 def _create_instances(self, **item) -> None:
447 for cls_ in self.types:
448 self._create_instance(cls_, **item)
450 def process(
451 self,
452 key: str,
453 value: str | Path | None = None,
454 **kwargs,
455 ) -> None:
456 skip_type_check: bool = kwargs.get("skip_type_check", False)
457 ensemble: bool = kwargs.get("ensemble", False)
458 value = value if value else kwargs.get("file", None)
459 if value and Path(value).exists() and Path(value).is_dir():
460 value = Path(value)
461 with change_dir(value):
462 try:
463 result_file = Path("results.yaml")
464 with open(result_file, "r") as f:
465 results: dict = yaml.safe_load(f)
466 partition = results.pop(
467 "collectra_results_metadata", dict()
468 ).get("partition", None)
469 if key not in results:
470 raise ValueError(
471 f"[red]{key}[/red] could not be found in {value}"
472 )
473 data = results[key]
474 if not data:
475 # Key exists but empty — pipeline processed it, found nothing
476 self.status = NodeStatus.READY
477 return
478 data = data if isinstance(data, list) else [data]
479 for item in data:
480 primitive_type = False
481 try:
482 primitive_type = not isinstance(item, dict) or not (
483 "type" in item
484 and (
485 "path" in item
486 or "data" in item
487 or item["type"].endswith(".Link")
488 )
489 )
490 if primitive_type:
491 raise ValueError(
492 f"Item must be a complex type with 'type' and ('path' or 'data')"
493 )
494 cls_ = load_class_from_string(item.pop("type"))
495 if not (
496 skip_type_check or ensemble
497 ) and not self.check_type(cls_):
498 raise TypeError(
499 f"{cls_} is not a subclass or not defined in {self.types}"
500 )
501 item["name"] = key
502 if (
503 "data" not in item
504 and "path" in item
505 and "path" not in cls_.__dataclass_fields__
506 ):
507 item["data"] = item.pop("path")
508 if "parents" in item and not isinstance(
509 item["parents"], list
510 ):
511 item["parents"] = [item["parents"]]
512 if (
513 partition is not None
514 and "partition" in cls_.all_attributes()
515 ):
516 item["partition"] = partition
517 if ensemble:
518 self._create_ensemble_instances(value, cls_, **item)
519 else:
520 self._create_instance(
521 cls_, source_file=value, **item
522 )
523 except Exception as e:
524 if not primitive_type:
525 logger.error(
526 "Error processing item in %s: %s",
527 self.name,
528 e,
529 exc_info=True,
530 )
531 else:
532 self._create_instances(name=key, data=str(item))
534 except Exception as e:
535 logger.error(
536 "Error loading data for %s from %s: %s",
537 key,
538 value,
539 e,
540 exc_info=True,
541 )
542 elif key:
543 self._create_instances(name=key, data=value)
545 @staticmethod
546 def batch_process(
547 item_file: Path,
548 artefact_nodes: list["ArtefactNode"],
549 *,
550 include_unlabelled: bool = False,
551 ) -> list[Artefact]:
552 names = [artefact_node.name for artefact_node in artefact_nodes]
553 with change_dir(item_file):
554 try:
555 data: list[Artefact] = list()
556 result_file = Path("results.yaml")
557 if not result_file.exists():
558 print(error_msg(f"Invalid file: {Path.cwd()}. Ignoring..."))
559 return data
560 with open(result_file, "r") as f:
561 file_data: dict = yaml.safe_load(f)
562 partition = file_data.get("collectra_results_metadata", dict()).get(
563 "partition", None
564 )
565 all_names_not_found = all(name not in file_data for name in names)
566 # Object detection can opt into base images with no annotations
567 # as negative samples. Classification must not treat an
568 # unlabelled image as a training example.
569 if all_names_not_found and include_unlabelled:
570 print(
571 f"[yellow]No matching labels found in [blue]{item_file}[/blue] "
572 f"for names: {', '.join(names)}. Using its base image as an "
573 "unlabelled detection sample."
574 )
575 found_base = False
576 for name, values in file_data.items():
577 values = values if isinstance(values, list) else [values]
578 for item in values:
579 if (
580 not isinstance(item, dict)
581 or "type" not in item
582 and ("data" not in item or "path" not in item)
583 ):
584 continue
585 cls_str = item.pop("type", "")
586 if cls_str == "collectra.Image":
587 cls_ = load_class_from_string(cls_str)
588 try:
589 item["name"] = name
590 item["data"] = (
591 item.pop("path")
592 if "path" in item
593 else item["data"]
594 )
595 if partition is not None:
596 item["partition"] = partition
597 instance = cls_(**item)
598 if not instance:
599 raise Warning(
600 f"Failed to load {item} with {cls_}"
601 )
602 found_base = True
603 data.append(instance)
604 break
605 except Exception as e:
606 traceback_error(
607 e,
608 f"Failed to load data item {name} from {item.get('data', '')}",
609 verbose=True,
610 )
611 if found_base:
612 break
613 elif all_names_not_found:
614 return data
615 for i, name in enumerate(names):
616 if name not in file_data:
617 continue
618 value = file_data[name]
619 if not value:
620 raise Warning(
621 f"Artefact seems to be empty for {name} in {item_file}. Provided: {value}"
622 )
623 value = value if isinstance(value, list) else [value]
624 for item in value:
625 if not isinstance(item, dict) or not (
626 "type" in item
627 and (
628 "data" in item
629 or "path" in item
630 or item["type"].endswith(".Link")
631 )
632 ):
633 continue
634 cls_ = load_class_from_string(item.pop("type"))
635 match = False
636 for type_ in artefact_nodes[i].types:
637 if issubclass(cls_, type_) or cls_ == type_:
638 match = True
639 break
640 if not match:
641 continue
642 item["name"] = name
643 if "path" in item and "path" not in cls_.__dataclass_fields__:
644 item["data"] = item.pop("path")
645 if partition is not None:
646 item["partition"] = partition
647 try:
648 instance = cls_(**item)
649 if instance:
650 instance.source_file = item_file
651 data.append(instance)
652 except Exception as e:
653 traceback_error(
654 e,
655 f"Failed to load data item {name} from {item.get('data', '')}: {e}",
656 )
657 return data
658 except Exception as e:
659 traceback_error(e, verbose=True)
660 return list()
663def load_artefact(item: dict[str, Any]):
664 artefact_class = load_class_from_string(item.pop("type"))
666 return artefact_class(**item)