#!/usr/bin/env python3 """ Load trace_profile.json files into PostgreSQL trace_profiles table. Usage: python3 scripts/load_trace_profiles.py [--output-dir OUTPUT_DIR] [--dry-run] """ import argparse import json import os import sys from pathlib import Path import psycopg2 from psycopg2.extras import execute_values def get_db_connection(): """Get PostgreSQL connection.""" database_url = os.environ.get( "DATABASE_URL", "postgresql://accusys@localhost:5432/momentry" ) return psycopg2.connect(database_url) def parse_trace_profile(json_path: Path) -> dict | None: """Parse trace_profile.json and extract relevant fields.""" try: with open(json_path, "r", encoding="utf-8") as f: data = json.load(f) except Exception as e: print(f" [ERROR] Failed to read {json_path}: {e}") return None # Extract file_uuid and trace_id from path parts = json_path.parts file_uuid = None trace_id = None for i, part in enumerate(parts): if part.startswith("trace_"): trace_id = int(part.replace("trace_", "")) if i > 0: file_uuid = parts[i - 1] break if not file_uuid or trace_id is None: print(f" [WARN] Cannot extract file_uuid/trace_id from {json_path}") return None return { "file_uuid": file_uuid, "trace_id": trace_id, "name": data.get("label") or data.get("name"), "start_frame": data.get("start_frame"), "end_frame": data.get("end_frame"), "frame_count": data.get("frame_count"), "key_frame": data.get("key_frame"), "status": data.get("status", "pending"), "avg_confidence": data.get("avg_confidence"), "vlm_description": data.get("vlm_description"), "vlm_clothing": data.get("vlm_clothing"), "vlm_tags": data.get("vlm_tags", []), "vlm_location": data.get("vlm_location"), "vlm_setting": data.get("vlm_setting"), "vlm_lighting": data.get("vlm_lighting"), "vlm_weather": data.get("vlm_weather"), "vlm_hand_objects": data.get("vlm_hand_objects"), "vlm_has_plants": data.get("vlm_has_plants"), "vlm_has_animals": data.get("vlm_has_animals"), "vlm_background": data.get("vlm_background"), "vlm_bg_tags": data.get("vlm_bg_tags", []), "vlm_model": data.get("vlm_model"), } def insert_trace_profiles(conn, profiles: list[dict], dry_run: bool = False): """Insert trace profiles into PostgreSQL.""" if not profiles: print("No profiles to insert") return 0 schema = os.environ.get("DATABASE_SCHEMA", "public") table = f"{schema}.trace_profiles" if schema != "public" else "trace_profiles" sql = f""" INSERT INTO {table} ( file_uuid, trace_id, name, start_frame, end_frame, frame_count, key_frame, status, avg_confidence, vlm_description, vlm_clothing, vlm_tags, vlm_location, vlm_setting, vlm_lighting, vlm_weather, vlm_hand_objects, vlm_has_plants, vlm_has_animals, vlm_background, vlm_bg_tags, vlm_model ) VALUES %s ON CONFLICT (file_uuid, trace_id) DO UPDATE SET name = EXCLUDED.name, start_frame = EXCLUDED.start_frame, end_frame = EXCLUDED.end_frame, frame_count = EXCLUDED.frame_count, key_frame = EXCLUDED.key_frame, status = EXCLUDED.status, avg_confidence = EXCLUDED.avg_confidence, vlm_description = EXCLUDED.vlm_description, vlm_clothing = EXCLUDED.vlm_clothing, vlm_tags = EXCLUDED.vlm_tags, vlm_location = EXCLUDED.vlm_location, vlm_setting = EXCLUDED.vlm_setting, vlm_lighting = EXCLUDED.vlm_lighting, vlm_weather = EXCLUDED.vlm_weather, vlm_hand_objects = EXCLUDED.vlm_hand_objects, vlm_has_plants = EXCLUDED.vlm_has_plants, vlm_has_animals = EXCLUDED.vlm_has_animals, vlm_background = EXCLUDED.vlm_background, vlm_bg_tags = EXCLUDED.vlm_bg_tags, vlm_model = EXCLUDED.vlm_model, updated_at = NOW() """ if dry_run: print(f"[DRY-RUN] Would insert {len(profiles)} profiles") return len(profiles) cursor = conn.cursor() # Prepare values values = [ ( p["file_uuid"], p["trace_id"], p["name"], p["start_frame"], p["end_frame"], p["frame_count"], p["key_frame"], p["status"], p["avg_confidence"], p["vlm_description"], p["vlm_clothing"], p["vlm_tags"], p["vlm_location"], p["vlm_setting"], p["vlm_lighting"], p["vlm_weather"], p["vlm_hand_objects"], p["vlm_has_plants"], p["vlm_has_animals"], p["vlm_background"], p["vlm_bg_tags"], p["vlm_model"], ) for p in profiles ] execute_values(cursor, sql, values) conn.commit() cursor.close() return len(profiles) def main(): parser = argparse.ArgumentParser(description="Load trace profiles into PostgreSQL") parser.add_argument( "--output-dir", default="/Users/accusys/momentry/output", help="Output directory containing trace_profile.json files", ) parser.add_argument( "--dry-run", action="store_true", help="Don't actually insert, just show what would be done", ) parser.add_argument( "--batch-size", type=int, default=100, help="Batch size for inserts", ) args = parser.parse_args() output_dir = Path(args.output_dir) if not output_dir.exists(): print(f"Error: Output directory {output_dir} does not exist") sys.exit(1) # Find all trace_profile.json files print(f"Scanning {output_dir} for trace_profile.json files...") trace_profiles = list(output_dir.glob("*/trace_*/trace_profile.json")) print(f"Found {len(trace_profiles)} trace_profile.json files") if not trace_profiles: print("No trace_profile.json files found") sys.exit(0) # Parse profiles print("\nParsing trace profiles...") profiles = [] for i, json_path in enumerate(trace_profiles): if (i + 1) % 500 == 0: print(f" Parsed {i + 1}/{len(trace_profiles)} files...") profile = parse_trace_profile(json_path) if profile: profiles.append(profile) print(f"Successfully parsed {len(profiles)} profiles") # Filter profiles with VLM data vlm_profiles = [ p for p in profiles if p.get("vlm_description") or p.get("vlm_clothing") or p.get("vlm_tags") ] print(f"Profiles with VLM data: {len(vlm_profiles)}") # Insert into PostgreSQL if not args.dry_run: print("\nConnecting to PostgreSQL...") conn = get_db_connection() else: conn = None print("\n[DRY-RUN] Skipping database connection") # Insert in batches batch_size = args.batch_size total_inserted = 0 for i in range(0, len(profiles), batch_size): batch = profiles[i : i + batch_size] if conn: inserted = insert_trace_profiles(conn, batch, args.dry_run) total_inserted += inserted if (i // batch_size + 1) % 10 == 0: print( f" Inserted batch {i // batch_size + 1} ({len(batch)} profiles)" ) else: total_inserted += len(batch) if conn: conn.close() print(f"\n✅ Done! Inserted {total_inserted} trace profiles") # Show sample if vlm_profiles: print("\nSample VLM profile:") sample = vlm_profiles[0] print(f" file_uuid: {sample['file_uuid']}") print(f" trace_id: {sample['trace_id']}") print(f" name: {sample['name']}") print(f" vlm_description: {sample.get('vlm_description', '')[:100]}...") print(f" vlm_tags: {sample.get('vlm_tags', [])[:5]}") if __name__ == "__main__": main()