Coverage for collectra/gui/data_display.py: 97.98%
247 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""
2Parses YAML structure, builds lineage relationships, and computes display values
3for ImageCrop annotations based on the rules:
51. Container crop (has ImageCrop children) → display blank
62. Leaf crop with Text child → traverse to deepest Text → display that text
73. Leaf crop without Text child → display blank
9"""
11import copy
12import uuid
13from dataclasses import dataclass, field
15import networkx as nx
16from pydantic import BaseModel, ConfigDict
18from collectra.commons.files import CollectraFile, CollectraResultsMetadata
20from .utils import normalise_items
23class CollectraNodeFactory:
25 @staticmethod
26 def create_node(data: dict) -> "CollectraNode":
27 """Factory method to create appropriate CollectraNode subclass."""
28 node_type = data.get("type", "")
29 from collectra.utils import load_class_from_string
30 from collectra.types.images import ImageCrop
32 try:
33 is_crop = issubclass(load_class_from_string(node_type), ImageCrop)
34 except (ImportError, AttributeError, ValueError):
35 is_crop = "ImageCrop" in node_type
36 if is_crop:
37 crop_region = CollectraCropRegion(
38 x_center=data["x_center"],
39 y_center=data["y_center"],
40 width_relative=data["width_relative"],
41 height_relative=data["height_relative"],
42 )
43 return CollectraAnnotationNode(**data, crop_region=crop_region)
44 return CollectraNode(**data)
47class CollectraNode(BaseModel):
48 """Node in the annotation graph."""
50 model_config = ConfigDict(extra="allow")
52 label: str
53 type: str
54 id: str
55 parents: list[str] = []
56 data: object = ""
57 embeddings: list[str] = []
58 orientation: str = "north"
59 name: str = "" # user-set per-instance name, distinct from label (the category)
61 def model_dump(self, *args, **kwargs) -> dict[str, str | float]:
62 data = super().model_dump(*args, **kwargs)
63 label = data.pop("label", None) # Exclude label from dump
64 if "path" in data and "data" not in self.model_fields_set:
65 data.pop("data", None)
66 if data["parents"]:
67 data["parents"] = (
68 data["parents"][0] if len(data["parents"]) == 1 else data["parents"]
69 )
70 else:
71 data.pop("parents", None)
72 if not data.get("name"):
73 data.pop("name", None) # Don't clutter files that never set one
74 if label is None:
75 raise ValueError("Label is missing from node data.")
76 return data
79class CollectraCropRegion(BaseModel):
80 """Crop region fields for annotation nodes."""
82 x_center: float
83 y_center: float
84 width_relative: float
85 height_relative: float
88class CollectraAnnotationNode(CollectraNode):
89 """Annotation node in the graph."""
91 crop_region: CollectraCropRegion
93 @property
94 def crop(self) -> dict[str, float]:
95 """Return crop region as a dict."""
96 return self.crop_region.model_dump()
98 def model_dump(self, *args, **kwargs) -> dict[str, str | float]:
99 data = super().model_dump(*args, **kwargs)
100 data.pop("crop_region", None) # Remove crop_region field
101 for key, value in self.crop.items():
102 data[key] = value
103 return data
106class NodeDisplayValue(BaseModel):
107 """Display value for a node."""
109 reason: str
110 value: str | None = None
111 source_id: str | None = None
112 crop_region: dict[str, float] | None = None
113 locked: bool = False
116@dataclass
117class CollectraGraph:
118 """
119 Graph structure for annotation lineage relationships.
121 Provides clean interface for traversal without external dependencies.
122 """
124 _graph: nx.DiGraph = field(default_factory=nx.DiGraph)
126 metadata: CollectraResultsMetadata = field(default_factory=CollectraResultsMetadata)
128 # Cached indexes for O(1) lookups
129 _children_index: dict[str, list[str]] = field(default_factory=dict, repr=False)
130 _parents_index: dict[str, list[str]] = field(default_factory=dict, repr=False)
132 # Labels seen in the source YAML, in original order — including ones whose
133 # value was `[]` (a pipeline task ran and found nothing). Those never
134 # become graph nodes, so nothing else remembers they exist; without this,
135 # to_data() silently drops them on every save.
136 _known_labels: list[str] = field(default_factory=list, repr=False)
138 @property
139 def nodes(self) -> list[str]:
140 """List of node IDs in the graph."""
141 if self._graph is None:
142 return []
143 return list(self._graph.nodes)
145 @classmethod
146 def from_yaml_data(cls, data: dict) -> "CollectraGraph":
147 """Build graph from parsed YAML data."""
148 data = copy.deepcopy(data)
149 metadata = CollectraResultsMetadata(
150 **data.pop("collectra_results_metadata", {})
151 )
152 return cls._from_data(data, metadata)
154 @classmethod
155 def from_collectra_file(cls, collectra_file: CollectraFile) -> "CollectraGraph":
156 """Build a display graph backed by a loaded Collectra file."""
157 return cls._from_data(
158 copy.deepcopy(collectra_file.data),
159 collectra_file.collectra_results_metadata,
160 )
162 @classmethod
163 def _from_data(
164 cls, data: dict, metadata: CollectraResultsMetadata
165 ) -> "CollectraGraph":
166 graph = cls(metadata=metadata)
168 for label, value in data.items():
169 graph._known_labels.append(label)
170 items = normalise_items(value)
171 for item in items:
172 if isinstance(item, dict) and "id" in item:
173 item["label"] = label # Store label for reverse lookup
174 graph.add_node(item)
176 return graph
178 def add_node(self, data: dict) -> None:
179 """Add a node and its parent edges."""
180 data["parents"] = normalise_items(data.get("parents", []))
181 node = CollectraNodeFactory.create_node(data)
182 self._graph.add_node(
183 node.id,
184 node=node,
185 )
187 for parent in data["parents"]:
188 self._graph.add_edge(parent, node.id)
190 def children(self, node_id: str) -> list[str]:
191 """Get immediate children of a node."""
192 return list(self._graph.successors(node_id))
194 def parents(self, node_id: str) -> list[str]:
195 """Get immediate parents of a node."""
196 return list(self._graph.predecessors(node_id))
198 def get_node(self, node_id: str) -> CollectraNode | None:
199 """Get the AnnotationNode object for a given node ID."""
200 node = self._graph.nodes.get(node_id, {}).get("node", None)
201 return node
203 def resolve_id(self, node_id: str) -> str:
204 """Translate a DAG label (e.g. "formation") to its real internal node id
205 (e.g. "formation-<uuid>"). The DAG only ever knows the label; nodes are
206 keyed internally by id. Returns node_id unchanged if it's already a real
207 id, or if no node with that label exists.
208 """
209 if node_id in self._graph.nodes:
210 return node_id
211 for candidate_id in self._graph.nodes:
212 node = self.get_node(candidate_id)
213 if node and node.label == node_id:
214 return candidate_id
215 return node_id
217 def get_type(self, node_id: str) -> str:
218 """Get type of a node."""
219 node = self.get_node(node_id)
220 return node.type if node else ""
222 def get_label(self, node_id: str) -> str:
223 """Get label of a node."""
224 node = self.get_node(node_id)
225 return node.label if node else ""
227 def get_data(self, node_id: str) -> str:
228 """Get data field of a node."""
229 node = self.get_node(node_id)
230 return node.data if node else ""
232 def get_name(self, node_id: str) -> str:
233 """Get name field of a node."""
234 node = self.get_node(node_id)
235 return node.name if node else ""
237 def get_crop_region(self, node_id: str) -> dict[str, float]:
238 """
239 Get the crop region coordinates of a node.
241 Args:
242 node_id: The node ID to get crop region from
244 Returns:
245 dict with x_center, y_center, width_relative, height_relative
247 Raises:
248 ValueError: If node is missing crop region fields
249 """
250 node = self.get_node(node_id)
251 if not isinstance(node, CollectraAnnotationNode):
252 raise ValueError(f"Node {node_id} is not an annotation node.")
254 return node.crop
256 def get_unique_labels(self) -> list[str]:
257 """Get all unique labels from nodes in the graph."""
258 labels = set()
259 for node_id in self._graph.nodes:
260 node = self.get_node(node_id)
261 if node and node.label:
262 labels.add(node.label)
263 return sorted(labels)
265 def count_nodes_by_label(self, type_filter: str = None) -> dict[str, int]:
266 """Count nodes grouped by label, optionally filtered by type."""
267 counts = {}
268 for node_id in self._graph.nodes:
269 node = self.get_node(node_id)
270 if node and node.label:
271 if type_filter is None or node.type == type_filter:
272 counts[node.label] = counts.get(node.label, 0) + 1
273 return counts
275 def set_label(self, node_id: str, label: str) -> None:
276 """Set label field of a node."""
277 node = self.get_node(node_id)
278 if node is None:
279 raise ValueError(f"Node {node_id} not found in graph.")
280 node.label = label
282 def set_name(self, node_id: str, name: str) -> None:
283 """Set name field of a node."""
284 node = self.get_node(node_id)
285 if node is None:
286 raise ValueError(f"Node {node_id} not found in graph.")
287 node.name = name
289 def set_data(self, node_id: str, data: str, crop_id: str | None = None) -> None:
290 """Set data field of a node."""
291 if node_id:
292 node = self.get_node(node_id)
293 if node is None:
294 raise ValueError(f"Node {node_id} not found in graph.")
295 node.data = data
296 return
297 if not crop_id:
298 raise ValueError(f"Node {node_id} not found and crop_id not provided.")
299 text_data = {
300 "label": "user_annotation_text",
301 "type": "collectra.Text",
302 "id": f"user_text_{uuid.uuid4().hex[:8]}",
303 "parents": crop_id,
304 "data": data,
305 }
306 self.add_node(text_data)
308 def set_crop_region(self, node_id: str, crop_region: dict[str, float]) -> None:
309 """Set crop region fields of a node (x_center, y_center, width_relative, height_relative)."""
310 node = self.get_node(node_id)
311 crop: CollectraCropRegion = CollectraCropRegion(**crop_region)
312 if not isinstance(node, CollectraAnnotationNode):
313 raise ValueError(f"Node {node_id} is not an annotation node.")
314 node.crop_region = crop
316 def remove_node(self, node_id: str) -> None:
317 """Remove a node and all edges connected to it."""
318 self._graph.remove_node(node_id)
320 def rename_id(self, old_id: str, new_id: str) -> None:
321 """Changes a node's id, re-pointing every parent/child edge to match (nx.relabel_nodes)."""
322 if old_id not in self._graph.nodes:
323 raise ValueError(f"Node {old_id} not found in graph.")
324 if new_id in self._graph.nodes:
325 raise ValueError(f"Node {new_id} already exists in graph.")
326 nx.relabel_nodes(self._graph, {old_id: new_id}, copy=False)
327 node = self.get_node(new_id)
328 if node:
329 node.id = new_id
331 def children_of_type(self, node_id: str, type_substr: str) -> list[str]:
332 """Get immediate children containing type_substr in their type."""
333 return [
334 child_id
335 for child_id in self.children(node_id)
336 if type_substr in self.get_type(child_id)
337 ]
339 def dfs_leaves(self, start: str, type_filter: str) -> list[str]:
340 """
341 DFS traversal following only edges to nodes matching type_filter.
342 Returns leaf nodes (nodes with no children of the filtered type).
343 """
344 leaves = []
345 visited = set()
346 stack = [start]
348 while stack:
349 node_id = stack.pop()
350 if node_id in visited:
351 continue
352 visited.add(node_id)
354 # Get children matching the type filter
355 typed_children = self.children_of_type(node_id, type_filter)
357 if not typed_children:
358 # This is a leaf in the filtered subgraph
359 leaves.append(node_id)
360 else:
361 stack.extend(typed_children)
363 return leaves
365 def find_deepest(self, start: str, type_filter: str) -> str:
366 """
367 Find the deepest node reachable from start following only type_filter edges.
368 Returns the first leaf found (DFS order).
369 """
370 leaves = self.dfs_leaves(start, type_filter)
371 return leaves[0] if leaves else ""
373 def to_yaml_data(self) -> dict:
374 """Convert graph back to YAML data structure."""
375 return {
376 "collectra_results_metadata": self.metadata.model_dump(),
377 **self.to_data(),
378 }
380 def to_data(self) -> dict:
381 """Convert graph nodes to the data held by a CollectraFile."""
382 yaml_data: dict = {}
384 for node_id in self._graph.nodes:
385 node = self.get_node(node_id)
386 if node is None:
387 continue
388 node_data = node.model_dump()
389 yaml_data.setdefault(node.label, []).append({**node_data})
391 def _collapse(value: list[dict]) -> list[dict] | dict:
392 return value[0] if len(value) == 1 else value
394 final_data: dict[str, list[dict] | dict] = {}
396 # Known labels first, in original order, so a label with no current
397 # nodes (never had any, or had all of them deleted) round-trips as
398 # `[]` instead of disappearing.
399 for label in self._known_labels:
400 final_data[label] = _collapse(yaml_data.pop(label, []))
402 # Anything left is a label created fresh this session (e.g. a new
403 # user-added Text node) that wasn't in the original data.
404 for key, value in yaml_data.items():
405 final_data[key] = _collapse(value)
407 return final_data
409 def compute_display_value(self, node_id: str) -> NodeDisplayValue:
410 """
411 Compute display value for an annotation.
413 Args:
414 node_id: The node ID to compute display value for
416 Returns:
417 NodeDisplayValue with value, source_id, crop_region, reason
418 """
419 if self.get_node(node_id) is None:
420 return NodeDisplayValue(reason=f"{node_id} Element not found")
422 node_type = self.get_type(node_id)
424 # Rule: Image type displays empty
425 if "Image" in node_type and "ImageCrop" not in node_type:
426 return NodeDisplayValue(
427 reason=f"{node_id} is of collectra.Image type: display blank",
428 locked=True,
429 )
431 # Rules for ImageCrop
432 if "ImageCrop" in node_type:
433 crop_data = self.get_crop_region(node_id)
435 # Rule 1: Container crop (has ImageCrop children)
436 if self.children_of_type(node_id, "ImageCrop"):
437 return NodeDisplayValue(
438 crop_region=crop_data,
439 reason=f"{node_id} is a container crop: display blank",
440 locked=True,
441 )
443 # Rule 2: Leaf crop with Text child
444 text_children = self.children_of_type(node_id, "Text")
445 text_children = text_children[0] if text_children else ""
446 if text_children:
447 deepest_id = self.find_deepest(text_children, "Text")
448 if deepest_id == "":
449 raise ValueError(
450 f"Expected to find deepest Text child for node {node_id}, but none found."
451 )
452 deepest_data = self.get_data(deepest_id) if deepest_id else ""
453 return NodeDisplayValue(
454 value=deepest_data,
455 source_id=deepest_id,
456 crop_region=crop_data,
457 reason=f"Leaf crop {node_id} with Text child: deepest Text is {deepest_id}",
458 )
460 # Rule 3: Leaf crop without Text child
461 return NodeDisplayValue(
462 value="",
463 source_id="",
464 crop_region=crop_data,
465 reason=f"Leaf crop {node_id} without Text child",
466 )
468 # Text elements (for completeness)
469 if "Text" in node_type:
470 return NodeDisplayValue(
471 value=self.get_data(node_id),
472 source_id=node_id,
473 crop_region=None,
474 reason="Text element: display data",
475 )
477 return NodeDisplayValue(
478 value="",
479 source_id="",
480 crop_region=None,
481 reason="Unknown type",
482 )