teach_csv

  1#!/usr/bin/env python3
  2"""Batch-teach recording CSVs into Atlas (no robot needed).
  3
  4Teaches labeled recordings collected earlier, e.g. after the P0 run:
  5
  6    python teach_csv.py recordings/zzzold/recording_20261029_153038_water_half.csv
  7    python teach_csv.py recordings/ --session thu
  8    python teach_csv.py recordings/ --dry-run     # show what would be taught
  9    python teach_csv.py --list                    # signature counts per label
 10
 11Label resolution: --label override > object_label column > filename.
 12Files labeled mystery/unlabeled are skipped unless --include-mystery.
 13Files already taught (matched by properties.csv) are skipped unless --force.
 14"""
 15
 16import argparse
 17import glob
 18import os
 19import sys
 20
 21import pandas as pd
 22
 23import atlas
 24import feature_extractor
 25
 26SKIP_LABELS = {"mystery", "unlabeled", "unknown"}
 27MOTION_TYPE = "swirl_macro"
 28
 29
 30def resolve_label(csv_path, df, override=None):
 31    """Determine the teach label for a CSV.
 32
 33    Precedence: explicit override, then the object_label column, then the
 34    recording_YYYYMMDD_HHMMSS_<label>.csv filename.
 35
 36    Args:
 37        csv_path: path to the recording CSV.
 38        df: the recording DataFrame.
 39        override: forced label, or None.
 40
 41    Returns:
 42        Label string, or None when none can be determined.
 43    """
 44    if override:
 45        return override
 46    if "object_label" in df.columns:
 47        lbl = str(df["object_label"].iloc[0]).strip()
 48        if lbl and lbl.lower() != "nan":
 49            return lbl
 50    base = os.path.basename(csv_path)
 51    # recording_YYYYMMDD_HHMMSS_<label>.csv
 52    parts = base.replace(".csv", "").split("_")
 53    if len(parts) >= 4 and parts[0] == "recording":
 54        return "_".join(parts[3:])
 55    return None
 56
 57
 58def teach_file(db, csv_path, session, label_override=None, dry_run=False,
 59               force=False, include_mystery=False, require_swirl=False):
 60    """Extract features from one CSV and teach it to Atlas.
 61
 62    Args:
 63        db: MongoDB database handle.
 64        csv_path: path to the recording CSV.
 65        session: session tag, e.g. "thu".
 66        label_override: forced label, or None.
 67        dry_run: report what would be taught without writing.
 68        force: re-teach files already in Atlas.
 69        include_mystery: also teach mystery-labeled files.
 70        require_swirl: skip recordings without a true SWIRL phase.
 71
 72    Returns:
 73        Tuple (status, filename, detail) where status is one of
 74        "taught", "dry", "skip", "error".
 75    """
 76    name = os.path.basename(csv_path)
 77    try:
 78        df = pd.read_csv(csv_path)
 79    except Exception as e:
 80        return ("error", name, f"unreadable: {e}")
 81    if not {"phase", "tau2_nm", "tau3_nm"}.issubset(df.columns):
 82        return ("skip", name, "not a grasp recording")
 83    label = resolve_label(csv_path, df, label_override)
 84    if not label:
 85        return ("skip", name, "no label found")
 86    if label.lower() in SKIP_LABELS and not include_mystery:
 87        return ("skip", name, f"label '{label}' (use --include-mystery to teach it)")
 88    try:
 89        feats = feature_extractor.extract_features(df)
 90    except Exception as e:
 91        return ("error", name, f"extraction failed: {e}")
 92    if dry_run:
 93        dq = feats["diagnostics"]
 94        proto = "swirl" if dq["phase_source"] == "SWIRL" else "fallback"
 95        if require_swirl and dq["phase_source"] != "SWIRL":
 96            return ("skip", name, "no SWIRL phase (fallback protocol)")
 97        return ("dry", name,
 98                f"would teach '{label}' [{proto}] "
 99                f"({dq['n_samples']} samples, tare={dq['tare_flags'] or 'OK'})")
100    if require_swirl and feats["diagnostics"]["phase_source"] != "SWIRL":
101        return ("skip", name, "no SWIRL phase (fallback protocol)")
102    if not force and db.signatures.count_documents({"properties.csv": name}, limit=1):
103        return ("skip", name, "already in Atlas (use --force to re-teach)")
104    dq = feats["diagnostics"]
105    saved = atlas.teach(
106        label, feats["e"], motion_type=MOTION_TYPE, session=session,
107        properties={"csv": name, "source": "teach_csv.py",
108                    "tare_flags": dq["tare_flags"],
109                    "phase_source": dq["phase_source"],
110                    "z": [round(float(v), 4) for v in feats["z"]],
111                    "n_samples": dq["n_samples"], "fs_hz": dq["fs_hz"]})
112    return ("taught", name, f"'{label}' sig_id={saved['sig_id']}")
113
114
115def list_library():
116    """Print signature counts per label in the Atlas library."""
117    db = atlas.get_db()
118    rows = list(db.signatures.aggregate([
119        {"$group": {"_id": "$label", "n": {"$sum": 1},
120                    "motion_types": {"$addToSet": "$motion_type"},
121                    "sessions": {"$addToSet": "$session"}}},
122        {"$sort": {"_id": 1}}]))
123    if not rows:
124        print("signatures collection is empty")
125        return
126    print(f"{'label':<24} {'n':>3}  motion_types  sessions")
127    for r in rows:
128        print(f"{r['_id']:<24} {r['n']:>3}  "
129              f"{','.join(sorted(r['motion_types']))}  {','.join(sorted(r['sessions']))}")
130
131
132def main():
133    """CLI entry: batch-teach recordings to Atlas."""
134    ap = argparse.ArgumentParser(
135        description="Batch-teach recording CSVs into Atlas.",
136        epilog="""examples:
137  python teach_csv.py recordings/ --require-swirl --dry-run   # preview
138  python teach_csv.py recordings/ --require-swirl --session fri
139  python teach_csv.py --list                                  # library contents
140  python teach_csv.py foo.csv --label water_075pct --force    # re-teach one file
141
142Label resolution: --label > object_label column > filename.
143Mystery/unlabeled and non-grasp CSVs are skipped; already-taught files are
144skipped unless --force. Needs MONGODB_URI in .env, and the Atlas IP
145allowlist must include this machine.""",
146        formatter_class=argparse.RawDescriptionHelpFormatter)
147    ap.add_argument("paths", nargs="*", help="CSV files, globs, or dirs (default: recordings/)")
148    ap.add_argument("--session", default="thu", help="session tag (default: thu)")
149    ap.add_argument("--label", default=None, help="force one label for all files")
150    ap.add_argument("--dry-run", action="store_true", help="show what would be taught")
151    ap.add_argument("--force", action="store_true", help="re-teach files already in Atlas")
152    ap.add_argument("--include-mystery", action="store_true")
153    ap.add_argument("--require-swirl", action="store_true",
154                    help="only teach recordings with a true SWIRL phase")
155    ap.add_argument("--list", action="store_true", help="show signature counts per label")
156    args = ap.parse_args()
157
158    if args.list:
159        list_library()
160        return
161
162    files = []
163    for p in args.paths or ["recordings/"]:
164        if os.path.isdir(p):
165            files += sorted(glob.glob(os.path.join(p, "**", "*.csv"), recursive=True))
166        elif os.path.isfile(p):
167            files.append(p)
168        else:
169            files += sorted(glob.glob(p, recursive=True))
170    if not files:
171        sys.exit("no CSV files found")
172
173    db = None if args.dry_run else atlas.get_db()
174    counts = {}
175    for f in files:
176        status, name, detail = teach_file(
177            db, f, args.session, label_override=args.label,
178            dry_run=args.dry_run, force=args.force,
179            include_mystery=args.include_mystery,
180            require_swirl=args.require_swirl)
181        counts[status] = counts.get(status, 0) + 1
182        print(f"[{status:>7}] {name}: {detail}")
183    print("\nsummary:", ", ".join(f"{k}={v}" for k, v in sorted(counts.items())))
184
185
186if __name__ == "__main__":
187    main()