app.py

"""
Video Annotation QC Agent
Tapi Tag Ontology · Phase II
Hugging Face Space — Gradio 4.44.0 compatible
"""

# ── Python 3.13 audioop fix (must be first) ──────────────────────────────────
import sys
try:
    import audioop  # noqa: F401
except ModuleNotFoundError:
    try:
        import audioop_lts as audioop  # type: ignore
        sys.modules["audioop"] = audioop
    except ModuleNotFoundError:
        pass

# ── Standard imports ─────────────────────────────────────────────────────────
import json
import random
from pathlib import Path

import gradio as gr

# ── Core QC imports ──────────────────────────────────────────────────────────
sys.path.insert(0, str(Path(__file__).parent))
from core.qc_agent import QCAgent, FrameAnnotation
from core.tapi_ontology import (
    ALL_VALID_LABELS,
    EGO_MOVEMENT_LABELS, VEHICLE_MOVEMENT_LABELS,
    PEDESTRIAN_MOVEMENT_LABELS, ANIMAL_MOVEMENT_LABELS,
    WEATHER_LABELS, TIME_OF_DAY_LABELS, VISIBILITY_LABELS,
    SETTING_AREA_LABELS, SETTING_INFRASTRUCTURE_LABELS,
    ROAD_PROPERTY_LABELS, EVENT_LABELS,
    TRAFFIC_SIGNAL_LABELS, TRAFFIC_SIGN_LABELS,
    ACTOR_VEHICLE_LABELS,
)

# ─────────────────────────────────────────────────────────────────────────────
#  Demo annotation generator — uses REAL Tapi labels
# ─────────────────────────────────────────────────────────────────────────────

# Realistic label pairs for demo
_DEMO_EGO    = ["ego drives on ego lane", "ego drives behind vehicle", "ego waits at signals"]
_DEMO_VEH    = ["vehicle drives on ego lane", "vehicle drives on non ego lane", "vehicle brakes in ego lane"]
_DEMO_ACTOR  = ["sedan", "suv", "truck", "bus"]
_DEMO_INFRA  = ["buildings", "gas station"]

def _make_demo_annotations(n_frames: int, n_objects: int) -> list[dict]:
    """Generate annotations using real Tapi labels with deliberate SOP violations."""
    anns = []
    track_labels = {f"T{i+1:04d}": random.choice(_DEMO_ACTOR) for i in range(n_objects)}

    for f in range(n_frames):
        objects = []

        # --- Environment labels (realistic) ---
        objects.append({"label": random.choice(list(WEATHER_LABELS)), "bbox": [0, 0, 10, 10]})
        objects.append({"label": random.choice(list(TIME_OF_DAY_LABELS)), "bbox": [0, 0, 10, 10]})

        # --- Ego movement ---
        ego_label = random.choice(_DEMO_EGO)
        # INJECT EGO-03 error at frame 12: ego waits + ego drives at same time
        if f == 12:
            objects.append({"label": "ego waits at signals",      "bbox": [0, 0, 10, 10]})
            objects.append({"label": "ego drives on ego lane",    "bbox": [0, 0, 10, 10]})
        else:
            objects.append({"label": ego_label, "bbox": [0, 0, 10, 10]})

        # --- Vehicle actors with bboxes ---
        for track_id, actor_label in track_labels.items():
            label = actor_label

            # INJECT LBL-01 at frame 5: unknown label
            if f == 5 and track_id == "T0001":
                label = "hovercar_xyz"

            # INJECT TRK-01 at frame 18: label changed on same track
            if f >= 18 and track_id == "T0002":
                label = "bus"

            # INJECT BOX-01 at frame 10: box too large
            if f == 10 and track_id == "T0001":
                bbox = [0, 0, 1900, 1070]
            # INJECT BOX-02 at frame 20: out of bounds
            elif f == 20 and track_id == list(track_labels.keys())[-1]:
                bbox = [1850, 950, 300, 300]
            else:
                x = random.randint(50, 800)
                y = random.randint(50, 500)
                w = random.randint(80, 300)
                h = random.randint(60, 200)
                bbox = [x, y, w, h]

            # INJECT FLICKER at frame 8: skip T0001 so it disappears then reappears
            if f == 8 and track_id == "T0001":
                continue

            objects.append({
                "track_id": track_id,
                "label": label,
                "bbox": bbox,
                "confidence": round(random.uniform(0.7, 0.99), 2),
            })

        # INJECT TRK-02 duplicate at frame 15
        if f == 15 and objects:
            actor_objs = [o for o in objects if o.get("track_id")]
            if actor_objs:
                dup = dict(actor_objs[0])
                dup["track_id"] = "T9999"
                objects.append(dup)

        # INJECT SET-03 at frame 3: urban area without buildings
        if f == 3:
            objects.append({"label": "urban", "bbox": [0, 0, 10, 10]})
            # deliberately NOT adding "buildings"

        anns.append({"frame_idx": f, "objects": objects})

    return anns


def _make_demo_ai_detections(n_frames: int) -> dict:
    idx = {}
    for f in range(n_frames):
        # Frame 7: AI found a pedestrian the human missed
        if f == 7:
            idx[f] = [{"bbox": [300, 200, 60, 130], "label": "adult", "confidence": 0.88}]
        else:
            idx[f] = []
    return idx


# ─────────────────────────────────────────────────────────────────────────────
#  QC runners
# ─────────────────────────────────────────────────────────────────────────────

def _build_ctx(flicker_gap: int) -> dict:
    return {
        "frame_width":       1920,
        "frame_height":      1080,
        "fps":               25,
        "flicker_gap":       int(flicker_gap),
        "track_history":     {},
        "track_frame_count": {},
        "track_last_seen":   {},
        "detection_index":   {},
        "ego_miss_streak":   0,
        "event_start":       {},
    }


def _format_report(report) -> tuple[str, str]:
    """Return (markdown_report, json_string)."""
    health = (
        "🟢 Good" if report.total_issues == 0
        else "🟡 Needs Review" if report.total_issues < 5
        else "🔴 Action Required"
    )

    md = f"""## 📋 QC Report — {report.total_frames} frames analysed

| Metric | Value |
|--------|-------|
| Total Issues | **{report.total_issues}** |
| Frames Affected | **{len(report.frames_affected)}** |
| Health | **{health}** |

### Summary
{report.summary}
"""
    if report.issues_by_type:
        rows = "\n".join(
            f"| `{k}` | {v} | {'🔴' if v >= 3 else '🟡'} |"
            for k, v in sorted(report.issues_by_type.items(), key=lambda x: -x[1])
        )
        md += f"\n### Issues by Rule\n| Rule | Count | |\n|------|-------|---|\n{rows}\n"
    else:
        md += "\n### ✅ No issues found!\n"

    if report.issues:
        lines = ["\n### Detailed Issues"]
        prev = -1
        for iss in sorted(report.issues, key=lambda i: i.frame):
            if iss.frame != prev:
                lines.append(f"\n**Frame {iss.frame}**")
                prev = iss.frame
            lines.append("  " + iss.user_message())
        md += "\n".join(lines)

    frames_str = ", ".join(str(f) for f in report.frames_affected[:25])
    if len(report.frames_affected) > 25:
        frames_str += f" … (+{len(report.frames_affected)-25} more)"
    md += f"\n\n### Frames Affected\n`{frames_str or 'None'}`"

    raw = {
        "summary": report.summary,
        "total_frames": report.total_frames,
        "total_issues": report.total_issues,
        "frames_affected": report.frames_affected,
        "issues_by_type": report.issues_by_type,
        "issues": [
            {"rule": i.check, "severity": i.severity,
             "frame": i.frame, "track_id": i.track_id, "message": i.message}
            for i in report.issues
        ],
    }
    return md, json.dumps(raw, indent=2)


def run_qc_demo(n_frames: int, n_objects: int, flicker_gap: int):
    ctx   = _build_ctx(flicker_gap)
    agent = QCAgent(ctx=ctx)
    agent.load_ai_detections(_make_demo_ai_detections(int(n_frames)))
    for ann in _make_demo_annotations(int(n_frames), int(n_objects)):
        agent.process_frame(FrameAnnotation(ann["frame_idx"], ann["objects"]))
    return _format_report(agent.generate_report())


def run_qc_on_json(annotation_json: str, flicker_gap: int):
    try:
        human_anns = json.loads(annotation_json)
    except json.JSONDecodeError as e:
        return f"❌ Invalid JSON: {e}", ""
    if not isinstance(human_anns, list):
        return "❌ JSON must be a list of frame annotation objects.", ""
    ctx   = _build_ctx(flicker_gap)
    agent = QCAgent(ctx=ctx)
    for ann in human_anns:
        if "frame_idx" not in ann:
            return "❌ Each object must have a `frame_idx` field.", ""
        agent.process_frame(FrameAnnotation(ann["frame_idx"], ann.get("objects", [])))
    return _format_report(agent.generate_report())


# ─────────────────────────────────────────────────────────────────────────────
#  Example JSON using REAL Tapi labels
# ─────────────────────────────────────────────────────────────────────────────
EXAMPLE_JSON = json.dumps([
    {
        "frame_idx": 0,
        "objects": [
            {"label": "daytime",                    "bbox": [0,0,10,10]},
            {"label": "clear",                      "bbox": [0,0,10,10]},
            {"label": "ego drives on ego lane",     "bbox": [0,0,10,10]},
            {"label": "urban",                      "bbox": [0,0,10,10]},
            {"label": "buildings",                  "bbox": [0,0,10,10]},
            {"track_id":"T001","label":"sedan",     "bbox":[100,200,150,80]},
            {"track_id":"T002","label":"adult",     "bbox":[400,300,60,120]},
        ]
    },
    {
        "frame_idx": 1,
        "objects": [
            {"label": "daytime",                    "bbox": [0,0,10,10]},
            {"label": "clear",                      "bbox": [0,0,10,10]},
            {"label": "ego drives on ego lane",     "bbox": [0,0,10,10]},
            {"label": "urban",                      "bbox": [0,0,10,10]},
            {"label": "buildings",                  "bbox": [0,0,10,10]},
            # TRK-01: T001 was sedan, now truck → label change error
            {"track_id":"T001","label":"truck",     "bbox":[110,205,150,80]},
            {"track_id":"T002","label":"adult",     "bbox":[405,302,60,120]},
        ]
    },
    {
        "frame_idx": 2,
        "objects": [
            {"label": "daytime",                    "bbox": [0,0,10,10]},
            {"label": "clear",                      "bbox": [0,0,10,10]},
            # EGO-03: waits + drives at same time
            {"label": "ego waits at signals",       "bbox": [0,0,10,10]},
            {"label": "ego drives on ego lane",     "bbox": [0,0,10,10]},
            {"label": "urban",                      "bbox": [0,0,10,10]},
            {"label": "buildings",                  "bbox": [0,0,10,10]},
            # BOX-01: box covers >90% frame
            {"track_id":"T001","label":"truck",     "bbox":[0,0,1900,1070]},
        ]
    },
    {
        "frame_idx": 3,
        "objects": [
            {"label": "daytime",                    "bbox": [0,0,10,10]},
            {"label": "clear",                      "bbox": [0,0,10,10]},
            {"label": "ego drives on ego lane",     "bbox": [0,0,10,10]},
            # SET-03: urban without buildings
            {"label": "urban",                      "bbox": [0,0,10,10]},
            # LBL-01: not in ontology
            {"track_id":"T001","label":"hovercar",  "bbox":[200,200,100,80]},
        ]
    },
], indent=2)


# ─────────────────────────────────────────────────────────────────────────────
#  Gradio UI
# ─────────────────────────────────────────────────────────────────────────────

CSS = """
.header {
    background: linear-gradient(135deg,#1e1b4b,#312e81,#1e3a5f);
    border-radius:12px; padding:24px 32px; margin-bottom:20px;
    border:1px solid #4338ca;
}
.header h1{color:#e0e7ff;font-size:1.7em;margin:0}
.header p {color:#a5b4fc;margin:6px 0 0;font-size:.9em}
.rule-chip {
    display:inline-block; background:#1e293b; border:1px solid #334155;
    border-radius:6px; padding:6px 12px; margin:3px;
    font-size:.8em; color:#94a3b8; font-family:monospace;
}
.rule-chip b{color:#818cf8}
"""

with gr.Blocks(title="Video Annotation QC Agent · Tapi Phase II", css=CSS) as demo:

    gr.HTML("""
    <div class="header">
      <h1>🎯 Video Annotation QC Agent</h1>
      <p>Tapi Tag Ontology · Phase II &nbsp;|&nbsp; SOP-Driven Quality Control &nbsp;|&nbsp; Human-in-the-Loop</p>
    </div>
    """)

    with gr.Tabs():

        # ── Tab 1: Demo ──────────────────────────────────────────────────────
        with gr.Tab("🚀 Demo — Try it now"):
            gr.Markdown(
                "Run QC on **auto-generated Tapi annotations** with deliberately injected SOP violations. "
                "All labels are from the real Phase II ontology."
            )
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### ⚙️ Parameters")
                    sl_frames  = gr.Slider(10, 100, value=30, step=5,  label="Number of frames")
                    sl_objects = gr.Slider(1,  6,   value=3,  step=1,  label="Number of tracked actors")
                    sl_flicker = gr.Slider(2,  15,  value=5,  step=1,  label="Flicker detection gap (frames)")
                    btn_demo   = gr.Button("▶  Run QC", variant="primary")
                with gr.Column(scale=2):
                    out_report = gr.Markdown()
                    with gr.Accordion("📥 Download JSON report", open=False):
                        out_json = gr.Code(language="json")

            btn_demo.click(
                fn=run_qc_demo,
                inputs=[sl_frames, sl_objects, sl_flicker],
                outputs=[out_report, out_json],
            )

        # ── Tab 2: Custom ────────────────────────────────────────────────────
        with gr.Tab("📝 Paste Your Annotations"):
            gr.Markdown("""
Paste your annotation JSON. Each frame must list its `objects` with a valid Tapi label.

```
[{"frame_idx": 0, "objects": [{"track_id":"T001","label":"sedan","bbox":[x,y,w,h]}, ...]}]
```
""")
            with gr.Row():
                with gr.Column(scale=1):
                    inp_json   = gr.Code(value=EXAMPLE_JSON, language="json",
                                         label="Annotation JSON", lines=22)
                    sl_flicker2 = gr.Slider(2, 15, value=5, step=1,
                                            label="Flicker detection gap (frames)")
                    btn_custom  = gr.Button("▶  Run QC on My Annotations", variant="primary")
                with gr.Column(scale=2):
                    out_custom  = gr.Markdown()
                    with gr.Accordion("📥 Raw JSON", open=False):
                        out_cjson = gr.Code(language="json")

            btn_custom.click(
                fn=run_qc_on_json,
                inputs=[inp_json, sl_flicker2],
                outputs=[out_custom, out_cjson],
            )

        # ── Tab 3: SOP Rules Reference ───────────────────────────────────────
        with gr.Tab("📋 SOP Rules"):
            gr.Markdown("## All QC Checks — derived from Tapi Phase II SOP (page 12)\n")
            gr.HTML("""
<div class="rule-chip"><b>ENV-01</b> Missing Environment label for any frames</div>
<div class="rule-chip"><b>ENV-02</b> Time of Day cannot change within a task</div>
<div class="rule-chip"><b>ENV-03</b> Contradictory Weather ↔ Visibility annotation</div>
<div class="rule-chip"><b>ENV-04</b> Multiple Visibility labels overlapping</div>
<div class="rule-chip"><b>SET-01</b> Missing Setting label for any frames</div>
<div class="rule-chip"><b>SET-02</b> Multiple Setting Area labels overlapping</div>
<div class="rule-chip"><b>SET-03</b> Urban/Residential area must have 'buildings'</div>
<div class="rule-chip"><b>ROAD-01</b> Missing Road Topology label for any frames</div>
<div class="rule-chip"><b>ROAD-02</b> Road Topology labels overlapping</div>
<div class="rule-chip"><b>EVT-01</b> Event tag longer than 5 seconds</div>
<div class="rule-chip"><b>EVT-02</b> Missing traffic signal during event frames</div>
<div class="rule-chip"><b>EGO-01</b> Ego movement missing for 10+ consecutive frames</div>
<div class="rule-chip"><b>EGO-02</b> Ego parked overlapping with other Ego movement</div>
<div class="rule-chip"><b>EGO-03</b> Ego waits overlapping with other Ego movement</div>
<div class="rule-chip"><b>EGO-04</b> Ego merges into + out of same location simultaneously</div>
<div class="rule-chip"><b>EGO-05</b> Ego turns left + right simultaneously</div>
<div class="rule-chip"><b>EGO-06</b> Ego changes lane left + right simultaneously</div>
<div class="rule-chip"><b>EGO-07</b> Ego merges into Rotary + out of Rotary simultaneously</div>
<div class="rule-chip"><b>MOV-01</b> Same label used for multiple annotations in same timeline</div>
<div class="rule-chip"><b>BOX-01</b> Bounding box too large (&gt;90%) or too small</div>
<div class="rule-chip"><b>BOX-02</b> Bounding box outside frame boundary</div>
<div class="rule-chip"><b>LBL-01</b> Label not in Tapi Tag Ontology Phase II</div>
<div class="rule-chip"><b>TRK-01</b> Track ID label changed across frames</div>
<div class="rule-chip"><b>TRK-02</b> Duplicate bounding boxes in same frame (IoU &gt; 0.70)</div>
<div class="rule-chip"><b>FLICKER</b> Object disappears and reappears within flicker gap</div>
<div class="rule-chip"><b>MISSED</b> AI detected object with no human annotation</div>
""")
            gr.Markdown("""
### Tapi Phase II Valid Labels (selected)

**Environment** — Weather: `clear` `cloudy` `rainy` `snowy` `foggy`
| Time of Day: `daytime` `nighttime` `dawn` `dusk`
| Visibility: `clear` `sun glare` `headlight glare` `poor visibility from weather`

**Setting** — Area: `urban` `residential` `rural` `industrial` `highway`
| Infrastructure: `buildings` `gas station` `sound barrier` `construction site`
| Nature: `trees` `river` `lake`

**Road Topology** — `median-separated` `wet` `snow` `hilly` `hov` `winding` `debris` `pothole`

**Events** — `green becomes yellow` `yellow becomes red` `red becomes green`

**Actors** — `ego car` `sedan` `suv` `other car` `van` `motorcycle` `bicycle` `bicyclist`
`truck` `bus` `police` `firetruck` `ambulance` `dog` `deer` `child` `adult` `wheel chair`

**Ego Movements** — `ego drives on ego lane` `ego brakes in ego lane` `ego waits at signals`
`ego turns left at intersection` `ego merges into rotary` *(+ 30 more)*

**Vehicle Movements** — `vehicle drives on ego lane` `vehicle changes lane right into ego lane`
`vehicle passes left ego` *(+ 40 more)*

**Pedestrian/Animal** — `pedestrian walks on crosswalk` `pedestrian stands on sidewalk`
`animal walks in ego lane` *(+ 15 more)*

**Static Objects** — `protected turn` `unprotected turn` `flashing signal` `walk signal`
`stop sign` `speed sign` `exit sign` `bus stop sign` `children crossing sign` `school zone sign`
""")

        # ── Tab 4: Architecture ──────────────────────────────────────────────
        with gr.Tab("🧩 Architecture"):
            gr.Markdown("""
## System Architecture

```
Video (MP4)
    │
    ▼
Frame Extractor (OpenCV)
    │
    ├──► YOLO Detector ──► IoU Tracker ──► AI Detection Index
    │                                             │
    │                                             ▼
Human Annotations ──────────────────► QC Agent (stateful)
                                             │
                              ┌──────────────┼──────────────┐
                              ▼              ▼              ▼
                       SOP Rule Engine   Flicker       Missed Ann.
                       (ENV/SET/ROAD/    Check         Check
                        EVT/EGO/MOV/
                        BOX/LBL/TRK)
                              │
                              ▼
                         QC Report (JSON + Markdown)
```

### Files
| File | Purpose |
|------|---------|
| `app.py` | Gradio UI — HF Space entry point |
| `core/tapi_ontology.py` | Complete Tapi Phase II label registry |
| `core/sop_rules.py` | 25 SOP rules → executable Python functions |
| `core/qc_agent.py` | Stateful cross-frame QC engine |
| `core/__init__.py` | Package exports |
| `requirements.txt` | Dependencies |
""")

        # ── Tab 5: Setup ─────────────────────────────────────────────────────
        with gr.Tab("⚡ Local Setup"):
            gr.Markdown("""
## Run Locally

```bash
# 1. Clone your HF Space
git clone https://huggingface.co/spaces/YOUR-USERNAME/video-annotation-qc
cd video-annotation-qc

# 2. Install dependencies
pip install -r requirements.txt
pip install gradio==4.44.0

# 3. Run
python app.py
```

## Python API

```python
from core.qc_agent import QCAgent, FrameAnnotation

ctx = {
    "frame_width": 1920, "frame_height": 1080, "fps": 25,
    "flicker_gap": 5, "track_history": {}, "track_frame_count": {},
    "track_last_seen": {}, "detection_index": {}, "ego_miss_streak": 0,
}

agent = QCAgent(ctx=ctx)

# Process each frame
agent.process_frame(FrameAnnotation(
    frame_idx=0,
    objects=[
        {"label": "daytime",             "bbox": [0,0,10,10]},
        {"label": "clear",               "bbox": [0,0,10,10]},
        {"label": "ego drives on ego lane","bbox":[0,0,10,10]},
        {"label": "urban",               "bbox": [0,0,10,10]},
        {"label": "buildings",           "bbox": [0,0,10,10]},
        {"track_id":"T001","label":"sedan","bbox":[100,200,150,80]},
    ]
))

report = agent.generate_report()
print(report.to_display())
```
""")

    gr.HTML("""
    <div style="text-align:center;padding:14px;color:#475569;
                font-size:.8em;margin-top:16px;border-top:1px solid #1e293b">
      Video Annotation QC Agent · Tapi Tag Ontology Phase II · Powered by Gradio
    </div>
    """)

if __name__ == "__main__":
    demo.launch()