Add files using upload-large-folder tool

runs/tubetoken_phase0/eval_stride8_n64_bidir.log

@@ -0,0 +1,13 @@
+
+# TubeToken Phase 0 Proposal Evaluation
+
+- all: n=3944, R@16=0.469, R@32=0.597, R@64=0.754, Oracle J&F=0.7491, miss=24.62%
+- area_unstable: n=18, R@16=0.556, R@32=0.556, R@64=0.889, Oracle J&F=0.7114, miss=11.11%
+- audio_keyword: n=2844, R@16=0.475, R@32=0.610, R@64=0.766, Oracle J&F=0.7569, miss=23.42%
+- h3_candidate: n=3932, R@16=0.469, R@32=0.597, R@64=0.754, Oracle J&F=0.7488, miss=24.64%
+- partial: n=8, R@16=0.250, R@32=0.250, R@64=1.000, Oracle J&F=0.8123, miss=0.00%
+- same_category: n=330, R@16=0.482, R@32=0.588, R@64=0.709, Oracle J&F=0.7261, miss=29.09%
+- small: n=1631, R@16=0.237, R@32=0.392, R@64=0.633, Oracle J&F=0.6367, miss=36.73%
+- spatial_keyword: n=965, R@16=0.331, R@32=0.476, R@64=0.658, Oracle J&F=0.6714, miss=34.20%
+- test_s: n=2288, R@16=0.326, R@32=0.483, R@64=0.657, Oracle J&F=0.6674, miss=34.27%
+- test_u: n=1656, R@16=0.665, R@32=0.755, R@64=0.887, Oracle J&F=0.8618, miss=11.29%

runs/tubetoken_phase0/miss_videos_r64.txt

@@ -0,0 +1,164 @@
+-wb3JWo4qTg_380000_390000
+0BsmPd44Bic_60000_70000
+2VqnoV67t0g_290000_300000
+2wmxck5D9Vw_102000_112000
+39xq5AATMp4_42000_52000
+39xq5AATMp4_60000_70000
+45oaGY3mzlQ_5000_15000
+4DCxlVTWN8g_3000_13000
+4GUrjmIZcIQ_82000_92000
+4YLjVRZHoZI_1750000_1760000
+4ZPJ867OBo4_75000_85000
+4k63MOSjOBw_21000_31000
+4rerQA8KMzM_10000_20000
+4tASCYb3ySA_1000_11000
+50Z4O0mzAXg_150000_160000
+5FYG_Wtet7U_120000_130000
+5YLQapaxA8o_31000_41000
+5YLQapaxA8o_80000_90000
+6KT2mbCNxO4_228000_238000
+6KT2mbCNxO4_434000_444000
+6WUjWMuoEhM_10000_20000
+6tTLxIKVtfE_26000_36000
+73QQbJIeB3Y_265000_275000
+73QQbJIeB3Y_353000_363000
+73QQbJIeB3Y_95000_105000
+79T0FclnIDw_0_10000
+79w-8fTYazw_15000_25000
+7PGfd8pg86w_15000_25000
+7pNykt6zACc_275000_285000
+7pWa_kcAoMg_49000_59000
+872Lwp3MOro_6000_16000
+8hr42lVW_gk_60000_70000
+9ZQ-T83pfWk_65000_75000
+9xp46AwF9BY_38000_48000
+A-e3dr7fTIs_0_10000
+A-e3dr7fTIs_318000_328000
+A-e3dr7fTIs_80000_90000
+AxjZti5_t94_90000_100000
+BRAQLY85x5U_75000_85000
+BSkcM4ScyEs_150000_160000
+CCFZT2_TJr0_15000_25000
+CPiOGMl59L4_30000_40000
+CVA8LpBW3Sc_76000_86000
+CjQmcO_Q5d8_42000_52000
+De5YhqX0jdI_0_10000
+DzXVMjsZloE_0_10000
+EKM1wu2vXy4_26000_36000
+EQWvv8WbU04_22000_32000
+ETmo71vP7tA_20000_30000
+EuU2PWB1t2g_20000_30000
+EzPcuNoSHMM_0_10000
+F1x1ck3OySg_10000_20000
+F3KWGQfJ2HM_26000_36000
+GG-PF_JxeW4_1640000_1650000
+HuFzVoyayj8_521000_531000
+JgkUToh3HeY_49000_59000
+KXvmc3dLn3E_720000_730000
+LOFX2UVozf8_80000_90000
+LitaFzObEEk_5000_15000
+MDyjY3uiWp0_273000_283000
+Maa21OL-40Q_20000_30000
+Mivqxr0RS8w_18000_28000
+OFgcrlxku9g_160000_170000
+OGCyzmaM_kE_10000_20000
+OMYuLiqSUxE_180000_190000
+OMYuLiqSUxE_30000_40000
+OPVEPq_r-vk_211000_221000
+OX6T2z4P9fA_23000_33000
+Ow9uE_v2AEg_28000_38000
+PSoKYh3ea1o_60000_70000
+PcdKAvd51l0_41000_51000
+PdHpl04tQV8_40000_50000
+Pe1LuVFTczE_106000_116000
+Pe1LuVFTczE_358000_368000
+QHcG-FDM75Q_113000_123000
+QNiHU290owU_55000_65000
+QTe-i0Pcn4s_37000_47000
+RHy5nC-gRV8_668000_678000
+RMF6sp6tWHM_100000_110000
+Ru7m8PyMlVM_120000_130000
+SFVZ2OklsVM_2715000_2725000
+SFVZ2OklsVM_2765000_2775000
+T9K1uy-G5qA_110000_120000
+TCcD-vOUtNc_99000_109000
+U1dZX1ReD88_48000_58000
+URZyjoh9lbc_500000_510000
+URZyjoh9lbc_560000_570000
+UYUH7Jmfp3g_13000_23000
+UYUH7Jmfp3g_76000_86000
+UlYU9z7Y8jY_68000_78000
+V9KZ5FCtG9A_15000_25000
+VbEvfbj_IxU_117000_127000
+VlPdfLr1FSo_7000_17000
+W5yveLPTD04_211000_221000
+WSy8ay1avew_60000_70000
+Y735cxoG5-4_270000_280000
+YAW2vMKV9pw_50000_60000
+ZPYqUww_x6k_293000_303000
+ZR3vnlhJuSE_101000_111000
+ZU0JSxWk1Po_16000_26000
+_-apT0tfo6U_16000_26000
+a1nWlW629TU_15000_25000
+aBmzZJZ_M8Y_9000_19000
+bELyeHxF7eA_42000_52000
+bcGfmy0X-CQ_30000_40000
+cIbVu0ixSAo_302000_312000
+dRa7aBGnStU_82000_92000
+dUcOkRkz6bA_387000_397000
+dVa49WwXzr8_1023000_1033000
+dk_xhLkWyDo_30000_40000
+ds3RKnNB-cY_8000_18000
+ehlPuuiNEd8_376000_386000
+f0jKjIOFzAY_243000_253000
+fTobKZBbMos_2000_12000
+fWvQqgSDUPU_10000_20000
+hYt2Qf438l8_40000_50000
+hornh-NQBHY_262000_272000
+iSBtK1T10Ew_415000_425000
+jj1UxRTBaNw_210000_220000
+k35blcO8Z7k_700000_710000
+kF2y7RIC7-Y_258000_268000
+kF2y7RIC7-Y_385000_395000
+lcOP60uXMeI_20000_30000
+lcOP60uXMeI_396000_406000
+lzzMHoi3r2w_50000_60000
+mUN93MlvX64_42000_52000
+nT0PHpAlvys_123000_133000
+nT0PHpAlvys_313000_323000
+nT0PHpAlvys_435000_445000
+oRSCL3149fI_146000_156000
+oRSCL3149fI_87000_97000
+oSMvY0tErC4_100000_110000
+oSMvY0tErC4_50000_60000
+oVK2QsKq8ak_186000_196000
+oYeir4FWq_8_6000_16000
+oaYIWYXFMqY_70000_80000
+oyi5pKPwz9Q_3000_13000
+p2u4OJKqMxE_130000_140000
+p5uRwtw7S3E_243000_253000
+pNV8MKNqOkI_20000_30000
+pf6ZpxTFL1Y_51000_61000
+puugfzdXYz4_30000_40000
+q3YnhPgt-rM_118000_128000
+qYtrnr4chfU_2000_12000
+r4NdM595K5c_40000_50000
+rozFJYWrLj0_120000_130000
+rzKjN2en0H8_830000_840000
+sHyhvtLTCbo_10000_20000
+sMsrz5VqchQ_130000_140000
+sfv_msSOYTo_10000_20000
+tMorLZku6Pc_194000_204000
+tZbh1cwwfv0_222000_232000
+vVJ-Zhj2HvU_22000_32000
+vifFbeL5rOo_9000_19000
+w7vA5f0vPvQ_15000_25000
+wfkdedUW-dk_50000_60000
+xUQ9rDswHdw_140000_150000
+xnx3u5YaNuc_402000_412000
+z3Q_mZgKLrM_20000_30000
+zM7QopQ3MgI_319000_329000
+zM7QopQ3MgI_410000_420000
+zM7QopQ3MgI_499000_509000
+zPMUL7f4OOU_420000_430000
+zstao4nIPmU_250000_260000

runs/tubetoken_phase0/proposals_stride8_n64_bidir.log

@@ -0,0 +1 @@
+

runs/tubetoken_phase_minus1/audit_full/audit_summary.json

@@ -0,0 +1,45 @@
+{
+  "area_unstable_expressions": 41,
+  "audio_keyword_expressions": 15890,
+  "audio_keyword_percent": 77.66753018231586,
+  "data_dir": "/workspace/SimToken/data",
+  "expressions_per_object": {
+    "ge2": 5836,
+    "ge3": 4206,
+    "max": 10,
+    "mean": 2.742125720412813,
+    "median": 3
+  },
+  "expressions_per_video": {
+    "ge2": 3521,
+    "ge3": 3381,
+    "max": 26,
+    "mean": 5.7243984331281474,
+    "median": 6.0
+  },
+  "h3_candidate_expressions": 18614,
+  "h3_candidate_objects": 5781,
+  "late_target_expressions": 0,
+  "mask_rows_audited": 20459,
+  "multi_expression_objects": 5836,
+  "multi_expression_videos": 3521,
+  "null_split_expressions": 1028,
+  "null_split_percent": 5.0246835133682,
+  "num_expressions": 20459,
+  "num_objects_vid_fid": 7461,
+  "num_videos": 3574,
+  "partial_target_expressions": 33,
+  "same_category_distractor_heuristic_expressions": 2563,
+  "same_category_distractor_heuristic_percent": 12.527494012415074,
+  "small_target_expressions": 10037,
+  "spatial_keyword_expressions": 5924,
+  "spatial_keyword_percent": 28.955471919448655,
+  "splits": {
+    "TODO": 25,
+    "test_n": 1028,
+    "test_s": 2288,
+    "test_u": 1656,
+    "train": 14113,
+    "val": 1349
+  }
+}

runs/tubetoken_phase_minus1/simtoken_eval/eval_null.log

@@ -0,0 +1,11 @@
+seg_token_idx:  32000
+
+
+model loaded
+
+Lora deployed
+trainable params: 4194304 || all params: 7709437232 || trainable%: 0.0544048012037826
+saved model loaded
+
+
+  valuate on test_n_refer, metric: 0.011791757307946682

tools/tubetoken/evaluate_phase0_proposals.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+"""Evaluate proposal recall and oracle tube J/F for TubeToken Phase 0."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List
+
+import numpy as np
+from tqdm import tqdm
+
+from phase0_common import (
+    bool_field,
+    evaluate_tube_jf,
+    load_audit_rows,
+    load_gt_tube,
+    read_metadata,
+    rows_by_video,
+    rows_by_video,
+    tube_iou_all,
+    tube_iou_visible,
+    video_id,
+    fid_value,
+    write_json,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Evaluate Phase 0 proposal cache.")
+    parser.add_argument("--data_dir", type=Path, required=True)
+    parser.add_argument("--proposal_dir", type=Path, required=True)
+    parser.add_argument("--out_dir", type=Path, required=True)
+    parser.add_argument("--audit_csv", type=Path, default=None)
+    parser.add_argument("--splits", type=str, default="test_s,test_u")
+    parser.add_argument("--frames", type=int, default=10)
+    parser.add_argument("--recall_ns", type=str, default="16,32,64,128")
+    parser.add_argument("--match_iou", type=float, default=0.5)
+    parser.add_argument("--limit_videos", type=int, default=0)
+    parser.add_argument("--video_list", type=Path, default=None)
+    parser.add_argument("--shard_id", type=int, default=0)
+    parser.add_argument("--num_shards", type=int, default=1)
+    parser.add_argument("--only_existing_proposals", action="store_true")
+    parser.add_argument("--skip_oracle_jf", action="store_true", help="Only compute Recall@N and tube IoU, useful for fast early checks.")
+    return parser.parse_args()
+
+
+def load_proposals(path: Path) -> np.ndarray:
+    data = np.load(path)
+    return data["masks"].astype(bool)
+
+
+def load_video_list(path: Path | None) -> List[str] | None:
+    if path is None:
+        return None
+    vids = []
+    for line in path.read_text().splitlines():
+        line = line.strip()
+        if line and not line.startswith("#"):
+            vids.append(line)
+    return vids
+
+
+def sample_subsets(row: dict, audit: Dict[str, dict]) -> List[str]:
+    out = ["all", row["split"]]
+    audit_row = audit.get(row["uid"])
+    if bool_field(audit_row, "small_target"):
+        out.append("small")
+    if bool_field(audit_row, "partial_target"):
+        out.append("partial")
+    if bool_field(audit_row, "area_unstable"):
+        out.append("area_unstable")
+    if bool_field(audit_row, "late_target"):
+        out.append("late_target")
+    if bool_field(audit_row, "is_audio_keyword"):
+        out.append("audio_keyword")
+    if bool_field(audit_row, "is_spatial_keyword"):
+        out.append("spatial_keyword")
+    if bool_field(audit_row, "same_category_distractor_heuristic"):
+        out.append("same_category")
+    if bool_field(audit_row, "h3_candidate"):
+        out.append("h3_candidate")
+    return out
+
+
+def empty_metrics(recall_ns: List[int]) -> dict:
+    return {
+        "count": 0,
+        "proposal_miss": 0,
+        "oracle_j": 0.0,
+        "oracle_f": 0.0,
+        "oracle_jf": 0.0,
+        "oracle_iou_visible": 0.0,
+        "oracle_iou_all": 0.0,
+        **{f"recall@{n}": 0 for n in recall_ns},
+    }
+
+
+def add_metrics(bucket: dict, sample: dict, recall_ns: List[int]) -> None:
+    bucket["count"] += 1
+    bucket["proposal_miss"] += int(not sample["covered"])
+    bucket["oracle_j"] += sample["oracle_j"]
+    bucket["oracle_f"] += sample["oracle_f"]
+    bucket["oracle_jf"] += sample["oracle_jf"]
+    bucket["oracle_iou_visible"] += sample["oracle_iou_visible"]
+    bucket["oracle_iou_all"] += sample["oracle_iou_all"]
+    for n in recall_ns:
+        bucket[f"recall@{n}"] += int(sample[f"recall@{n}"])
+
+
+def finalize(bucket: dict, recall_ns: List[int]) -> dict:
+    count = bucket["count"]
+    if count == 0:
+        return dict(bucket)
+    out = dict(bucket)
+    out["proposal_miss_percent"] = 100.0 * bucket["proposal_miss"] / count
+    for key in ["oracle_j", "oracle_f", "oracle_jf", "oracle_iou_visible", "oracle_iou_all"]:
+        out[key] = bucket[key] / count
+    for n in recall_ns:
+        out[f"recall@{n}"] = bucket[f"recall@{n}"] / count
+    return out
+
+
+def main() -> None:
+    args = parse_args()
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    splits = [s.strip() for s in args.splits.split(",") if s.strip()]
+    recall_ns = [int(x) for x in args.recall_ns.split(",") if x.strip()]
+    audit = load_audit_rows(args.audit_csv) if args.audit_csv else {}
+    rows = read_metadata(args.data_dir, splits)
+    selected_vids = load_video_list(args.video_list)
+    if selected_vids is not None:
+        selected = set(selected_vids)
+        rows = [row for row in rows if video_id(row) in selected]
+    if args.num_shards < 1:
+        raise ValueError("--num_shards must be >= 1")
+    if args.shard_id < 0 or args.shard_id >= args.num_shards:
+        raise ValueError("--shard_id must be in [0, num_shards)")
+    if args.num_shards > 1:
+        vids = sorted(rows_by_video(rows).keys())
+        selected = {vid for idx, vid in enumerate(vids) if idx % args.num_shards == args.shard_id}
+        rows = [row for row in rows if video_id(row) in selected]
+    if args.limit_videos:
+        vids = sorted(rows_by_video(rows).keys())[: args.limit_videos]
+        rows = [row for row in rows if video_id(row) in set(vids)]
+    if args.only_existing_proposals:
+        rows = [row for row in rows if (args.proposal_dir / f"{video_id(row)}.npz").exists()]
+
+    sample_rows: List[dict] = []
+    summary = defaultdict(lambda: empty_metrics(recall_ns))
+
+    video_groups = rows_by_video(rows)
+    total_objects = sum(len({fid_value(row) for row in group}) for group in video_groups.values())
+
+    with tqdm(total=total_objects, desc="Evaluating proposal objects") as pbar:
+        for vid, video_rows in video_groups.items():
+            prop_path = args.proposal_dir / f"{vid}.npz"
+            if not prop_path.exists():
+                raise FileNotFoundError(f"Missing proposal cache: {prop_path}")
+            proposals = load_proposals(prop_path)
+            object_cache = {}
+
+            for row in video_rows:
+                key = fid_value(row)
+                if key in object_cache:
+                    base_sample = object_cache[key]
+                else:
+                    gt = load_gt_tube(args.data_dir, vid, key, args.frames)
+                    visible_ious = np.array([tube_iou_visible(tube, gt) for tube in proposals], dtype=np.float32)
+                    all_ious = np.array([tube_iou_all(tube, gt) for tube in proposals], dtype=np.float32)
+                    if len(visible_ious) == 0:
+                        best_idx = -1
+                        best_visible = 0.0
+                        best_all = 0.0
+                        oracle_j = oracle_f = oracle_jf = 0.0
+                    else:
+                        best_idx = int(visible_ious.argmax())
+                        best_visible = float(visible_ious[best_idx])
+                        best_all = float(all_ious[best_idx])
+                        if args.skip_oracle_jf:
+                            oracle_j = oracle_f = oracle_jf = 0.0
+                        else:
+                            oracle_j, oracle_f, oracle_jf = evaluate_tube_jf(proposals[best_idx], gt)
+
+                    base_sample = {
+                        "vid": vid,
+                        "fid": key,
+                        "best_idx": best_idx,
+                        "num_tubes": int(proposals.shape[0]),
+                        "covered": best_visible >= args.match_iou,
+                        "oracle_iou_visible": best_visible,
+                        "oracle_iou_all": best_all,
+                        "oracle_j": oracle_j,
+                        "oracle_f": oracle_f,
+                        "oracle_jf": oracle_jf,
+                    }
+                    for n in recall_ns:
+                        top = visible_ious[: min(n, len(visible_ious))]
+                        base_sample[f"recall@{n}"] = bool(len(top) and float(top.max()) >= args.match_iou)
+                    object_cache[key] = base_sample
+                    pbar.update(1)
+
+                sample = dict(base_sample)
+                sample.update({"uid": row["uid"], "split": row["split"]})
+                sample_rows.append(sample)
+                for subset in sample_subsets(row, audit):
+                    add_metrics(summary[subset], sample, recall_ns)
+
+    with (args.out_dir / "sample_metrics.csv").open("w", newline="") as f:
+        fieldnames = list(sample_rows[0].keys()) if sample_rows else []
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        writer.writerows(sample_rows)
+
+    final_summary = {name: finalize(bucket, recall_ns) for name, bucket in sorted(summary.items())}
+    write_json(args.out_dir / "summary.json", final_summary)
+
+    md = ["# TubeToken Phase 0 Proposal Evaluation", ""]
+    for name, metrics in final_summary.items():
+        if metrics["count"] == 0:
+            continue
+        recall_text = ", ".join(f"R@{n}={metrics[f'recall@{n}']:.3f}" for n in recall_ns)
+        md.append(
+            f"- {name}: n={metrics['count']}, {recall_text}, "
+            f"Oracle J&F={metrics['oracle_jf']:.4f}, miss={metrics['proposal_miss_percent']:.2f}%"
+        )
+    (args.out_dir / "report.md").write_text("\n".join(md) + "\n")
+    print("\n".join(md))
+
+
+if __name__ == "__main__":
+    main()

tools/tubetoken/generate_sam2_proposals.py

@@ -0,0 +1,356 @@
+#!/usr/bin/env python3
+"""Generate SAM2 proposal tubes for TubeToken Phase 0.
+
+The cache format is one NPZ per video:
+  masks: uint8 [N, T, H, W]
+  scores: float32 [N]
+  keyframes: int64 [N]
+  boxes_xyxy: float32 [N, 4]
+"""
+
+from __future__ import annotations
+
+import argparse
+from contextlib import nullcontext
+import os
+import sys
+import time
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+from PIL import Image
+from tqdm import tqdm
+
+from phase0_common import read_metadata, rows_by_video, video_id
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate SAM2 proposal tubes.")
+    parser.add_argument("--data_dir", type=Path, required=True)
+    parser.add_argument("--out_dir", type=Path, required=True)
+    parser.add_argument("--splits", type=str, default="test_s,test_u")
+    parser.add_argument("--sam2_repo", type=Path, default=None, help="Path to a local facebookresearch/sam2 clone.")
+    parser.add_argument("--model_cfg", type=str, default="configs/sam2.1/sam2.1_hiera_l.yaml")
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    parser.add_argument("--seed_proposal_dir", type=Path, default=None, help="Reuse boxes/keyframes/scores from an existing proposal cache and only rerun propagation.")
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--amp_dtype", type=str, default="bf16", choices=["none", "bf16", "fp16"])
+    parser.add_argument("--frames", type=int, default=10)
+    parser.add_argument("--stride", type=int, default=8)
+    parser.add_argument("--max_tubes", type=int, default=128)
+    parser.add_argument("--amg_points_per_side", type=int, default=32)
+    parser.add_argument("--amg_pred_iou_thresh", type=float, default=0.80)
+    parser.add_argument("--amg_stability_score_thresh", type=float, default=0.88)
+    parser.add_argument("--min_mask_area", type=int, default=64)
+    parser.add_argument("--limit_videos", type=int, default=0)
+    parser.add_argument("--video_list", type=Path, default=None)
+    parser.add_argument("--shard_id", type=int, default=0)
+    parser.add_argument("--num_shards", type=int, default=1)
+    parser.add_argument("--quiet_sam2", action="store_true")
+    parser.add_argument("--bidirectional", action="store_true", default=True)
+    parser.set_defaults(group_by_keyframe=False)
+    parser.add_argument("--group_by_keyframe", dest="group_by_keyframe", action="store_true")
+    parser.add_argument("--no_group_by_keyframe", dest="group_by_keyframe", action="store_false")
+    parser.add_argument("--overwrite", action="store_true")
+    return parser.parse_args()
+
+
+def import_sam2(repo: Path | None):
+    if repo is not None:
+        sys.path.insert(0, str(repo))
+    from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
+    from sam2.build_sam import build_sam2, build_sam2_video_predictor
+
+    return SAM2AutomaticMaskGenerator, build_sam2, build_sam2_video_predictor
+
+
+def keyframes_for_stride(frames: int, stride: int) -> List[int]:
+    if stride <= 0:
+        raise ValueError("--stride must be positive")
+    keyframes = list(range(0, frames, stride))
+    mid = frames // 2
+    if mid not in keyframes:
+        keyframes.append(mid)
+    return sorted(set(k for k in keyframes if 0 <= k < frames))
+
+
+def load_rgb(path: Path) -> np.ndarray:
+    with Image.open(path) as img:
+        return np.array(img.convert("RGB"))
+
+
+def proposal_score(item: dict) -> float:
+    pred_iou = float(item.get("predicted_iou", 0.0))
+    stability = float(item.get("stability_score", 0.0))
+    area = float(item.get("area", 0.0))
+    return pred_iou + 0.1 * stability + min(area / 1_000_000.0, 0.01)
+
+
+def xywh_to_xyxy(box: List[float]) -> np.ndarray:
+    x, y, w, h = box
+    return np.array([x, y, x + w, y + h], dtype=np.float32)
+
+
+def amp_context(device: str, amp_dtype: str):
+    if not device.startswith("cuda") or amp_dtype == "none":
+        return nullcontext()
+    dtype = torch.bfloat16 if amp_dtype == "bf16" else torch.float16
+    return torch.autocast("cuda", dtype=dtype)
+
+
+def collect_keyframe_masks(mask_generator, data_dir: Path, vid: str, keyframes: List[int], min_area: int, device: str, amp_dtype: str) -> List[dict]:
+    proposals: List[dict] = []
+    for kf in keyframes:
+        image = load_rgb(data_dir / "media" / vid / "frames" / f"{kf}.jpg")
+        with torch.inference_mode(), amp_context(device, amp_dtype):
+            masks = mask_generator.generate(image)
+        for m in masks:
+            area = int(m.get("area", np.asarray(m["segmentation"]).sum()))
+            if area < min_area:
+                continue
+            bbox = xywh_to_xyxy(m["bbox"])
+            proposals.append(
+                {
+                    "keyframe": kf,
+                    "box": bbox,
+                    "score": proposal_score(m),
+                    "area": area,
+                }
+            )
+    proposals.sort(key=lambda x: x["score"], reverse=True)
+    return proposals
+
+
+def load_seed_proposals(seed_dir: Path, vid: str) -> List[dict] | None:
+    path = seed_dir / f"{vid}.npz"
+    if not path.exists():
+        return None
+    data = np.load(path)
+    boxes = data["boxes_xyxy"]
+    keyframes = data["keyframes"]
+    scores = data["scores"]
+    proposals = []
+    for box, keyframe, score in zip(boxes, keyframes, scores):
+        proposals.append(
+            {
+                "keyframe": int(keyframe),
+                "box": box.astype(np.float32),
+                "score": float(score),
+                "area": 0,
+            }
+        )
+    proposals.sort(key=lambda x: x["score"], reverse=True)
+    return proposals
+
+
+def add_box_prompt(predictor, state, frame_idx: int, obj_id: int, box: np.ndarray):
+    try:
+        return predictor.add_new_points_or_box(
+            inference_state=state,
+            frame_idx=frame_idx,
+            obj_id=obj_id,
+            box=box,
+        )
+    except TypeError:
+        return predictor.add_new_points_or_box(state, frame_idx, obj_id, box=box)
+
+
+def logits_to_mask(logits) -> np.ndarray:
+    if hasattr(logits, "detach"):
+        logits = logits.detach().cpu().numpy()
+    arr = np.asarray(logits)
+    while arr.ndim > 2:
+        arr = arr[0]
+    return arr > 0
+
+
+def run_propagation(state, predictor, masks_by_obj: Dict[int, List[np.ndarray]], frames: int, reverse: bool, start_frame_idx: int | None = None) -> None:
+    for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(
+        state,
+        start_frame_idx=start_frame_idx,
+        reverse=reverse,
+    ):
+        if hasattr(out_obj_ids, "detach"):
+            out_obj_ids = out_obj_ids.detach().cpu().tolist()
+        for pos, obj_id in enumerate(list(out_obj_ids)):
+            if 0 <= int(out_frame_idx) < frames:
+                masks_by_obj[int(obj_id)][int(out_frame_idx)] = logits_to_mask(out_mask_logits[pos])
+
+
+def propagate_proposal_group(predictor, video_dir: Path, proposals: List[dict], frames: int, device: str, amp_dtype: str, bidirectional: bool) -> np.ndarray:
+    with torch.inference_mode(), amp_context(device, amp_dtype):
+        state = predictor.init_state(video_path=str(video_dir))
+        for obj_id, proposal in enumerate(proposals):
+            add_box_prompt(predictor, state, int(proposal["keyframe"]), obj_id, proposal["box"])
+
+        masks_by_obj: Dict[int, List[np.ndarray]] = {
+            obj_id: [None for _ in range(frames)] for obj_id in range(len(proposals))
+        }
+        run_propagation(state, predictor, masks_by_obj, frames, reverse=False)
+        if bidirectional:
+            start_frame_idx = max(int(p["keyframe"]) for p in proposals)
+            run_propagation(state, predictor, masks_by_obj, frames, reverse=True, start_frame_idx=start_frame_idx)
+
+    # Some SAM2 versions only propagate forward. Fill missing frames with the
+    # nearest available mask so Phase 0 can still score temporal purity.
+    tube_masks = []
+    for obj_id in range(len(proposals)):
+        masks = masks_by_obj[obj_id]
+        known = [i for i, m in enumerate(masks) if m is not None]
+        if not known:
+            continue
+        for t in range(frames):
+            if masks[t] is None:
+                nearest = min(known, key=lambda k: abs(k - t))
+                masks[t] = masks[nearest]
+        tube_masks.append(np.stack(masks, axis=0))
+    if not tube_masks:
+        return np.zeros((0, frames, 1, 1), dtype=np.uint8)
+    return np.stack(tube_masks, axis=0).astype(np.uint8)
+
+
+def propagate_boxes(
+    predictor,
+    video_dir: Path,
+    proposals: List[dict],
+    frames: int,
+    device: str,
+    amp_dtype: str,
+    bidirectional: bool,
+    group_by_keyframe: bool,
+) -> np.ndarray:
+    if not group_by_keyframe:
+        return propagate_proposal_group(predictor, video_dir, proposals, frames, device, amp_dtype, bidirectional)
+
+    grouped: Dict[int, List[Tuple[int, dict]]] = {}
+    for idx, proposal in enumerate(proposals):
+        grouped.setdefault(int(proposal["keyframe"]), []).append((idx, proposal))
+
+    ordered_masks: List[np.ndarray | None] = [None for _ in proposals]
+    for _, indexed_group in sorted(grouped.items()):
+        group_indices = [idx for idx, _ in indexed_group]
+        group_props = [proposal for _, proposal in indexed_group]
+        group_masks = propagate_proposal_group(predictor, video_dir, group_props, frames, device, amp_dtype, bidirectional)
+        for local_idx, global_idx in enumerate(group_indices[: group_masks.shape[0]]):
+            ordered_masks[global_idx] = group_masks[local_idx]
+
+    known = [mask for mask in ordered_masks if mask is not None]
+    if not known:
+        return np.zeros((0, frames, 1, 1), dtype=np.uint8)
+    h, w = known[0].shape[-2:]
+    final = [mask if mask is not None else np.zeros((frames, h, w), dtype=np.uint8) for mask in ordered_masks]
+    return np.stack(final, axis=0).astype(np.uint8)
+
+
+def load_video_list(path: Path | None) -> List[str] | None:
+    if path is None:
+        return None
+    vids = []
+    for line in path.read_text().splitlines():
+        line = line.strip()
+        if line and not line.startswith("#"):
+            vids.append(line)
+    return vids
+
+
+def main() -> None:
+    args = parse_args()
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    SAM2AutomaticMaskGenerator, build_sam2, build_sam2_video_predictor = import_sam2(args.sam2_repo)
+
+    mask_generator = None
+    if args.seed_proposal_dir is None:
+        image_model = build_sam2(args.model_cfg, str(args.checkpoint), device=args.device)
+        mask_generator = SAM2AutomaticMaskGenerator(
+            image_model,
+            points_per_side=args.amg_points_per_side,
+            pred_iou_thresh=args.amg_pred_iou_thresh,
+            stability_score_thresh=args.amg_stability_score_thresh,
+        )
+    video_predictor = build_sam2_video_predictor(args.model_cfg, str(args.checkpoint), device=args.device)
+
+    splits = [s.strip() for s in args.splits.split(",") if s.strip()]
+    rows = read_metadata(args.data_dir, splits)
+    vids = sorted(rows_by_video(rows).keys())
+    selected_vids = load_video_list(args.video_list)
+    if selected_vids is not None:
+        selected = set(selected_vids)
+        vids = [vid for vid in vids if vid in selected]
+    if args.num_shards < 1:
+        raise ValueError("--num_shards must be >= 1")
+    if args.shard_id < 0 or args.shard_id >= args.num_shards:
+        raise ValueError("--shard_id must be in [0, num_shards)")
+    if args.num_shards > 1:
+        vids = [vid for idx, vid in enumerate(vids) if idx % args.num_shards == args.shard_id]
+    if args.limit_videos:
+        vids = vids[: args.limit_videos]
+    keyframes = keyframes_for_stride(args.frames, args.stride)
+
+    manifest = {
+        "data_dir": str(args.data_dir),
+        "splits": splits,
+        "model_cfg": args.model_cfg,
+        "checkpoint": str(args.checkpoint),
+        "stride": args.stride,
+        "keyframes": keyframes,
+        "max_tubes": args.max_tubes,
+        "videos": len(vids),
+        "items": [],
+    }
+
+    for vid in tqdm(vids, desc="Generating SAM2 proposals"):
+        out_path = args.out_dir / f"{vid}.npz"
+        if out_path.exists() and not args.overwrite:
+            manifest["items"].append({"vid": vid, "path": str(out_path), "skipped": True})
+            continue
+
+        start = time.perf_counter()
+        proposals = load_seed_proposals(args.seed_proposal_dir, vid) if args.seed_proposal_dir is not None else None
+        if proposals is None:
+            proposals = collect_keyframe_masks(mask_generator, args.data_dir, vid, keyframes, args.min_mask_area, args.device, args.amp_dtype)
+        proposals = proposals[: args.max_tubes]
+        if args.quiet_sam2:
+            from contextlib import redirect_stdout, redirect_stderr
+            import io
+
+            with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()):
+                if proposals:
+                    masks = propagate_boxes(video_predictor, args.data_dir / "media" / vid / "frames", proposals, args.frames, args.device, args.amp_dtype, args.bidirectional, args.group_by_keyframe)
+                else:
+                    first = load_rgb(args.data_dir / "media" / vid / "frames" / "0.jpg")
+                    h, w = first.shape[:2]
+                    masks = np.zeros((0, args.frames, h, w), dtype=np.uint8)
+        else:
+            if proposals:
+                masks = propagate_boxes(video_predictor, args.data_dir / "media" / vid / "frames", proposals, args.frames, args.device, args.amp_dtype, args.bidirectional, args.group_by_keyframe)
+            else:
+                first = load_rgb(args.data_dir / "media" / vid / "frames" / "0.jpg")
+                h, w = first.shape[:2]
+                masks = np.zeros((0, args.frames, h, w), dtype=np.uint8)
+
+        n = min(len(proposals), masks.shape[0])
+        proposals = proposals[:n]
+        masks = masks[:n]
+        scores = np.array([p["score"] for p in proposals], dtype=np.float32)
+        boxes = np.stack([p["box"] for p in proposals], axis=0).astype(np.float32) if proposals else np.zeros((0, 4), dtype=np.float32)
+        proposal_keyframes = np.array([p["keyframe"] for p in proposals], dtype=np.int64)
+        np.savez_compressed(
+            out_path,
+            masks=masks,
+            scores=scores,
+            keyframes=proposal_keyframes,
+            boxes_xyxy=boxes,
+        )
+        elapsed = time.perf_counter() - start
+        manifest["items"].append({"vid": vid, "path": str(out_path), "tubes": int(n), "seconds": elapsed})
+
+    with (args.out_dir / "manifest.json").open("w") as f:
+        import json
+
+        json.dump(manifest, f, indent=2)
+
+
+if __name__ == "__main__":
+    main()

tools/tubetoken/phase0_common.py

@@ -0,0 +1,214 @@
+"""Shared utilities for TubeToken Phase 0 experiments."""
+
+from __future__ import annotations
+
+import csv
+import json
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+import numpy as np
+from PIL import Image
+
+try:
+    import cv2
+except Exception:  # pragma: no cover
+    cv2 = None
+
+
+VALID_EVAL_SPLITS = {"test_s", "test_u"}
+VALID_DATA_SPLITS = {"train", "val", "test_s", "test_u", "test_n"}
+
+
+def read_metadata(data_dir: Path, splits: Optional[Iterable[str]] = None) -> List[dict]:
+    wanted = set(splits) if splits else None
+    with (data_dir / "metadata.csv").open("r", newline="") as f:
+        rows = list(csv.DictReader(f))
+    if wanted is not None:
+        rows = [r for r in rows if r["split"] in wanted]
+    return rows
+
+
+def video_id(row: dict) -> str:
+    return row.get("vid") or row["uid"].rsplit("_", 2)[0]
+
+
+def fid_value(row: dict) -> str:
+    return str(row.get("fid", "")).strip()
+
+
+def category_from_uid(row: dict) -> str:
+    vid = video_id(row)
+    uid = row.get("uid", "")
+    suffix = uid[len(vid) + 1 :] if uid.startswith(vid + "_") else uid.rsplit("_", 2)[-2]
+    return suffix.rsplit("_", 1)[0] if "_" in suffix else suffix
+
+
+def rows_by_video(rows: Sequence[dict]) -> Dict[str, List[dict]]:
+    out: Dict[str, List[dict]] = defaultdict(list)
+    for row in rows:
+        out[video_id(row)].append(row)
+    return out
+
+
+def load_mask(path: Path) -> np.ndarray:
+    with Image.open(path) as img:
+        return np.array(img.convert("L")) > 0
+
+
+def load_gt_tube(data_dir: Path, vid: str, fid: str, frames: int = 10) -> np.ndarray:
+    masks = []
+    for t in range(frames):
+        path = data_dir / "gt_mask" / vid / f"fid_{fid}" / f"0000{t}.png"
+        masks.append(load_mask(path))
+    return np.stack(masks, axis=0)
+
+
+def mask_iou(pred: np.ndarray, gt: np.ndarray) -> float:
+    pred = pred.astype(bool)
+    gt = gt.astype(bool)
+    union = np.logical_or(pred, gt).sum()
+    if union == 0:
+        return 1.0
+    inter = np.logical_and(pred, gt).sum()
+    return float(inter / union)
+
+
+def tube_iou_visible(pred_tube: np.ndarray, gt_tube: np.ndarray) -> float:
+    visible = gt_tube.reshape(gt_tube.shape[0], -1).sum(axis=1) > 0
+    if not visible.any():
+        return 0.0
+    vals = [mask_iou(pred_tube[t], gt_tube[t]) for t in np.where(visible)[0]]
+    return float(np.mean(vals)) if vals else 0.0
+
+
+def tube_iou_all(pred_tube: np.ndarray, gt_tube: np.ndarray) -> float:
+    vals = [mask_iou(pred_tube[t], gt_tube[t]) for t in range(gt_tube.shape[0])]
+    return float(np.mean(vals)) if vals else 0.0
+
+
+def db_eval_iou(annotation: np.ndarray, segmentation: np.ndarray) -> float:
+    annotation = annotation.astype(bool)
+    segmentation = segmentation.astype(bool)
+    if annotation.sum() == 0 and segmentation.sum() == 0:
+        return 1.0
+    if annotation.sum() == 0 or segmentation.sum() == 0:
+        return 0.0
+    inter = np.logical_and(annotation, segmentation).sum()
+    union = np.logical_or(annotation, segmentation).sum()
+    return float(inter / union) if union > 0 else 0.0
+
+
+def db_eval_boundary(annotation: np.ndarray, segmentation: np.ndarray, bound_th: float = 0.008) -> float:
+    annotation = annotation.astype(bool)
+    segmentation = segmentation.astype(bool)
+    if annotation.sum() == 0 and segmentation.sum() == 0:
+        return 1.0
+    if annotation.sum() == 0 or segmentation.sum() == 0:
+        return 0.0
+
+    bound_pix = max(1, int(round(bound_th * np.linalg.norm(annotation.shape))))
+    if cv2 is not None:
+        fg_boundary = mask_to_boundary_cv2(annotation, bound_pix)
+        seg_boundary = mask_to_boundary_cv2(segmentation, bound_pix)
+        kernel = np.ones((2 * bound_pix + 1, 2 * bound_pix + 1), dtype=np.uint8)
+        fg_dil = cv2.dilate(fg_boundary.astype(np.uint8), kernel, iterations=1).astype(bool)
+        seg_dil = cv2.dilate(seg_boundary.astype(np.uint8), kernel, iterations=1).astype(bool)
+        gt_match = np.logical_and(fg_boundary, seg_dil).sum()
+        pred_match = np.logical_and(seg_boundary, fg_dil).sum()
+        n_fg = fg_boundary.sum()
+        n_pred = seg_boundary.sum()
+        if n_fg == 0 and n_pred == 0:
+            return 1.0
+        if n_fg == 0 or n_pred == 0:
+            return 0.0
+        precision = pred_match / n_pred
+        recall = gt_match / n_fg
+        if precision + recall == 0:
+            return 0.0
+        return float(2 * precision * recall / (precision + recall))
+
+    fg_boundary = mask_to_boundary(annotation, bound_pix)
+    seg_boundary = mask_to_boundary(segmentation, bound_pix)
+    fg_dil = binary_dilate(fg_boundary, bound_pix)
+    seg_dil = binary_dilate(seg_boundary, bound_pix)
+
+    gt_match = np.logical_and(fg_boundary, seg_dil).sum()
+    pred_match = np.logical_and(seg_boundary, fg_dil).sum()
+    n_fg = fg_boundary.sum()
+    n_pred = seg_boundary.sum()
+    if n_fg == 0 and n_pred == 0:
+        return 1.0
+    if n_fg == 0 or n_pred == 0:
+        return 0.0
+    precision = pred_match / n_pred
+    recall = gt_match / n_fg
+    if precision + recall == 0:
+        return 0.0
+    return float(2 * precision * recall / (precision + recall))
+
+
+def mask_to_boundary_cv2(mask: np.ndarray, dilation: int) -> np.ndarray:
+    kernel = np.ones((2 * dilation + 1, 2 * dilation + 1), dtype=np.uint8)
+    eroded = cv2.erode(mask.astype(np.uint8), kernel, iterations=1).astype(bool)
+    return np.logical_xor(mask.astype(bool), eroded)
+
+
+def mask_to_boundary(mask: np.ndarray, dilation: int) -> np.ndarray:
+    eroded = binary_erode(mask, dilation)
+    return np.logical_xor(mask, eroded)
+
+
+def binary_erode(mask: np.ndarray, radius: int) -> np.ndarray:
+    padded = np.pad(mask.astype(bool), radius, mode="constant", constant_values=False)
+    out = np.ones_like(mask, dtype=bool)
+    size = 2 * radius + 1
+    for dy in range(size):
+        for dx in range(size):
+            out &= padded[dy : dy + mask.shape[0], dx : dx + mask.shape[1]]
+    return out
+
+
+def binary_dilate(mask: np.ndarray, radius: int) -> np.ndarray:
+    padded = np.pad(mask.astype(bool), radius, mode="constant", constant_values=False)
+    out = np.zeros_like(mask, dtype=bool)
+    size = 2 * radius + 1
+    for dy in range(size):
+        for dx in range(size):
+            out |= padded[dy : dy + mask.shape[0], dx : dx + mask.shape[1]]
+    return out
+
+
+def evaluate_tube_jf(pred_tube: np.ndarray, gt_tube: np.ndarray) -> Tuple[float, float, float]:
+    js = [db_eval_iou(gt_tube[t], pred_tube[t]) for t in range(gt_tube.shape[0])]
+    fs = [db_eval_boundary(gt_tube[t], pred_tube[t]) for t in range(gt_tube.shape[0])]
+    j = float(np.mean(js)) if js else 0.0
+    f = float(np.mean(fs)) if fs else 0.0
+    return j, f, (j + f) / 2
+
+
+def bbox_from_mask(mask: np.ndarray) -> Optional[Tuple[int, int, int, int]]:
+    ys, xs = np.where(mask.astype(bool))
+    if len(xs) == 0:
+        return None
+    return int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())
+
+
+def load_audit_rows(audit_csv: Path) -> Dict[str, dict]:
+    if not audit_csv.exists():
+        return {}
+    with audit_csv.open("r", newline="") as f:
+        return {row["uid"]: row for row in csv.DictReader(f)}
+
+
+def bool_field(row: Optional[dict], key: str) -> bool:
+    if not row:
+        return False
+    return str(row.get(key, "")).lower() in {"1", "true", "yes"}
+
+
+def write_json(path: Path, obj: dict) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w") as f:
+        json.dump(obj, f, indent=2, sort_keys=True)