chore: extract generic temporal nested logic

Author: rishisurana-labelbox

libs/labelbox/src/labelbox/data/serialization/ndjson/label.py

@@ -28,6 +28,7 @@
 from ...annotation_types.audio import (
     AudioClassificationAnnotation,
 )
+from .temporal import create_audio_ndjson_annotations
 from labelbox.types import DocumentRectangle, DocumentEntity
 from .classification import (
     NDChecklistSubclass,
@@ -169,190 +170,25 @@ def _create_video_annotations(
     def _create_audio_annotations(
         cls, label: Label
     ) -> Generator[BaseModel, None, None]:
-        """Create audio annotations with nested classifications (v3-like),
-        while preserving v2 behavior for non-nested cases.
-
-        Strategy:
-        - Group audio annotations by classification (schema_id or name)
-        - Identify root groups (not fully contained by another group's frames)
-        - For each root group, build answer items grouped by value with frames
-        - Recursively attach nested classifications by time containment
-        """
-
-        # 1) Collect all audio annotations grouped by classification key
-        #    Use feature_schema_id when present, otherwise fall back to name
-        audio_by_group: Dict[str, List[AudioClassificationAnnotation]] = defaultdict(list)
-        for annot in label.annotations:
-            if isinstance(annot, AudioClassificationAnnotation):
-                audio_by_group[annot.feature_schema_id or annot.name].append(annot)
-
-        if not audio_by_group:
+        """Create audio annotations with nested classifications using modular hierarchy builder."""
+        # Extract audio annotations from the label
+        audio_annotations = [
+            annot for annot in label.annotations 
+            if isinstance(annot, AudioClassificationAnnotation)
+        ]
+        
+        if not audio_annotations:
             return
-
-        # Helper: produce a user-facing classification name for a group
-        def group_display_name(group_key: str, anns: List[AudioClassificationAnnotation]) -> str:
-            # Prefer the first non-empty annotation name
-            for a in anns:
-                if a.name:
-                    return a.name
-            # Fallback to group key (may be schema id)
-            return group_key
-
-        # Helper: compute whether group A is fully contained by any other group by time
-        def is_group_nested(group_key: str) -> bool:
-            anns = audio_by_group[group_key]
-            for ann in anns:
-                # An annotation is considered nested if there exists any container in other groups
-                contained = False
-                for other_key, other_anns in audio_by_group.items():
-                    if other_key == group_key:
-                        continue
-                    for parent in other_anns:
-                        if parent.start_frame <= ann.start_frame and (
-                            parent.end_frame is not None
-                            and ann.end_frame is not None
-                            and parent.end_frame >= ann.end_frame
-                        ):
-                            contained = True
-                            break
-                    if contained:
-                        break
-                if not contained:
-                    # If any annotation in this group is not contained, group is a root
-                    return False
-            # All annotations were contained somewhere → nested group
-            return True
-
-        # Helper: group annotations by logical value and produce answer entries
-        def group_by_value(annotations: List[AudioClassificationAnnotation]) -> List[Dict[str, Any]]:
-            value_buckets: Dict[str, List[AudioClassificationAnnotation]] = defaultdict(list)
-
-            for ann in annotations:
-                # Compute grouping key depending on classification type
-                if hasattr(ann.value, "answer"):
-                    if isinstance(ann.value.answer, list):
-                        # Checklist: stable key from selected option names
-                        key = str(sorted([opt.name for opt in ann.value.answer]))
-                    elif hasattr(ann.value.answer, "name"):
-                        # Radio: option name
-                        key = ann.value.answer.name
-                    else:
-                        # Text: the string value
-                        key = ann.value.answer
-                else:
-                    key = str(ann.value)
-                value_buckets[key].append(ann)
-
-            entries: List[Dict[str, Any]] = []
-            for _, anns in value_buckets.items():
-                first = anns[0]
-                frames = [{"start": a.start_frame, "end": a.end_frame} for a in anns]
-
-                if hasattr(first.value, "answer") and isinstance(first.value.answer, list):
-                    # Checklist: emit one entry per distinct option present in this bucket
-                    # Since bucket is keyed by the combination, take names from first
-                    for opt_name in sorted([o.name for o in first.value.answer]):
-                        entries.append({"name": opt_name, "frames": frames})
-                elif hasattr(first.value, "answer") and hasattr(first.value.answer, "name"):
-                    # Radio
-                    entries.append({"name": first.value.answer.name, "frames": frames})
-                else:
-                    # Text
-                    entries.append({"value": first.value.answer, "frames": frames})
-
-            return entries
-
-        # Helper: check if child ann is inside any of the parent frames list
-        def ann_within_frames(ann: AudioClassificationAnnotation, frames: List[Dict[str, int]]) -> bool:
-            for fr in frames:
-                if fr["start"] <= ann.start_frame and (
-                    ann.end_frame is not None and fr["end"] is not None and fr["end"] >= ann.end_frame
-                ):
-                    return True
-            return False
-
-        # Helper: recursively build nested classifications for a specific parent frames list
-        def build_nested_for_frames(parent_frames: List[Dict[str, int]], exclude_group: str) -> List[Dict[str, Any]]:
-            nested: List[Dict[str, Any]] = []
-
-            # Collect all annotations within parent frames across all groups except the excluded one
-            all_contained: List[AudioClassificationAnnotation] = []
-            for gk, ga in audio_by_group.items():
-                if gk == exclude_group:
-                    continue
-                all_contained.extend([a for a in ga if ann_within_frames(a, parent_frames)])
-
-            def strictly_contains(container: AudioClassificationAnnotation, inner: AudioClassificationAnnotation) -> bool:
-                if container is inner:
-                    return False
-                if container.end_frame is None or inner.end_frame is None:
-                    return False
-                return container.start_frame <= inner.start_frame and container.end_frame >= inner.end_frame and (
-                    container.start_frame < inner.start_frame or container.end_frame > inner.end_frame
-                )
-
-            for group_key, anns in audio_by_group.items():
-                if group_key == exclude_group:
-                    continue
-                # Do not nest groups that are roots themselves to avoid duplicating top-level groups inside others
-                if group_key in root_group_keys:
-                    continue
-
-                # Filter annotations that are contained by any parent frame
-                candidate_anns = [a for a in anns if ann_within_frames(a, parent_frames)]
-                if not candidate_anns:
-                    continue
-
-                # Keep only immediate children (those not strictly contained by another contained annotation)
-                child_anns = []
-                for a in candidate_anns:
-                    has_closer_container = any(strictly_contains(b, a) for b in all_contained)
-                    if not has_closer_container:
-                        child_anns.append(a)
-                if not child_anns:
-                    continue
-
-                # Build this child classification block
-                child_entries = group_by_value(child_anns)
-                # Recurse: for each answer entry, compute further nested
-                for entry in child_entries:
-                    entry_frames = entry.get("frames", [])
-                    child_nested = build_nested_for_frames(entry_frames, group_key)
-                    if child_nested:
-                        entry["classifications"] = child_nested
-
-                nested.append({
-                    "name": group_display_name(group_key, anns),
-                    "answer": child_entries,
-                })
-
-            return nested
-
-        # 2) Determine root groups (not fully contained by other groups)
-        root_group_keys = [k for k in audio_by_group.keys() if not is_group_nested(k)]
-
-        # 3) Emit one NDJSON object per root classification group
-        class AudioNDJSON(BaseModel):
-            name: str
-            answer: List[Dict[str, Any]]
-            dataRow: Dict[str, str]
-
-        for group_key in root_group_keys:
-            anns = audio_by_group[group_key]
-            top_entries = group_by_value(anns)
-
-            # Attach nested to each top-level answer entry
-            for entry in top_entries:
-                frames = entry.get("frames", [])
-                children = build_nested_for_frames(frames, group_key)
-                if children:
-                    entry["classifications"] = children
-
-            yield AudioNDJSON(
-                name=group_display_name(group_key, anns),
-                answer=top_entries,
-                dataRow={"globalKey": label.data.global_key},
-            )
+        
+        # Use the modular hierarchy builder to create NDJSON annotations
+        ndjson_annotations = create_audio_ndjson_annotations(
+            audio_annotations, 
+            label.data.global_key
+        )
+        
+        # Yield each NDJSON annotation
+        for annotation in ndjson_annotations:
+            yield annotation