Accusys
39a2cbc65b
- get_face_groups_handler: COALESCE(tp.name, tn.label) for name consistency - sync_file_status: compare JSON vs pre_chunks (not chunk table) - face consistency: compare frames.len() not total_faces - cleanup 2 ghost records with NULL file_name/file_path - replace identity_agent with face_dedup in pipeline stages - remove identity_agent_api.rs and all references - update required_processors to match actual processors - update AGENTS.md with team responsibilities - add Studio pipeline changes documentation
257 lines
8.1 KiB
Python
257 lines
8.1 KiB
Python
#!/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() |