feat: ASRX hybrid pipeline, identity history, worker fixes, checkpoint system
This commit is contained in:
Accusys
397
scripts/extract_face_crops.py
Normal file
397
scripts/extract_face_crops.py
Normal file
@@ -0,0 +1,397 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
extract_face_crops.py - 批量提取 face crops
|
||||
|
||||
Usage:
|
||||
python3 scripts/extract_face_crops.py --uuid <file_uuid>
|
||||
python3 scripts/extract_face_crops.py --uuid <file_uuid> --video <video_path>
|
||||
|
||||
儲存位置: {OUTPUT_DIR}/.faces/{file_uuid}/{trace_id}/{frame}.jpg
|
||||
|
||||
條件:
|
||||
- trace_id != None and trace_id != 0
|
||||
- landmarks.left_eye or landmarks.right_eye
|
||||
|
||||
品檢:
|
||||
- file_size > 500 bytes
|
||||
- mean_brightness > 5
|
||||
- std_deviation > 10
|
||||
|
||||
Retry: 最多 3 次
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import subprocess
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Optional, Tuple, Set
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
import threading
|
||||
|
||||
# Constants
|
||||
MAX_RETRIES = 3
|
||||
MIN_FILE_SIZE = 500
|
||||
MIN_BRIGHTNESS = 5
|
||||
MIN_STD_DEV = 10
|
||||
FFMPEG_TIMEOUT = 30
|
||||
MAX_WORKERS = 8 # Parallel threads for ffmpeg
|
||||
|
||||
|
||||
class FaceCropExtractor:
|
||||
def __init__(self, output_dir: str):
|
||||
self.output_dir = Path(output_dir)
|
||||
self.faces_dir = self.output_dir / ".faces"
|
||||
self.faces_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.stats = {"total_faces": 0, "qualified": 0, "successful": 0, "failed": 0, "skipped": 0, "low_confidence": 0, "too_small": 0}
|
||||
self.stats_lock = threading.Lock()
|
||||
|
||||
def process_video(self, uuid: str, video_path: str) -> dict:
|
||||
"""處理單一影片"""
|
||||
face_json = self.output_dir / f"{uuid}.face.json"
|
||||
traced_json = self.output_dir / f"{uuid}.face_traced.json"
|
||||
|
||||
if not face_json.exists():
|
||||
print(f"[ERROR] face.json not found: {uuid}")
|
||||
return {"error": "face.json not found"}
|
||||
|
||||
if not os.path.exists(video_path):
|
||||
print(f"[ERROR] Video not found: {video_path}")
|
||||
return {"error": "video not found"}
|
||||
|
||||
# Load face.json (landmarks)
|
||||
print(f"[LOAD] Reading {face_json}")
|
||||
with open(face_json) as f:
|
||||
face_data = json.load(f)
|
||||
|
||||
# Load face_traced.json if exists (trace_id)
|
||||
traced_data = {}
|
||||
if traced_json.exists():
|
||||
print(f"[LOAD] Reading {traced_json}")
|
||||
with open(traced_json) as f:
|
||||
traced_data = json.load(f)
|
||||
|
||||
# Build lookup: (frame, x, y) -> trace_id from traced_data
|
||||
trace_lookup: Dict[Tuple[int, int, int], int] = {}
|
||||
frames = traced_data.get("frames", {})
|
||||
if isinstance(frames, dict):
|
||||
for fnum, frm in frames.items():
|
||||
faces = frm.get("faces", [])
|
||||
if faces is None:
|
||||
continue
|
||||
for face in faces:
|
||||
if face is None:
|
||||
continue
|
||||
trace_id = face.get("trace_id")
|
||||
if trace_id and trace_id != 0:
|
||||
x = face.get("x", 0)
|
||||
y = face.get("y", 0)
|
||||
key = (int(fnum), x, y)
|
||||
trace_lookup[key] = trace_id
|
||||
|
||||
# Create output directory
|
||||
uuid_dir = self.faces_dir / uuid
|
||||
uuid_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
results = {"successful": [], "failed": []}
|
||||
processed: Set[Tuple[int, int]] = set() # (trace_id, frame)
|
||||
trace_counts: Dict[int, int] = {} # trace_id -> count
|
||||
|
||||
# Process faces from face.json
|
||||
frames = face_data.get("frames", {})
|
||||
if isinstance(frames, dict):
|
||||
frame_items = frames.items()
|
||||
elif isinstance(frames, list):
|
||||
frame_items = [(frm.get("frame"), frm) for frm in frames]
|
||||
else:
|
||||
frame_items = []
|
||||
|
||||
# Collect extraction tasks
|
||||
tasks = []
|
||||
for fnum, frm in frame_items:
|
||||
if fnum is None:
|
||||
continue
|
||||
|
||||
faces = frm.get("faces", [])
|
||||
if faces is None:
|
||||
continue
|
||||
|
||||
for face in faces:
|
||||
if face is None:
|
||||
continue
|
||||
|
||||
self.stats["total_faces"] += 1
|
||||
|
||||
bb = face.get("bbox", face)
|
||||
x = bb.get("x", 0) if isinstance(bb, dict) else 0
|
||||
y = bb.get("y", 0) if isinstance(bb, dict) else 0
|
||||
w = bb.get("width", 0) if isinstance(bb, dict) else 0
|
||||
h = bb.get("height", 0) if isinstance(bb, dict) else 0
|
||||
confidence = face.get("confidence", 0.0)
|
||||
|
||||
# Quality filtering: confidence + size
|
||||
MIN_CONFIDENCE = 0.6
|
||||
MIN_SIZE = 20
|
||||
|
||||
if confidence < MIN_CONFIDENCE:
|
||||
self.stats["low_confidence"] += 1
|
||||
continue
|
||||
if w < MIN_SIZE or h < MIN_SIZE:
|
||||
self.stats["too_small"] += 1
|
||||
continue
|
||||
|
||||
# Check: has eyes
|
||||
lm = face.get("landmarks")
|
||||
has_eyes = False
|
||||
if lm:
|
||||
if isinstance(lm, dict):
|
||||
has_eyes = lm.get("left_eye") or lm.get("right_eye")
|
||||
elif isinstance(lm, list) and len(lm) >= 2:
|
||||
has_eyes = True
|
||||
|
||||
if not has_eyes:
|
||||
self.stats["skipped"] += 1
|
||||
continue
|
||||
|
||||
self.stats["qualified"] += 1
|
||||
|
||||
# Check: in trace
|
||||
key = (int(fnum), x, y)
|
||||
trace_id = trace_lookup.get(key) or face.get("trace_id")
|
||||
|
||||
# Determine output directory
|
||||
if trace_id and trace_id != 0:
|
||||
output_dir = uuid_dir / str(trace_id)
|
||||
crop_key = (trace_id, int(fnum))
|
||||
else:
|
||||
# No trace_id → unbound directory
|
||||
output_dir = uuid_dir / "unbound"
|
||||
crop_key = ("unbound", int(fnum), x, y) # unique key for unbound
|
||||
|
||||
if crop_key in processed:
|
||||
continue
|
||||
processed.add(crop_key)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
output_path = output_dir / f"{fnum}.jpg"
|
||||
|
||||
tasks.append({
|
||||
"trace_id": trace_id or "unbound",
|
||||
"frame": int(fnum),
|
||||
"x": x, "y": y, "w": w, "h": h,
|
||||
"output_path": output_path
|
||||
})
|
||||
|
||||
# Parallel extraction
|
||||
print(f"[EXTRACT] Processing {len(tasks)} faces with {MAX_WORKERS} threads...")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
||||
futures = {
|
||||
executor.submit(
|
||||
self.extract_face, video_path, t["frame"],
|
||||
t["x"], t["y"], t["w"], t["h"], t["output_path"]
|
||||
): t for t in tasks
|
||||
}
|
||||
|
||||
for i, future in enumerate(as_completed(futures)):
|
||||
t = futures[future]
|
||||
result = future.result()
|
||||
|
||||
if result["success"]:
|
||||
self.stats["successful"] += 1
|
||||
results["successful"].append({
|
||||
"trace_id": t["trace_id"],
|
||||
"frame": t["frame"],
|
||||
"path": str(t["output_path"])
|
||||
})
|
||||
trace_counts[t["trace_id"]] = trace_counts.get(t["trace_id"], 0) + 1
|
||||
else:
|
||||
self.stats["failed"] += 1
|
||||
results["failed"].append({
|
||||
"trace_id": t["trace_id"],
|
||||
"frame": t["frame"],
|
||||
"bbox": {"x": t["x"], "y": t["y"], "w": t["w"], "h": t["h"]},
|
||||
"reason": result.get("reason", "unknown")
|
||||
})
|
||||
|
||||
# Progress every 1000
|
||||
if (i + 1) % 1000 == 0:
|
||||
print(f" Progress: {i+1}/{len(tasks)} ({self.stats['successful']} OK, {self.stats['failed']} fail)")
|
||||
|
||||
# Write summary
|
||||
self.write_summary(uuid, trace_counts, results)
|
||||
|
||||
return results
|
||||
|
||||
def extract_face(self, video_path: str, frame: int, x: int, y: int,
|
||||
w: int, h: int, output_path: Path) -> dict:
|
||||
"""提取 face crop（含 retry，使用 -ss 快速 seek）"""
|
||||
for attempt in range(MAX_RETRIES):
|
||||
try:
|
||||
ts = frame / 24.0 # FPS is always 24 for this video
|
||||
cmd = [
|
||||
"ffmpeg", "-y", "-ss", f"{ts:.3f}",
|
||||
"-i", video_path,
|
||||
"-vf", f"crop={w}:{h}:{x}:{y}",
|
||||
"-frames:v", "1",
|
||||
"-q:v", "2", # 高品質 JPEG
|
||||
str(output_path)
|
||||
]
|
||||
|
||||
proc = subprocess.run(
|
||||
cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
||||
timeout=FFMPEG_TIMEOUT
|
||||
)
|
||||
|
||||
if proc.returncode != 0:
|
||||
if attempt < MAX_RETRIES - 1:
|
||||
continue
|
||||
return {"success": False, "reason": "ffmpeg_failed"}
|
||||
|
||||
# Quality check
|
||||
quality = self.check_quality(output_path)
|
||||
if quality["ok"]:
|
||||
return {"success": True, "path": str(output_path)}
|
||||
|
||||
# Quality failed, retry
|
||||
if attempt < MAX_RETRIES - 1:
|
||||
# Remove bad file
|
||||
if output_path.exists():
|
||||
output_path.unlink()
|
||||
continue
|
||||
|
||||
return {"success": False, "reason": quality.get("reason", "quality_failed")}
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
if attempt < MAX_RETRIES - 1:
|
||||
continue
|
||||
return {"success": False, "reason": "timeout"}
|
||||
except Exception as e:
|
||||
return {"success": False, "reason": str(e)}
|
||||
|
||||
return {"success": False, "reason": "max_retries"}
|
||||
|
||||
def check_quality(self, path: Path) -> dict:
|
||||
"""品檢"""
|
||||
if not path.exists():
|
||||
return {"ok": False, "reason": "file_not_exist"}
|
||||
|
||||
file_size = path.stat().st_size
|
||||
if file_size < MIN_FILE_SIZE:
|
||||
return {"ok": False, "reason": f"empty_file ({file_size}B)"}
|
||||
|
||||
try:
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
img = Image.open(path)
|
||||
arr = np.array(img.convert('RGB'))
|
||||
|
||||
mean_brightness = arr.mean()
|
||||
if mean_brightness < MIN_BRIGHTNESS:
|
||||
return {"ok": False, "reason": f"black_frame (mean={mean_brightness:.1f})"}
|
||||
|
||||
std_dev = arr.std()
|
||||
if std_dev < MIN_STD_DEV:
|
||||
return {"ok": False, "reason": f"low_contrast (std={std_dev:.1f})"}
|
||||
|
||||
return {"ok": True}
|
||||
|
||||
except ImportError:
|
||||
# PIL not available, skip advanced quality check
|
||||
return {"ok": True}
|
||||
except Exception as e:
|
||||
return {"ok": False, "reason": str(e)}
|
||||
|
||||
def write_summary(self, uuid: str, trace_counts: Dict[int, int], results: dict):
|
||||
"""寫摘要報告"""
|
||||
summary_path = self.faces_dir / uuid / "_summary.json"
|
||||
|
||||
summary = {
|
||||
"file_uuid": uuid,
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"stats": self.stats,
|
||||
"trace_counts": trace_counts,
|
||||
"total_traces": len(trace_counts),
|
||||
"failed_count": len(results["failed"]),
|
||||
"failed_faces": results["failed"] if results["failed"] else None
|
||||
}
|
||||
|
||||
with open(summary_path, "w") as f:
|
||||
json.dump(summary, f, indent=2, ensure_ascii=False)
|
||||
|
||||
print(f"\n[SUMMARY] Written to {summary_path}")
|
||||
|
||||
def print_stats(self):
|
||||
"""印統計"""
|
||||
print(f"\n=== Statistics ===")
|
||||
print(f"Total faces scanned: {self.stats['total_faces']}")
|
||||
print(f"Filtered (low confidence < 0.6): {self.stats['low_confidence']}")
|
||||
print(f"Filtered (too small < 20px): {self.stats['too_small']}")
|
||||
print(f"Qualified (trace_id + eyes): {self.stats['qualified']}")
|
||||
print(f"Successfully extracted: {self.stats['successful']}")
|
||||
print(f"Failed: {self.stats['failed']}")
|
||||
print(f"Skipped (no trace/eyes): {self.stats['skipped']}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Extract face crops from videos",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__
|
||||
)
|
||||
|
||||
parser.add_argument("--uuid", type=str, required=True,
|
||||
help="File UUID to process")
|
||||
parser.add_argument("--video", type=str,
|
||||
help="Video file path (optional, will check DB if not provided)")
|
||||
parser.add_argument("--output-dir", type=str,
|
||||
default="/Users/accusys/momentry/output_dev",
|
||||
help="Output directory (default: output_dev)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get video path
|
||||
video_path = args.video
|
||||
if not video_path:
|
||||
# Query from DB
|
||||
video_path = query_video_path_from_db(args.uuid)
|
||||
if not video_path:
|
||||
print(f"[ERROR] Video path not found for UUID: {args.uuid}")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"=== Face Crop Extraction ===")
|
||||
print(f"UUID: {args.uuid}")
|
||||
print(f"Video: {video_path}")
|
||||
print(f"Output: {args.output_dir}/.faces/{args.uuid}/")
|
||||
print()
|
||||
|
||||
extractor = FaceCropExtractor(args.output_dir)
|
||||
results = extractor.process_video(args.uuid, video_path)
|
||||
extractor.print_stats()
|
||||
|
||||
|
||||
def query_video_path_from_db(uuid: str) -> Optional[str]:
|
||||
"""從 PostgreSQL 取得影片路徑"""
|
||||
psql_path = "/opt/homebrew/Cellar/libpq/18.3/bin/psql"
|
||||
|
||||
if not os.path.exists(psql_path):
|
||||
return None
|
||||
|
||||
cmd = [
|
||||
psql_path, "-U", "accusys", "-d", "momentry", "-t", "-A",
|
||||
"-c", f"SELECT file_path FROM public.videos WHERE file_uuid = '{uuid}' LIMIT 1"
|
||||
]
|
||||
|
||||
try:
|
||||
proc = subprocess.run(cmd, capture_output=True, text=True, timeout=5)
|
||||
path = proc.stdout.strip()
|
||||
return path if path else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user