icassp2027-experiment/prepare_ego4d_av_eot_int_review.py

#!/usr/bin/env python3
"""Build Ego4D AV wearer-entry/INT candidates for LLM and human review.

The input is one or more official ``av_train.json``/``av_val.json`` files.
Official voice boundaries are never moved or merged.  This deterministic stage
only attaches nearby speaker activity, transcriptions, and social ``is_at_me``
facts.  Floor-entry and interruption-yield semantics remain unresolved until
LLM pre-labeling and human review.
"""

from __future__ import annotations

import argparse
from collections import Counter
from dataclasses import asdict, dataclass
import json
from pathlib import Path
from typing import Any, Iterable, Sequence


SCHEMA_VERSION = 1
PROTOCOL_VERSION = "ego4d-av-wearer-entry-int-candidates-v1"


@dataclass(frozen=True)
class CandidateConfig:
    context_seconds: float = 4.0
    relation_window_seconds: float = 3.0
    same_speaker_pause_max_seconds: float = 2.0
    other_handover_max_seconds: float = 1.5
    max_terminal_overrun_seconds: float = 1.1


@dataclass(frozen=True)
class SpeechSegment:
    uid: str
    person_id: str
    role: str
    start: float
    end: float
    transcription: str
    addressed_to_wearer: bool | None
    source: dict[str, Any]


def _number(value: Any, *, field: str) -> float:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"{field} must be numeric, got {value!r}")
    result = float(value)
    if result < 0:
        raise ValueError(f"{field} must be non-negative, got {result}")
    return result


def _segment_times(row: dict[str, Any]) -> tuple[float, float]:
    if row.get("start_time") is not None and row.get("end_time") is not None:
        start = _number(row["start_time"], field="start_time")
        end = _number(row["end_time"], field="end_time")
    elif row.get("start_time_sec") is not None and row.get("end_time_sec") is not None:
        start = _number(row["start_time_sec"], field="start_time_sec")
        end = _number(row["end_time_sec"], field="end_time_sec")
    else:
        raise ValueError("voice segment lacks clip-local start/end time")
    if end < start:
        raise ValueError(f"voice segment has negative duration: {start}..{end}")
    return start, end


def _person_key(value: Any) -> str:
    if value is None:
        return ""
    return str(value).strip()


def _best_transcription(
    rows: Sequence[dict[str, Any]], person_id: str, start: float, end: float
) -> str:
    matches: list[tuple[float, float, str]] = []
    for row in rows:
        if _person_key(row.get("person_id", row.get("person"))) != person_id:
            continue
        try:
            left = _number(row.get("start_time_sec", row.get("start_time")), field="transcript start")
            right = _number(row.get("end_time_sec", row.get("end_time")), field="transcript end")
        except ValueError:
            continue
        overlap = max(0.0, min(end, right) - max(start, left))
        if overlap > 0:
            matches.append((left, overlap, str(row.get("transcription", "")).strip()))
    matches.sort()
    return " ".join(text for _, _, text in matches if text)


def _addressed_to_wearer(
    rows: Sequence[dict[str, Any]], person_id: str, start: float, end: float
) -> bool | None:
    overlaps: list[bool] = []
    for row in rows:
        if _person_key(row.get("person", row.get("person_id"))) != person_id:
            continue
        try:
            left, right = _segment_times(row)
        except ValueError:
            continue
        if min(end, right) - max(start, left) > 0:
            value = row.get("is_at_me")
            if isinstance(value, bool):
                overlaps.append(value)
    if not overlaps:
        return None
    return any(overlaps)


def load_clip_segments(
    clip: dict[str, Any], *, audit_out: dict[str, Any] | None = None,
) -> tuple[list[SpeechSegment], str]:
    """Return validated clip-local speaker segments and the wearer person ID."""
    clip_uid = str(clip.get("clip_uid", "")).strip()
    if not clip_uid:
        raise ValueError("clip lacks clip_uid")
    transcripts = clip.get("transcriptions") or []
    social = clip.get("social_segments_talking") or []
    if not isinstance(transcripts, list) or not isinstance(social, list):
        raise ValueError(f"{clip_uid}: transcriptions/social_segments_talking must be lists")

    people: list[dict[str, Any]] = []
    wearer = clip.get("camera_wearer")
    if isinstance(wearer, dict):
        people.append(wearer)
    raw_people = clip.get("persons") or []
    if not isinstance(raw_people, list):
        raise ValueError(f"{clip_uid}: persons must be a list")
    people.extend(person for person in raw_people if isinstance(person, dict))

    wearer_ids = {
        _person_key(person.get("person_id"))
        for person in people
        if person.get("camera_wearer") is True
    }
    wearer_ids.discard("")
    if len(wearer_ids) != 1:
        raise ValueError(f"{clip_uid}: expected exactly one camera wearer, found {sorted(wearer_ids)}")
    wearer_id = next(iter(wearer_ids))

    segments: list[SpeechSegment] = []
    seen: set[tuple[str, float, float]] = set()
    zero_duration: list[dict[str, Any]] = []
    ordinal = 0
    for person in people:
        person_id = _person_key(person.get("person_id"))
        if not person_id:
            raise ValueError(f"{clip_uid}: person lacks person_id")
        voice = person.get("voice_segments") or []
        if not isinstance(voice, list):
            raise ValueError(f"{clip_uid}/{person_id}: voice_segments must be a list")
        for source in voice:
            if not isinstance(source, dict):
                raise ValueError(f"{clip_uid}/{person_id}: voice segment must be an object")
            start, end = _segment_times(source)
            key = (person_id, start, end)
            if key in seen:
                continue
            seen.add(key)
            if end == start:
                zero_duration.append({
                    "person_id": person_id,
                    "participant_role": "wearer" if person_id == wearer_id else "other",
                    "time_seconds": start,
                    "source": dict(source),
                    "policy": "ignored-zero-duration-artifact-not-expanded",
                })
                continue
            segments.append(
                SpeechSegment(
                    uid=f"{clip_uid}/speech/{ordinal:05d}",
                    person_id=person_id,
                    role="wearer" if person_id == wearer_id else "other",
                    start=start,
                    end=end,
                    transcription=_best_transcription(transcripts, person_id, start, end),
                    addressed_to_wearer=(
                        None if person_id == wearer_id
                        else _addressed_to_wearer(social, person_id, start, end)
                    ),
                    source=dict(source),
                )
            )
            ordinal += 1
    segments.sort(key=lambda row: (row.start, row.end, row.person_id, row.uid))
    for left, right in zip(segments, segments[1:]):
        if left.person_id == right.person_id and left.start == right.start and left.end == right.end:
            raise ValueError(f"{clip_uid}: duplicate voice segment after normalization")
    if audit_out is not None:
        audit_out.update({
            "zero_duration_voice_segment_count": len(zero_duration),
            "zero_duration_voice_segments": zero_duration,
            "zero_duration_policy": "ignore-with-audit; never clamp or expand",
        })
    return segments, wearer_id


def _timeline_row(row: SpeechSegment) -> dict[str, Any]:
    return {
        "uid": row.uid,
        "participant_id": row.person_id,
        "participant_role": row.role,
        "start_seconds": row.start,
        "end_seconds": row.end,
        "duration_seconds": row.end - row.start,
        "transcription": row.transcription,
        "addressed_to_wearer": row.addressed_to_wearer,
    }


def _context(segments: Sequence[SpeechSegment], anchor: float, seconds: float) -> list[dict[str, Any]]:
    return [
        _timeline_row(row) for row in segments
        if row.end > anchor - seconds and row.start < anchor + seconds
    ]


def _relation(segments: Sequence[SpeechSegment], anchor: float, wearer_id: str, window: float) -> dict[str, Any]:
    others = [row for row in segments if row.person_id != wearer_id]
    previous = [row for row in others if row.end <= anchor and anchor - row.end <= window]
    following = [row for row in others if row.start >= anchor and row.start - anchor <= window]
    active = [row for row in others if row.start < anchor < row.end]
    prev = max(previous, key=lambda row: row.end, default=None)
    nxt = min(following, key=lambda row: row.start, default=None)
    return {
        "previous_speaker_id": None if prev is None else prev.person_id,
        "previous_speaker_offset_seconds": None if prev is None else prev.end,
        "gap_from_previous_speaker_offset_seconds": None if prev is None else anchor - prev.end,
        "active_other_ids": sorted({row.person_id for row in active}),
        "next_speaker_id": None if nxt is None else nxt.person_id,
        "next_other_onset_seconds": None if nxt is None else nxt.start,
        "next_other_delay_seconds": None if nxt is None else nxt.start - anchor,
    }


def _opportunity_options(
    segments: Sequence[SpeechSegment], wearer: SpeechSegment, wearer_id: str, window: float
) -> list[dict[str, Any]]:
    result = []
    for row in segments:
        if row.person_id == wearer_id or row.end > wearer.start:
            continue
        delay = wearer.start - row.end
        if delay > window:
            continue
        intervening = [
            other for other in segments
            if other.person_id != wearer_id and other.uid != row.uid
            and other.start >= row.end and other.start < wearer.start
        ]
        result.append({
            "uid": row.uid,
            "participant_id": row.person_id,
            "anchor_seconds": row.end,
            "response_delay_seconds": delay,
            "transcription": row.transcription,
            "addressed_to_wearer": row.addressed_to_wearer,
            "other_onset_between_release_and_onset": bool(intervening),
        })
    return sorted(result, key=lambda row: (-float(row["anchor_seconds"]), str(row["uid"])))


def _base_candidate(
    *, clip: dict[str, Any], split: str, source: str, task: str, anchor_type: str,
    anchor: float, region_start: float, region_end: float, sequence: int,
    segments: Sequence[SpeechSegment], evidence: dict[str, Any], rule_type: str,
    automatic_label: str = "candidate", automatic_subtype: str | None = None,
) -> dict[str, Any]:
    clip_uid = str(clip["clip_uid"])
    uid = f"{clip_uid}/{source.lower()}/{sequence:05d}"
    return {
        "schema_version": SCHEMA_VERSION,
        "protocol_version": PROTOCOL_VERSION,
        "dataset": "ego4d_av",
        "uid": uid,
        "event_id": uid,
        "session": clip_uid,
        "clip_uid": clip_uid,
        "video_uid": clip.get("video_uid"),
        "split": split,
        "task": task,
        "event_type": rule_type,
        "candidate_source": source,
        "rule_candidate_type": rule_type,
        "automatic_protocol_label": automatic_label,
        "automatic_hard_negative_type": automatic_subtype,
        "automatic_confidence": "deterministic-boundary-only",
        "anchor_type": anchor_type,
        "anchor_dense_seconds": anchor,
        "event_time": anchor,
        "region_start_seconds": region_start,
        "region_end_seconds": region_end,
        "clip_stem": clip_uid,
        "temporal_evidence": evidence,
        "speaker_timeline": _context(segments, anchor, max(anchor - region_start, region_end - anchor)),
        "context": _context(segments, anchor, max(anchor - region_start, region_end - anchor)),
        "ai_verdict": None,
        "ai_event_type": None,
        "ai_subtype": None,
        "ai_confidence": None,
        "ai_reason": None,
        "human_status": "UNREVIEWED",
        "human_verdict": None,
        "human_event_type": None,
        "human_subtype": None,
        "exclude_reason": None,
    }


def derive_clip_candidates(
    clip: dict[str, Any], *, split: str, config: CandidateConfig,
    audit_out: dict[str, Any] | None = None,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    clip_audit: dict[str, Any] = {}
    segments, wearer_id = load_clip_segments(clip, audit_out=clip_audit)
    declared_clip_end = _number(
        clip.get("clip_end_sec", max((row.end for row in segments), default=0.0)),
        field="clip_end_sec",
    )
    maximum_segment_end = max((row.end for row in segments), default=declared_clip_end)
    terminal_overrun = max(0.0, maximum_segment_end - declared_clip_end)
    if terminal_overrun > config.max_terminal_overrun_seconds + 1e-9:
        first = next(row for row in segments if row.end > declared_clip_end)
        raise ValueError(
            f"{clip.get('clip_uid')}: voice segment {first.uid} ends at {first.end:.6f}s "
            f"after clip_end_sec={declared_clip_end:.6f}; terminal overrun "
            f"{terminal_overrun:.6f}s exceeds the audited limit "
            f"{config.max_terminal_overrun_seconds:.6f}s; refusing to clamp or drop it"
        )
    # Ego4D v1 has quantized terminal voice endpoints a few milliseconds past
    # the 30-fps video endpoint (for example 300.008 vs 299.966667 seconds).
    # Preserve that official endpoint and make the bounded exception explicit.
    clip_end = max(declared_clip_end, maximum_segment_end)
    candidates: list[dict[str, Any]] = []
    wearer_segments = [row for row in segments if row.person_id == wearer_id]

    for sequence, row in enumerate(wearer_segments):
        previous_wearer = max(
            (other for other in wearer_segments if other.end <= row.start),
            key=lambda other: other.end,
            default=None,
        )
        evidence = _relation(segments, row.start, wearer_id, config.relation_window_seconds)
        options = _opportunity_options(segments, row, wearer_id, config.relation_window_seconds)
        evidence.update({
            "behavior_anchor_seconds": row.start,
            "wearer_segment_uid": row.uid,
            "wearer_transcription": row.transcription,
            "wearer_speech_duration_after_onset_seconds": row.end - row.start,
            "previous_wearer_offset_seconds": None if previous_wearer is None else previous_wearer.end,
            "previous_wearer_silence_duration_seconds": None if previous_wearer is None else row.start - previous_wearer.end,
            "opportunity_options": options,
            "preceding_release_options_present": bool(options),
            "official_voice_segments_are_unmodified": True,
        })
        candidates.append(_base_candidate(
            clip=clip, split=split, source="WEARER_ONSET", task="wearer_entry",
            anchor_type="wearer_onset", anchor=row.start,
            region_start=max(0.0, row.start - config.context_seconds),
            region_end=min(clip_end, row.start + config.context_seconds), sequence=sequence,
            segments=segments, evidence=evidence, rule_type="WEARER_ENTRY_CANDIDATE",
        ))

        active_before = [other for other in segments if other.person_id != wearer_id and other.start < row.end < other.end]
        new_during = [other for other in segments if other.person_id != wearer_id and row.start < other.start < row.end]
        next_other = min(
            (other for other in segments if other.person_id != wearer_id and other.start >= row.end),
            key=lambda other: other.start,
            default=None,
        )
        offset_evidence = _relation(segments, row.end, wearer_id, config.relation_window_seconds)
        offset_evidence.update({
            "behavior_anchor_seconds": row.end,
            "wearer_activity_run_start_seconds": row.start,
            "wearer_segment_uid": row.uid,
            "wearer_transcription": row.transcription,
            "other_crosses_wearer_offset": bool(active_before),
            "crossing_other_ids": sorted({other.person_id for other in active_before}),
            "new_other_onsets_during_wearer": [_timeline_row(other) for other in new_during],
            "new_other_onset_during_wearer": bool(new_during),
            "next_speaker_ids": [] if next_other is None else [next_other.person_id],
            "next_speech_onset_seconds": None if next_other is None else next_other.start,
            "official_voice_segments_are_unmodified": True,
        })
        candidates.append(_base_candidate(
            clip=clip, split=split, source="WEARER_OFFSET", task="int_yield",
            anchor_type="wearer_offset", anchor=row.end,
            region_start=max(0.0, row.end - config.context_seconds),
            region_end=min(clip_end, row.end + config.context_seconds), sequence=sequence,
            segments=segments, evidence=offset_evidence, rule_type="YIELD_CANDIDATE",
        ))

    others = [row for row in segments if row.person_id != wearer_id]
    control_sequence = 0
    for previous in others:
        following = sorted(
            (row for row in others if row.start >= previous.end and row.uid != previous.uid),
            key=lambda row: (row.start, row.end),
        )
        if not following:
            continue
        nxt = following[0]
        gap = nxt.start - previous.end
        wearer_in_gap = any(row.start < nxt.start and row.end > previous.end for row in wearer_segments)
        subtype = None
        source = None
        max_gap = None
        if nxt.person_id == previous.person_id:
            source, subtype = "SAME_SPEAKER_PAUSE", "ENTRY_OTHER_MIDTURN_PAUSE"
            max_gap = config.same_speaker_pause_max_seconds
        elif nxt.person_id != previous.person_id:
            source, subtype = "OTHER_TO_OTHER", "ENTRY_OTHER_TO_OTHER"
            max_gap = config.other_handover_max_seconds
        if gap > float(max_gap) or wearer_in_gap:
            continue
        evidence = _relation(segments, previous.end, wearer_id, config.relation_window_seconds)
        evidence.update({
            "previous_speaker_id": previous.person_id,
            "incoming_speaker_id": nxt.person_id,
            "previous_segment_uid": previous.uid,
            "incoming_segment_uid": nxt.uid,
            "gap_seconds": gap,
            "previous_transcription": previous.transcription,
            "incoming_transcription": nxt.transcription,
            "wearer_activity_in_gap": False,
            "official_voice_segments_are_unmodified": True,
        })
        candidates.append(_base_candidate(
            clip=clip, split=split, source=str(source), task="wearer_entry",
            anchor_type="other_offset", anchor=previous.end,
            region_start=max(0.0, previous.end - config.context_seconds),
            region_end=min(clip_end, previous.end + config.context_seconds),
            sequence=control_sequence, segments=segments, evidence=evidence,
            rule_type=str(subtype), automatic_label="hard_negative",
            automatic_subtype=str(subtype),
        ))
        control_sequence += 1
    candidates.sort(key=lambda row: (float(row["event_time"]), str(row["uid"])))
    for row in candidates:
        row["temporal_evidence"].update({
            "declared_clip_end_seconds": declared_clip_end,
            "maximum_voice_segment_end_seconds": maximum_segment_end,
            "terminal_voice_overrun_seconds": terminal_overrun,
            "max_allowed_terminal_overrun_seconds": config.max_terminal_overrun_seconds,
            "terminal_voice_endpoint_preserved_not_clamped": True,
            "zero_duration_voice_segment_count": clip_audit[
                "zero_duration_voice_segment_count"
            ],
            "zero_duration_voice_segments": clip_audit["zero_duration_voice_segments"],
            "zero_duration_voice_policy": clip_audit["zero_duration_policy"],
        })
    clip_audit.update({
        "declared_clip_end_seconds": declared_clip_end,
        "maximum_voice_segment_end_seconds": maximum_segment_end,
        "terminal_voice_overrun_seconds": terminal_overrun,
        "max_allowed_terminal_overrun_seconds": config.max_terminal_overrun_seconds,
    })
    if audit_out is not None:
        audit_out.update(clip_audit)
    return candidates, [_timeline_row(row) for row in segments]


def iter_clips(paths: Iterable[Path]) -> Iterable[tuple[dict[str, Any], str]]:
    seen: set[str] = set()
    for path in paths:
        payload = json.loads(path.read_text(encoding="utf-8"))
        videos = payload.get("videos") if isinstance(payload, dict) else None
        if not isinstance(videos, list):
            raise ValueError(f"{path}: expected a videos list")
        for video in videos:
            if not isinstance(video, dict):
                continue
            split = str(video.get("split") or path.stem.removeprefix("av_") or "unknown")
            for clip in video.get("clips") or []:
                if not isinstance(clip, dict) or clip.get("valid") is False:
                    continue
                uid = str(clip.get("clip_uid", ""))
                if uid in seen:
                    raise ValueError(f"duplicate clip_uid across inputs: {uid}")
                seen.add(uid)
                yield clip, split


def _write_jsonl(path: Path, rows: Sequence[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".partial")
    with temporary.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    temporary.replace(path)


def _write_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".partial")
    temporary.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    temporary.replace(path)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--annotations", type=Path, nargs="+", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--timeline-output", type=Path)
    parser.add_argument("--clip-uid", action="append", help="Restrict to one or more clip UIDs")
    parser.add_argument("--limit", type=int, help="Process at most N selected clips")
    parser.add_argument("--context-seconds", type=float, default=4.0)
    parser.add_argument("--relation-window-seconds", type=float, default=3.0)
    parser.add_argument("--same-speaker-pause-max-seconds", type=float, default=2.0)
    parser.add_argument("--other-handover-max-seconds", type=float, default=1.5)
    parser.add_argument(
        "--max-terminal-overrun-seconds", type=float, default=1.1,
        help="Allow and audit a bounded official voice endpoint past clip_end_sec (default: %(default)s).",
    )
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    numeric = [
        args.context_seconds, args.relation_window_seconds,
        args.same_speaker_pause_max_seconds, args.other_handover_max_seconds,
        args.max_terminal_overrun_seconds,
    ]
    if any(value < 0 for value in numeric):
        raise ValueError("all time windows must be non-negative")
    if args.limit is not None and args.limit <= 0:
        raise ValueError("--limit must be positive")
    paths = [path.expanduser().resolve() for path in args.annotations]
    for path in paths:
        if not path.is_file():
            raise FileNotFoundError(path)
    selected = list(iter_clips(paths))
    if args.clip_uid:
        requested = set(args.clip_uid)
        known = {str(clip["clip_uid"]) for clip, _ in selected}
        if requested - known:
            raise ValueError(f"unknown clip UID(s): {sorted(requested - known)}")
        selected = [(clip, split) for clip, split in selected if clip["clip_uid"] in requested]
    if args.limit is not None:
        selected = selected[: args.limit]
    config = CandidateConfig(*numeric)
    rows: list[dict[str, Any]] = []
    timelines: dict[str, list[dict[str, Any]]] = {}
    clip_audits: dict[str, dict[str, Any]] = {}
    for clip, split in selected:
        audit: dict[str, Any] = {}
        derived, timeline = derive_clip_candidates(
            clip, split=split, config=config, audit_out=audit
        )
        rows.extend(derived)
        clip_uid = str(clip["clip_uid"])
        timelines[clip_uid] = timeline
        clip_audits[clip_uid] = audit
    if not rows:
        raise ValueError("selected annotations produced no candidates")
    missing_voice_count_by_clip = {
        str(clip["clip_uid"]): len(clip.get("missing_voice_segments") or [])
        for clip, _ in selected
    }
    for row in rows:
        row["temporal_evidence"]["clip_missing_voice_segments_count"] = (
            missing_voice_count_by_clip[row["clip_uid"]]
        )
    output = args.output.expanduser().resolve()
    timeline_output = (
        args.timeline_output.expanduser().resolve() if args.timeline_output
        else output.with_name(output.stem + "_timelines.json")
    )
    existing = [path for path in (output, timeline_output) if path.exists()]
    if existing and not args.overwrite:
        raise FileExistsError(
            "Refusing to overwrite existing output(s): "
            + ", ".join(str(path) for path in existing)
            + "; pass --overwrite"
        )
    _write_jsonl(output, rows)
    _write_json(
        timeline_output,
        {
            "schema_version": 1,
            "dataset": "ego4d_av",
            "sessions": timelines,
            "clip_audits": clip_audits,
        },
    )
    counts = Counter(row["candidate_source"] for row in rows)
    print(f"clips={len(selected)} candidates={len(rows)} output={output}")
    print("candidate_sources=" + json.dumps(dict(sorted(counts.items())), ensure_ascii=False))
    print(
        "zero_duration_voice_segments_ignored_with_audit="
        + str(sum(audit["zero_duration_voice_segment_count"] for audit in clip_audits.values()))
    )
    print(f"timelines={timeline_output}")


if __name__ == "__main__":
    main()