feat: add TKG node marking for Identity Agent suggestions
TKG Helper (scripts/utils/tkg_helper.py): - mark_face_track_suggested(): Mark node as 'suggested' with pending identity info - mark_face_track_confirmed(): Mark node as 'confirmed' with identity_ref - mark_face_track_stranger(): Mark node as 'stranger' with stranger_ref - batch_mark_suggestions(): Batch mark multiple traces - batch_mark_strangers(): Batch mark stranger clusters - get_face_track_nodes(): Get all face_track nodes for a file - get_pending_face_tracks(): Get nodes with status='pending' - get_suggested_face_tracks(): Get nodes with status='suggested' Identity Matcher updates: - Add --mark-tkg flag to update TKG nodes after matching - Integrates with tkg_helper for batch operations Node properties schema: - status: pending | suggested | confirmed | stranger - pending_identity_name/uuid/id: suggested identity info - suggested_by: tmdb | propagation | manual - confidence: matching score - identity_ref: confirmed identity reference
This commit is contained in:
Accusys
421
scripts/utils/tkg_helper.py
Normal file
421
scripts/utils/tkg_helper.py
Normal file
@@ -0,0 +1,421 @@
|
||||
#!/opt/homebrew/bin/python3.11
|
||||
"""
|
||||
TKG Helper - PostgreSQL TKG node operations for Identity Agent
|
||||
|
||||
Functions:
|
||||
- mark_face_track_suggested(): Mark face_track node as 'suggested'
|
||||
- mark_face_track_confirmed(): Mark face_track node as 'confirmed'
|
||||
- mark_face_track_stranger(): Mark face_track node as 'stranger'
|
||||
- get_face_track_nodes(): Get all face_track nodes for a file
|
||||
- get_pending_face_tracks(): Get face_track nodes with status='pending'
|
||||
|
||||
TKG face_track node properties schema:
|
||||
{
|
||||
"trace_id": int,
|
||||
"frame_count": int,
|
||||
"start_frame": int,
|
||||
"end_frame": int,
|
||||
"avg_bbox": {...},
|
||||
"avg_pose": {...},
|
||||
|
||||
// Identity binding states
|
||||
"status": "pending" | "suggested" | "confirmed" | "stranger",
|
||||
"pending_identity_name": str | null,
|
||||
"pending_identity_uuid": str | null,
|
||||
"suggested_by": "tmdb" | "propagation" | "manual" | null,
|
||||
"confidence": float,
|
||||
|
||||
// Confirmed fields
|
||||
"identity_uuid": str | null,
|
||||
"identity_ref": str | null,
|
||||
"stranger_ref": str | null
|
||||
}
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import psycopg2
|
||||
import psycopg2.extras
|
||||
from typing import Optional, Dict, List
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
DB_URL = os.environ.get("DATABASE_URL", "postgresql://accusys@localhost:5432/momentry")
|
||||
SCHEMA = os.environ.get("DATABASE_SCHEMA", "dev")
|
||||
|
||||
|
||||
def get_conn():
|
||||
"""Get PostgreSQL connection"""
|
||||
return psycopg2.connect(DB_URL)
|
||||
|
||||
|
||||
def table_name(table: str) -> str:
|
||||
"""Get schema-prefixed table name"""
|
||||
if SCHEMA == "public":
|
||||
return table
|
||||
return f"{SCHEMA}.{table}"
|
||||
|
||||
|
||||
def mark_face_track_suggested(
|
||||
file_uuid: str,
|
||||
trace_id: int,
|
||||
identity_id: int,
|
||||
identity_uuid: str,
|
||||
name: str,
|
||||
confidence: float,
|
||||
suggested_by: str = "tmdb",
|
||||
) -> bool:
|
||||
"""Mark face_track node as 'suggested'
|
||||
|
||||
Args:
|
||||
file_uuid: Video file UUID
|
||||
trace_id: Face trace ID
|
||||
identity_id: PG identity.id
|
||||
identity_uuid: Identity UUID
|
||||
name: Identity name
|
||||
confidence: Matching confidence score
|
||||
suggested_by: 'tmdb' | 'propagation' | 'manual'
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
conn = get_conn()
|
||||
cur = conn.cursor()
|
||||
|
||||
tkg_table = table_name("tkg_nodes")
|
||||
external_id = f"face_track_{trace_id}"
|
||||
|
||||
props = {
|
||||
"status": "suggested",
|
||||
"pending_identity_name": name,
|
||||
"pending_identity_uuid": identity_uuid,
|
||||
"pending_identity_id": identity_id,
|
||||
"suggested_by": suggested_by,
|
||||
"confidence": round(confidence, 4),
|
||||
}
|
||||
|
||||
try:
|
||||
cur.execute(
|
||||
f"""
|
||||
UPDATE {tkg_table}
|
||||
SET properties = properties || %s::jsonb
|
||||
WHERE file_uuid = %s AND node_type = 'face_track' AND external_id = %s
|
||||
""",
|
||||
(json.dumps(props), file_uuid, external_id),
|
||||
)
|
||||
conn.commit()
|
||||
updated = cur.rowcount > 0
|
||||
if updated:
|
||||
print(f"[TKG] Marked trace {trace_id} as suggested: {name} (confidence={confidence:.4f})")
|
||||
return updated
|
||||
except Exception as e:
|
||||
print(f"[TKG] Error marking trace {trace_id}: {e}")
|
||||
conn.rollback()
|
||||
return False
|
||||
finally:
|
||||
cur.close()
|
||||
conn.close()
|
||||
|
||||
|
||||
def mark_face_track_confirmed(
|
||||
file_uuid: str,
|
||||
trace_id: int,
|
||||
identity_id: int,
|
||||
identity_uuid: str,
|
||||
name: str,
|
||||
) -> bool:
|
||||
"""Mark face_track node as 'confirmed'
|
||||
|
||||
Args:
|
||||
file_uuid: Video file UUID
|
||||
trace_id: Face trace ID
|
||||
identity_id: PG identity.id
|
||||
identity_uuid: Identity UUID
|
||||
name: Identity name
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
conn = get_conn()
|
||||
cur = conn.cursor()
|
||||
|
||||
tkg_table = table_name("tkg_nodes")
|
||||
external_id = f"face_track_{trace_id}"
|
||||
identity_ref = f"{file_uuid}:identity_{identity_id}"
|
||||
|
||||
props = {
|
||||
"status": "confirmed",
|
||||
"identity_uuid": identity_uuid,
|
||||
"identity_id": identity_id,
|
||||
"identity_ref": identity_ref,
|
||||
"identity_name": name,
|
||||
}
|
||||
|
||||
# Remove pending fields
|
||||
remove_keys = ["pending_identity_name", "pending_identity_uuid", "pending_identity_id", "suggested_by", "confidence"]
|
||||
|
||||
try:
|
||||
# Build JSONB update: add new props, remove pending fields
|
||||
props_json = json.dumps(props)
|
||||
cur.execute(
|
||||
f"""
|
||||
UPDATE {tkg_table}
|
||||
SET properties = (properties || %s::jsonb)
|
||||
- 'pending_identity_name' - 'pending_identity_uuid'
|
||||
- 'pending_identity_id' - 'suggested_by' - 'confidence' - 'stranger_ref'
|
||||
WHERE file_uuid = %s AND node_type = 'face_track' AND external_id = %s
|
||||
""",
|
||||
(props_json, file_uuid, external_id),
|
||||
)
|
||||
conn.commit()
|
||||
updated = cur.rowcount > 0
|
||||
if updated:
|
||||
print(f"[TKG] Marked trace {trace_id} as confirmed: {name}")
|
||||
return updated
|
||||
except Exception as e:
|
||||
print(f"[TKG] Error confirming trace {trace_id}: {e}")
|
||||
conn.rollback()
|
||||
return False
|
||||
finally:
|
||||
cur.close()
|
||||
conn.close()
|
||||
|
||||
|
||||
def mark_face_track_stranger(
|
||||
file_uuid: str,
|
||||
trace_id: int,
|
||||
stranger_cluster_id: int,
|
||||
) -> bool:
|
||||
"""Mark face_track node as 'stranger'
|
||||
|
||||
Args:
|
||||
file_uuid: Video file UUID
|
||||
trace_id: Face trace ID
|
||||
stranger_cluster_id: Stranger cluster ID
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
conn = get_conn()
|
||||
cur = conn.cursor()
|
||||
|
||||
tkg_table = table_name("tkg_nodes")
|
||||
external_id = f"face_track_{trace_id}"
|
||||
stranger_ref = f"stranger_{stranger_cluster_id}"
|
||||
|
||||
props = {
|
||||
"status": "stranger",
|
||||
"stranger_id": stranger_cluster_id,
|
||||
"stranger_ref": stranger_ref,
|
||||
}
|
||||
|
||||
try:
|
||||
cur.execute(
|
||||
f"""
|
||||
UPDATE {tkg_table}
|
||||
SET properties = (properties || %s::jsonb)
|
||||
- 'pending_identity_name' - 'pending_identity_uuid'
|
||||
- 'pending_identity_id' - 'suggested_by' - 'confidence'
|
||||
- 'identity_uuid' - 'identity_ref'
|
||||
WHERE file_uuid = %s AND node_type = 'face_track' AND external_id = %s
|
||||
""",
|
||||
(json.dumps(props), file_uuid, external_id),
|
||||
)
|
||||
conn.commit()
|
||||
updated = cur.rowcount > 0
|
||||
if updated:
|
||||
print(f"[TKG] Marked trace {trace_id} as stranger cluster {stranger_cluster_id}")
|
||||
return updated
|
||||
except Exception as e:
|
||||
print(f"[TKG] Error marking stranger trace {trace_id}: {e}")
|
||||
conn.rollback()
|
||||
return False
|
||||
finally:
|
||||
cur.close()
|
||||
conn.close()
|
||||
|
||||
|
||||
def get_face_track_nodes(file_uuid: str) -> List[Dict]:
|
||||
"""Get all face_track nodes for a file
|
||||
|
||||
Args:
|
||||
file_uuid: Video file UUID
|
||||
|
||||
Returns:
|
||||
List of face_track nodes with properties
|
||||
"""
|
||||
conn = get_conn()
|
||||
cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
|
||||
|
||||
tkg_table = table_name("tkg_nodes")
|
||||
|
||||
try:
|
||||
cur.execute(
|
||||
f"""
|
||||
SELECT id, external_id, label, properties, created_at
|
||||
FROM {tkg_table}
|
||||
WHERE file_uuid = %s AND node_type = 'face_track'
|
||||
ORDER BY external_id
|
||||
""",
|
||||
(file_uuid,),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
finally:
|
||||
cur.close()
|
||||
conn.close()
|
||||
|
||||
|
||||
def get_pending_face_tracks(file_uuid: str) -> List[Dict]:
|
||||
"""Get face_track nodes with status='pending' or NULL status
|
||||
|
||||
Args:
|
||||
file_uuid: Video file UUID
|
||||
|
||||
Returns:
|
||||
List of pending face_track nodes
|
||||
"""
|
||||
conn = get_conn()
|
||||
cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
|
||||
|
||||
tkg_table = table_name("tkg_nodes")
|
||||
|
||||
try:
|
||||
cur.execute(
|
||||
f"""
|
||||
SELECT id, external_id, label, properties, created_at
|
||||
FROM {tkg_table}
|
||||
WHERE file_uuid = %s AND node_type = 'face_track'
|
||||
AND (properties->>'status' IS NULL OR properties->>'status' = 'pending')
|
||||
ORDER BY external_id
|
||||
""",
|
||||
(file_uuid,),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
finally:
|
||||
cur.close()
|
||||
conn.close()
|
||||
|
||||
|
||||
def get_suggested_face_tracks(file_uuid: str) -> List[Dict]:
|
||||
"""Get face_track nodes with status='suggested'
|
||||
|
||||
Args:
|
||||
file_uuid: Video file UUID
|
||||
|
||||
Returns:
|
||||
List of suggested face_track nodes
|
||||
"""
|
||||
conn = get_conn()
|
||||
cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
|
||||
|
||||
tkg_table = table_name("tkg_nodes")
|
||||
|
||||
try:
|
||||
cur.execute(
|
||||
f"""
|
||||
SELECT id, external_id, label, properties, created_at
|
||||
FROM {tkg_table}
|
||||
WHERE file_uuid = %s AND node_type = 'face_track'
|
||||
AND properties->>'status' = 'suggested'
|
||||
ORDER BY external_id
|
||||
""",
|
||||
(file_uuid,),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
finally:
|
||||
cur.close()
|
||||
conn.close()
|
||||
|
||||
|
||||
def clear_face_track_status(file_uuid: str, trace_id: int) -> bool:
|
||||
"""Clear identity binding status from face_track node
|
||||
|
||||
Args:
|
||||
file_uuid: Video file UUID
|
||||
trace_id: Face trace ID
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
conn = get_conn()
|
||||
cur = conn.cursor()
|
||||
|
||||
tkg_table = table_name("tkg_nodes")
|
||||
external_id = f"face_track_{trace_id}"
|
||||
|
||||
try:
|
||||
cur.execute(
|
||||
f"""
|
||||
UPDATE {tkg_table}
|
||||
SET properties = properties
|
||||
- 'status' - 'pending_identity_name' - 'pending_identity_uuid'
|
||||
- 'pending_identity_id' - 'suggested_by' - 'confidence'
|
||||
- 'identity_uuid' - 'identity_ref' - 'identity_id' - 'identity_name'
|
||||
- 'stranger_id' - 'stranger_ref'
|
||||
WHERE file_uuid = %s AND node_type = 'face_track' AND external_id = %s
|
||||
""",
|
||||
(file_uuid, external_id),
|
||||
)
|
||||
conn.commit()
|
||||
return cur.rowcount > 0
|
||||
except Exception as e:
|
||||
print(f"[TKG] Error clearing trace {trace_id}: {e}")
|
||||
conn.rollback()
|
||||
return False
|
||||
finally:
|
||||
cur.close()
|
||||
conn.close()
|
||||
|
||||
|
||||
def batch_mark_suggestions(file_uuid: str, suggestions: Dict) -> int:
|
||||
"""Batch mark multiple face_track nodes as 'suggested'
|
||||
|
||||
Args:
|
||||
file_uuid: Video file UUID
|
||||
suggestions: {trace_id: {identity_id, identity_uuid, name, score, suggested_by}}
|
||||
|
||||
Returns:
|
||||
Number of nodes updated
|
||||
"""
|
||||
updated = 0
|
||||
for trace_id_str, suggestion in suggestions.items():
|
||||
trace_id = int(trace_id_str)
|
||||
success = mark_face_track_suggested(
|
||||
file_uuid,
|
||||
trace_id,
|
||||
suggestion.get("identity_id"),
|
||||
suggestion.get("identity_uuid"),
|
||||
suggestion.get("name"),
|
||||
suggestion.get("score", 0.0),
|
||||
suggestion.get("suggested_by", "tmdb"),
|
||||
)
|
||||
if success:
|
||||
updated += 1
|
||||
|
||||
print(f"[TKG] Batch marked {updated}/{len(suggestions)} traces as suggested")
|
||||
return updated
|
||||
|
||||
|
||||
def batch_mark_strangers(file_uuid: str, stranger_clusters: Dict) -> int:
|
||||
"""Batch mark multiple face_track nodes as 'stranger'
|
||||
|
||||
Args:
|
||||
file_uuid: Video file UUID
|
||||
stranger_clusters: {cluster_id: [trace_ids]}
|
||||
|
||||
Returns:
|
||||
Number of nodes updated
|
||||
"""
|
||||
updated = 0
|
||||
for cluster_id, trace_ids in stranger_clusters.items():
|
||||
for trace_id in trace_ids:
|
||||
success = mark_face_track_stranger(file_uuid, trace_id, cluster_id)
|
||||
if success:
|
||||
updated += 1
|
||||
|
||||
print(f"[TKG] Batch marked {updated} traces as strangers in {len(stranger_clusters)} clusters")
|
||||
return updated
|
||||
Reference in New Issue
Block a user