Coverage for collectra/pipelines/ensemble.py: 49.39%
245 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Ensemble processing module for Collectra pipelines.
3This module provides functionality for ensembling multiple collectra result files
4into a single combined result using various techniques:
5- Weighted Box Fusion (WBF) for ImageCrop ensembling
6- Edit distance clustering for Text ensembling
7- IoU-based matching for provenance tracking
9Classes:
10 EnsembleProcessor: Main class for processing ensemble operations
12Functions:
13 calculate_iou: Calculate Intersection over Union between two boxes
14 find_centroid_text: Find the medoid text in a cluster
15 cluster_texts: Cluster similar texts using single-linkage clustering
16 find_contributing_boxes: Find source boxes that contributed to a fused box
17"""
19__all__ = [
20 "EnsembleProcessor",
21 "calculate_iou",
22 "find_centroid_text",
23 "cluster_texts",
24 "find_contributing_boxes",
25]
27import datetime
28import shutil
29from pathlib import Path
31import editdistance
32import yaml
33from rich.progress import track
35from ..types.base import Artefact, ArtefactNode
36from ..types.images import Image, ImageCrop
37from ..types.texts import Text
39# =============================================================================
40# Pure Utility Functions (no class dependencies)
41# =============================================================================
44def weighted_boxes_fusion(*args, **kwargs):
45 try:
46 from ensemble_boxes import weighted_boxes_fusion as fuse
47 except ModuleNotFoundError as error:
48 if error.name != "ensemble_boxes":
49 raise
50 raise ImportError(
51 'Box fusion requires optional ML dependencies. Install with: pip install "collectra[ml]"'
52 ) from error
53 return fuse(*args, **kwargs)
56def calculate_iou(box1: list[float], box2: list[float]) -> float:
57 """Calculate IoU (Intersection over Union) between two boxes.
59 Args:
60 box1: Box 1 in format [x1, y1, x2, y2]
61 box2: Box 2 in format [x1, y1, x2, y2]
63 Returns:
64 IoU value between 0 and 1
65 """
66 x1 = max(box1[0], box2[0])
67 y1 = max(box1[1], box2[1])
68 x2 = min(box1[2], box2[2])
69 y2 = min(box1[3], box2[3])
71 inter = max(0, x2 - x1) * max(0, y2 - y1)
72 area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
73 area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
74 union = area1 + area2 - inter
76 return inter / union if union > 0 else 0
79def find_centroid_text(texts: list[str]) -> str:
80 """Find the text with minimum total edit distance to all others (medoid).
82 Args:
83 texts: List of text strings to find centroid from
85 Returns:
86 The text that is most representative of the cluster
87 """
88 if len(texts) == 1:
89 return texts[0]
91 min_total_distance = float("inf")
92 centroid = texts[0]
94 for candidate in texts:
95 total_distance = sum(editdistance.eval(candidate, other) for other in texts)
96 if total_distance < min_total_distance:
97 min_total_distance = total_distance
98 centroid = candidate
100 return centroid
103def cluster_texts(
104 texts: list[str],
105 provenance: list[str],
106 num_clusters: int,
107 similarity_threshold: float,
108) -> list[dict]:
109 """Cluster similar texts and find centroid for each cluster.
111 Uses single-linkage clustering: a text joins a cluster if it's similar
112 to ANY existing member, allowing transitive similarity chains.
114 Args:
115 texts: List of text strings to cluster
116 provenance: List of source references for each text
117 num_clusters: Maximum number of clusters to create
118 similarity_threshold: Minimum similarity (0-1) to consider texts as same cluster
120 Returns:
121 List of cluster dicts with 'centroid' and 'contributors' keys
122 """
123 clusters = []
124 used: set[int] = set()
126 while len(used) < len(texts) and len(clusters) < num_clusters:
127 for i, text in enumerate(texts):
128 if i in used:
129 continue
130 # Start new cluster
131 used.add(i)
132 cluster = [text]
133 prov = [provenance[i]]
135 # Iteratively find all similar texts (single-linkage)
136 # Keep expanding until no more texts can be added
137 for j, other in enumerate(texts):
138 if j in used:
139 continue
140 max_len = max(len(text), len(other)) or 1
141 similarity = 1 - (editdistance.eval(text, other) / max_len)
142 if similarity >= similarity_threshold:
143 cluster.append(other)
144 prov.append(provenance[j])
145 used.add(j)
147 centroid = find_centroid_text(cluster)
148 clusters.append(
149 {
150 "centroid": centroid,
151 "contributors": prov,
152 }
153 )
154 break # Restart fresh to find next cluster
156 return clusters
159def find_contributing_boxes(
160 fused_box: list[float],
161 source_boxes: list[list[float]],
162 provenance: list[str],
163 iou_threshold: float = 0.5,
164) -> list[str]:
165 """Find source boxes that contributed to a fused box via IoU.
167 Args:
168 fused_box: The fused box in format [x1, y1, x2, y2]
169 source_boxes: List of source boxes
170 provenance: List of source references for each box
171 iou_threshold: Minimum IoU to consider as contributor
173 Returns:
174 List of source references that contributed to the fused box
176 Raises:
177 ValueError: If no contributing boxes are found
178 """
179 contributors = list()
180 for box, prov in zip(source_boxes, provenance):
181 if calculate_iou(fused_box, box) >= iou_threshold:
182 contributors.append(prov)
183 if len(contributors) == 0:
184 raise ValueError("No contributing boxes found for fused box")
185 return list(contributors)
188# =============================================================================
189# EnsembleProcessor Class
190# =============================================================================
193class EnsembleProcessor:
194 """Processor for ensembling multiple collectra result files.
196 This class handles the ensemble logic for combining results from multiple
197 sources into a single merged result with provenance tracking.
199 Attributes:
200 node_manager: Node manager for resolving artefact nodes
201 ext: File extension for collectra result folders
202 name: Workflow name
203 version: Workflow version
204 """
206 def __init__(self, node_manager, ext: str, name: str, version: str):
207 """Initialize the EnsembleProcessor.
209 Args:
210 node_manager: Node manager for resolving artefact nodes
211 ext: File extension for collectra result folders
212 name: Workflow name
213 version: Workflow version
214 """
215 self.node_manager = node_manager
216 self.ext = ext
217 self.name = name
218 self.version = version
220 def ensemble(self, folders: list[Path], output: Path) -> None:
221 """Create ensembled files from source folders.
223 Args:
224 folders: List of source folders containing collectra results
225 output: Output folder path for ensembled results
226 """
227 file_linkage = self.create_file_linkage(folders)
228 layers: list[list[ArtefactNode]] = self.generate_ensembling_layers()
229 output_folder = Path(output)
230 for filename, source_paths in track(
231 file_linkage.items(), description="Ensembling files..."
232 ):
233 self._reset_nodes()
234 self.ensemble_single_file(filename, source_paths, layers, output_folder)
235 with open(output_folder / "linkage.yml", "w") as f:
236 yaml.dump(file_linkage, f, sort_keys=False)
238 def create_file_linkage(self, folders: list[Path]) -> dict[str, list[str]]:
239 """Identify unique filenames across folders and create linkage.
241 Args:
242 folders: List of source folders to scan
244 Returns:
245 Dictionary mapping filename to list of source paths
246 Example: {"file1": ["folder1/file1.collectra", "folder2/file1.collectra"]}
247 """
248 file_linkage: dict[str, list[str]] = {}
250 # Scan all folders for files with the target extension
251 for folder in folders:
252 if not folder.exists():
253 continue
255 for item in folder.iterdir():
256 if item.is_dir() and item.name.endswith(self.ext):
257 # Extract the filename without extension
258 filename = item.name
259 file_linkage.setdefault(filename, []).append(str(item))
261 return file_linkage
263 def generate_ensembling_layers(self) -> list[list[ArtefactNode]]:
264 """Generate layers of artefact nodes for ensemble processing.
266 Returns:
267 List of layers, where each layer contains ArtefactNode objects
268 that should be processed together
269 """
270 layers = []
271 all_names = self.node_manager.get_node_names()
272 first_node = self.node_manager.resolve_node(all_names[0])
273 root_nodes = self.node_manager.get_parents(first_node)
274 layers.append(root_nodes)
275 for task_node in self.node_manager.get_task_nodes():
276 output_nodes = self.node_manager.get_children_artefact(task_node)
277 layers.append(output_nodes)
278 return layers
280 def ensemble_single_file(
281 self,
282 filename: str,
283 source_paths: list[str],
284 layers: list[list[ArtefactNode]],
285 output_folder: Path,
286 ) -> None:
287 """Ensemble a single file from source paths and save to output path.
289 Args:
290 filename: Name of the file being ensembled
291 source_paths: List of paths to source collectra folders
292 layers: List of artefact node layers to process
293 output_folder: Output folder for the ensembled result
294 """
295 output_path = output_folder / filename
296 output_path.mkdir(parents=True, exist_ok=True)
297 ensembled_data = dict()
298 for layer in layers:
299 for artefact_node in layer:
300 ensembled_data[artefact_node.name] = self.ensemble_at_layer(
301 artefact_node, source_paths, ensembled_data
302 )
303 full_data = dict()
304 full_data["collectra_results_metadata"] = {
305 "workflow": self.name,
306 "version": self.version,
307 "timestamp": datetime.datetime.now().isoformat(),
308 }
309 full_data.update(ensembled_data)
310 with open(output_path / "results.yaml", "w") as f:
311 for key, value in ensembled_data.items():
312 if value:
313 yaml.dump({key: value}, f, sort_keys=False)
314 f.write("\n")
315 self.copy_artifact_files(Path(source_paths[0]), output_path)
317 def ensemble_at_layer(
318 self, artefact_node: ArtefactNode, source_paths: list[str], ensembled_data: dict
319 ) -> dict | list[dict]:
320 """Ensemble data at a specific layer from source paths.
322 Args:
323 artefact_node: The artefact node to ensemble
324 source_paths: List of source paths to read data from
325 ensembled_data: Previously ensembled data for reference
327 Returns:
328 Ensembled data as dict or list of dicts
329 """
330 artefact_node.ensemble = True
331 for source_path in source_paths:
332 artefact_node.process(
333 artefact_node.name, Path(source_path), skip_type_check=True
334 )
335 label_data = self.ensemble_artefact_node(artefact_node, ensembled_data)
336 if label_data is None:
337 return []
338 if isinstance(label_data, list) and len(label_data) == 1:
339 return label_data[0]
340 else:
341 return label_data
343 def ensemble_artefact_node(
344 self, artefact_node: ArtefactNode, ensembled_data: dict
345 ) -> list[dict] | dict | None:
346 """Ensemble an artefact node based on its item types.
348 Args:
349 artefact_node: The artefact node containing ensemble items
350 ensembled_data: Previously ensembled data for parent resolution
352 Returns:
353 Ensembled result or None if no items
354 """
355 items: dict[str, Artefact] = artefact_node.ensemble_items
357 if not items:
358 return None
360 if all(type(item) == Image for item in items.values()):
361 first_item: Image = list(items.values())[0]
362 return {
363 "type": "collectra.Image",
364 "id": f"{artefact_node.name}_ensemble",
365 "data": first_item.get_path().name,
366 "ensemble": list(items.keys()),
367 }
369 elif all(type(item) == ImageCrop for item in items.values()):
370 return self.ensemble_image_crops(artefact_node, ensembled_data)
372 elif all(type(item) == Text for item in items.values()):
373 return self.ensemble_text(artefact_node, ensembled_data)
375 else:
376 raise Warning(
377 f"Ensembling not supported for data type in node {artefact_node.name}"
378 )
380 def ensemble_image_crops(
381 self,
382 artefact_node: ArtefactNode,
383 ensembled_data: dict,
384 ) -> dict | list[dict]:
385 """Ensemble ImageCrop items using Weighted Box Fusion.
387 Args:
388 artefact_node: The artefact node containing ImageCrop items
389 ensembled_data: Previously ensembled data for parent resolution
391 Returns:
392 Ensembled ImageCrop result(s)
393 """
394 sources = list(
395 set([key.split("::")[0] for key in artefact_node.ensemble_items.keys()])
396 )
398 source_crops = dict()
400 for idx, item in artefact_node.ensemble_items.items():
401 for source in sources:
402 if not idx.startswith(source):
403 continue
404 source_crops.setdefault(source, []).append(item)
406 # Collect boxes for WBF
407 boxes_per_source: list[list[list[float]]] = []
408 scores_per_source: list[list[float]] = []
409 labels_per_source: list[list[int]] = []
410 prov_per_source: list[list[str]] = []
411 data_source: str = ""
413 for crops in source_crops.values():
414 src_boxes, src_scores, src_labels, src_prov = [], [], [], []
415 for crop in crops:
416 # Convert center-format to corner-format (normalized 0-1)
417 x1 = max(0, crop.x_center - crop.width_relative / 2)
418 y1 = max(0, crop.y_center - crop.height_relative / 2)
419 x2 = min(1, crop.x_center + crop.width_relative / 2)
420 y2 = min(1, crop.y_center + crop.height_relative / 2)
422 src_boxes.append([x1, y1, x2, y2])
423 src_scores.append(1.0)
424 src_labels.append(0)
425 src_prov.append(crop.id)
426 if data_source == "":
427 data_source = crop.get_path().name
428 else:
429 if data_source == crop.get_path().name:
430 continue
431 raise ValueError("Mismatched data sources in crops for ensembling")
433 boxes_per_source.append(src_boxes)
434 scores_per_source.append(src_scores)
435 labels_per_source.append(src_labels)
436 prov_per_source.append(src_prov)
438 # WBF call
439 weights = [1] * len(boxes_per_source)
440 fused_boxes, _, _ = weighted_boxes_fusion(
441 boxes_per_source,
442 scores_per_source,
443 labels_per_source,
444 weights=weights,
445 iou_thr=0.5,
446 skip_box_thr=0.0,
447 )
449 # Build result
450 result = []
451 flat_prov = [p for sp in prov_per_source for p in sp]
452 flat_boxes = [b for sb in boxes_per_source for b in sb]
454 for i, fused_box in enumerate(fused_boxes):
455 # Convert back to center format
456 x1, y1, x2, y2 = fused_box
457 x_center = (x1 + x2) / 2
458 y_center = (y1 + y2) / 2
459 width_rel = x2 - x1
460 height_rel = y2 - y1
462 # Find contributing boxes via IoU
463 ensemble_refs = find_contributing_boxes(
464 list(fused_box), flat_boxes, flat_prov, iou_threshold=0.5
465 )
467 ensemble_ref_nodes = [
468 artefact_node.ensemble_items[key]
469 for key in artefact_node.ensemble_items.keys()
470 if key in ensemble_refs
471 ]
473 # Resolve parents from contributing crops
474 ensembled_parents = self.resolve_ensemble_parents(
475 ensemble_ref_nodes, ensembled_data
476 )
478 result.append(
479 {
480 "type": "collectra.ImageCrop",
481 "id": f"{artefact_node.name}_ensembled_{i+1}",
482 "data": data_source,
483 "ensemble": ensemble_refs,
484 "parents": ensembled_parents,
485 "x_center": float(x_center),
486 "y_center": float(y_center),
487 "width_relative": float(width_rel),
488 "height_relative": float(height_rel),
489 }
490 )
492 return result[0] if len(result) == 1 else result
494 def ensemble_text(
495 self, artefact_node: ArtefactNode, ensembled_data: dict
496 ) -> dict | list[dict]:
497 """Ensemble Text items using edit distance clustering.
499 Args:
500 artefact_node: The artefact node containing Text items
501 ensembled_data: Previously ensembled data for parent resolution
503 Returns:
504 Ensembled Text result(s) with provenance
505 """
506 sources = list(
507 set([key.split("::")[0] for key in artefact_node.ensemble_items.keys()])
508 )
509 source_texts: dict[str, list] = {}
511 for idx, item in artefact_node.ensemble_items.items():
512 for source in sources:
513 if idx.startswith(source):
514 source_texts.setdefault(source, []).append(item)
516 # Flatten texts with provenance
517 flat_texts = []
518 flat_prov = []
519 src_for_clustering = dict()
520 for key, texts in source_texts.items():
521 for text in texts:
522 src_for_clustering.setdefault(key, []).append(text)
523 flat_texts.append(text.data)
524 flat_prov.append(text.id)
526 num_clusters = 0
527 for items in src_for_clustering.values():
528 if len(items) > num_clusters:
529 num_clusters = len(items)
531 # Cluster similar texts and find centroids
532 clusters = cluster_texts(
533 flat_texts, flat_prov, num_clusters, similarity_threshold=0.7
534 )
536 # Build results
537 result = []
538 for i, cluster in enumerate(clusters):
539 ensemble_refs = cluster["contributors"]
541 ensemble_ref_nodes = [
542 artefact_node.ensemble_items[key]
543 for key in artefact_node.ensemble_items.keys()
544 if key in ensemble_refs
545 ]
547 ensembled_parents = self.resolve_ensemble_parents(
548 ensemble_ref_nodes, ensembled_data
549 )
551 result.append(
552 {
553 "type": "collectra.Text",
554 "id": f"{artefact_node.name}_ensembled_{i+1}",
555 "ensemble": ensemble_refs,
556 "parents": ensembled_parents,
557 "data": cluster["centroid"],
558 }
559 )
561 return result[0] if len(result) == 1 else result
563 def resolve_ensemble_parents(
564 self,
565 ensemble_ref_nodes: list[Artefact],
566 ensemble_data: dict,
567 ) -> list[str]:
568 """Resolve parent references from ensemble data.
570 Args:
571 ensemble_ref_nodes: List of data objects with parent references
572 ensemble_data: Previously ensembled data to search for parents
574 Returns:
575 List of parent IDs or single parent ID if only one
576 """
577 ensembled_parents = set()
578 for ensemble in ensemble_ref_nodes:
579 parents = ensemble.parents
580 for parent in parents:
581 for key in ensemble_data:
582 if isinstance(ensemble_data[key], list):
583 for item in ensemble_data[key]:
584 if "ensemble" in item and parent in item["ensemble"]:
585 ensembled_parents.add(item["id"])
586 else:
587 if (
588 "ensemble" in ensemble_data[key]
589 and parent in ensemble_data[key]["ensemble"]
590 ):
591 ensembled_parents.add(ensemble_data[key]["id"])
592 ensembled_parents = list(ensembled_parents)
593 ensembled_parents = (
594 ensembled_parents[0] if len(ensembled_parents) == 1 else ensembled_parents
595 )
596 return ensembled_parents
598 def copy_artifact_files(self, source_path: Path, dest_path: Path) -> None:
599 """Copy artifact files from first source to ensemble folder.
601 Copies all files except results.yaml from the source collectra folder
602 to the destination folder.
604 Args:
605 source_path: Path to first source collectra folder
606 dest_path: Path to ensemble collectra folder
607 """
608 if not source_path.exists():
609 return
610 for item in source_path.iterdir():
611 # Skip results.yaml
612 if item.name == "results.yaml":
613 continue
614 # Copy artifact files
615 if item.is_file():
616 dest_file = dest_path / item.name
617 shutil.copy2(item, dest_file)
619 def _reset_nodes(self) -> None:
620 """Reset all artefact nodes for fresh processing."""
621 self.node_manager.reset_all_nodes()