397 lines
14 KiB
Python
397 lines
14 KiB
Python
#!/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() |