feat: update Python processors and add utility scripts
- Update ASR, face, OCR, pose processors - Add release pre-flight check script - Add synonym generation, chunk processing scripts - Add face recognition, stamp search utilities
This commit is contained in:
Warren
372
scripts/face_registration.py
Normal file
372
scripts/face_registration.py
Normal file
@@ -0,0 +1,372 @@
|
||||
#!/opt/homebrew/bin/python3.11
|
||||
"""
|
||||
Face Registration Script
|
||||
Register new faces to the face database
|
||||
"""
|
||||
|
||||
import sys
|
||||
import json
|
||||
import argparse
|
||||
import os
|
||||
import numpy as np
|
||||
import time
|
||||
from typing import Dict, Any, Optional
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
|
||||
class FaceRegistration:
|
||||
def __init__(self):
|
||||
self.face_model = None
|
||||
self.face_database = {}
|
||||
self.database_path = None
|
||||
|
||||
def load_models(self, use_mps: bool = False):
|
||||
"""Load InsightFace models with MPS support"""
|
||||
try:
|
||||
import insightface
|
||||
from insightface.app import FaceAnalysis
|
||||
|
||||
# Determine execution providers based on configuration
|
||||
providers = ["CPUExecutionProvider"]
|
||||
|
||||
if use_mps:
|
||||
try:
|
||||
# Try to import MPS provider
|
||||
import onnxruntime as ort
|
||||
|
||||
available_providers = ort.get_available_providers()
|
||||
|
||||
if "CoreMLExecutionProvider" in available_providers:
|
||||
print(
|
||||
"[INFO] Using CoreMLExecutionProvider for MPS acceleration"
|
||||
)
|
||||
providers = ["CoreMLExecutionProvider", "CPUExecutionProvider"]
|
||||
elif "CUDAExecutionProvider" in available_providers:
|
||||
print("[INFO] Using CUDAExecutionProvider")
|
||||
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
else:
|
||||
print("[INFO] MPS/CUDA not available, using CPU")
|
||||
providers = ["CPUExecutionProvider"]
|
||||
|
||||
except ImportError:
|
||||
print("[WARNING] ONNX Runtime not available, using CPU")
|
||||
providers = ["CPUExecutionProvider"]
|
||||
|
||||
print(f"[INFO] Using execution providers: {providers}")
|
||||
|
||||
# Initialize face analysis app
|
||||
self.face_model = FaceAnalysis(
|
||||
name="buffalo_l", # or 'buffalo_s' for smaller model
|
||||
providers=providers,
|
||||
)
|
||||
|
||||
# For MPS/CoreML, we need to adjust context
|
||||
ctx_id = -1 # Default for CPU
|
||||
if use_mps and "CoreMLExecutionProvider" in providers:
|
||||
ctx_id = 0 # CoreML uses device 0
|
||||
|
||||
self.face_model.prepare(ctx_id=ctx_id, det_size=(640, 640))
|
||||
|
||||
print("[INFO] InsightFace models loaded successfully")
|
||||
return True
|
||||
|
||||
except ImportError as e:
|
||||
print(f"[ERROR] Failed to import InsightFace: {e}")
|
||||
print("[INFO] Install with: pip install insightface")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Failed to load models: {e}")
|
||||
return False
|
||||
|
||||
def load_database(self, database_path: str):
|
||||
"""Load existing face database"""
|
||||
self.database_path = database_path
|
||||
|
||||
if os.path.exists(database_path):
|
||||
try:
|
||||
with open(database_path, "r") as f:
|
||||
self.face_database = json.load(f)
|
||||
print(f"[INFO] Loaded {len(self.face_database)} faces from database")
|
||||
except Exception as e:
|
||||
print(f"[WARNING] Failed to load database: {e}")
|
||||
self.face_database = {}
|
||||
else:
|
||||
print("[INFO] Creating new face database")
|
||||
self.face_database = {}
|
||||
|
||||
def save_database(self):
|
||||
"""Save face database to file"""
|
||||
if not self.database_path:
|
||||
print("[ERROR] No database path specified")
|
||||
return False
|
||||
|
||||
try:
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(os.path.dirname(self.database_path), exist_ok=True)
|
||||
|
||||
with open(self.database_path, "w") as f:
|
||||
json.dump(self.face_database, f, indent=2)
|
||||
|
||||
print(f"[INFO] Saved {len(self.face_database)} faces to database")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Failed to save database: {e}")
|
||||
return False
|
||||
|
||||
def register_face(
|
||||
self, image_path: str, name: str, metadata: Optional[Dict[str, Any]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""Register a new face from image"""
|
||||
# Check if image exists
|
||||
if not os.path.exists(image_path):
|
||||
return {
|
||||
"success": False,
|
||||
"message": f"Image not found: {image_path}",
|
||||
"face_id": None,
|
||||
"embedding": None,
|
||||
"attributes": None,
|
||||
}
|
||||
|
||||
# Load models if not already loaded
|
||||
if self.face_model is None:
|
||||
if not self.load_models():
|
||||
return {
|
||||
"success": False,
|
||||
"message": "Failed to load face models",
|
||||
"face_id": None,
|
||||
"embedding": None,
|
||||
"attributes": None,
|
||||
}
|
||||
|
||||
# Load image
|
||||
try:
|
||||
import cv2
|
||||
|
||||
image = cv2.imread(image_path)
|
||||
if image is None:
|
||||
return {
|
||||
"success": False,
|
||||
"message": f"Failed to load image: {image_path}",
|
||||
"face_id": None,
|
||||
"embedding": None,
|
||||
"attributes": None,
|
||||
}
|
||||
except ImportError:
|
||||
return {
|
||||
"success": False,
|
||||
"message": "OpenCV not installed",
|
||||
"face_id": None,
|
||||
"embedding": None,
|
||||
"attributes": None,
|
||||
}
|
||||
|
||||
# Detect faces
|
||||
try:
|
||||
faces = self.face_model.get(image)
|
||||
|
||||
if len(faces) == 0:
|
||||
return {
|
||||
"success": False,
|
||||
"message": "No faces detected in image",
|
||||
"face_id": None,
|
||||
"embedding": None,
|
||||
"attributes": None,
|
||||
}
|
||||
|
||||
if len(faces) > 1:
|
||||
print(f"[WARNING] Multiple faces detected, using the first one")
|
||||
|
||||
# Use the first face
|
||||
face = faces[0]
|
||||
|
||||
# Get embedding
|
||||
embedding = face.embedding.tolist() if hasattr(face, "embedding") else None
|
||||
|
||||
if embedding is None:
|
||||
return {
|
||||
"success": False,
|
||||
"message": "Failed to extract face embedding",
|
||||
"face_id": None,
|
||||
"embedding": None,
|
||||
"attributes": None,
|
||||
}
|
||||
|
||||
# Get attributes
|
||||
attributes = {}
|
||||
if hasattr(face, "age") and face.age is not None:
|
||||
attributes["age"] = int(face.age)
|
||||
if hasattr(face, "gender") and face.gender is not None:
|
||||
attributes["gender"] = "female" if face.gender == 0 else "male"
|
||||
|
||||
# Get pose if available
|
||||
pose = None
|
||||
if hasattr(face, "pose") and face.pose is not None:
|
||||
pose = {
|
||||
"yaw": float(face.pose[0]),
|
||||
"pitch": float(face.pose[1]),
|
||||
"roll": float(face.pose[2]),
|
||||
}
|
||||
|
||||
# Generate face ID
|
||||
import uuid
|
||||
|
||||
face_id = str(uuid.uuid4())
|
||||
|
||||
# Create face record
|
||||
face_record = {
|
||||
"face_id": face_id,
|
||||
"name": name,
|
||||
"embedding": embedding,
|
||||
"attributes": {
|
||||
"age": attributes.get("age"),
|
||||
"gender": attributes.get("gender"),
|
||||
"emotion": None,
|
||||
"glasses": None,
|
||||
"mask": None,
|
||||
"pose": pose,
|
||||
}
|
||||
if any([attributes.get("age"), attributes.get("gender"), pose])
|
||||
else None,
|
||||
"metadata": metadata or {},
|
||||
"registration_time": time.time(),
|
||||
"image_path": image_path,
|
||||
}
|
||||
|
||||
# Add to database
|
||||
self.face_database[face_id] = face_record
|
||||
|
||||
# Save database
|
||||
if not self.save_database():
|
||||
return {
|
||||
"success": False,
|
||||
"message": "Failed to save face to database",
|
||||
"face_id": face_id,
|
||||
"embedding": embedding,
|
||||
"attributes": face_record["attributes"],
|
||||
}
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"message": f"Face registered successfully as '{name}'",
|
||||
"face_id": face_id,
|
||||
"embedding": embedding,
|
||||
"attributes": face_record["attributes"],
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"success": False,
|
||||
"message": f"Face registration failed: {str(e)}",
|
||||
"face_id": None,
|
||||
"embedding": None,
|
||||
"attributes": None,
|
||||
}
|
||||
|
||||
def list_faces(self) -> Dict[str, Any]:
|
||||
"""List all registered faces"""
|
||||
faces = []
|
||||
|
||||
for face_id, face_data in self.face_database.items():
|
||||
faces.append(
|
||||
{
|
||||
"face_id": face_id,
|
||||
"name": face_data.get("name", "Unknown"),
|
||||
"registration_time": face_data.get("registration_time"),
|
||||
"metadata": face_data.get("metadata", {}),
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"message": f"Found {len(faces)} registered faces",
|
||||
"faces": faces,
|
||||
}
|
||||
|
||||
def delete_face(self, face_id: str) -> Dict[str, Any]:
|
||||
"""Delete a face from database"""
|
||||
if face_id not in self.face_database:
|
||||
return {"success": False, "message": f"Face ID not found: {face_id}"}
|
||||
|
||||
# Remove from database
|
||||
deleted_face = self.face_database.pop(face_id)
|
||||
|
||||
# Save database
|
||||
if not self.save_database():
|
||||
# Try to restore
|
||||
self.face_database[face_id] = deleted_face
|
||||
return {
|
||||
"success": False,
|
||||
"message": "Failed to save database after deletion",
|
||||
}
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"message": f"Face '{deleted_face.get('name', 'Unknown')}' deleted successfully",
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Face Registration Tool")
|
||||
parser.add_argument("image_path", help="Path to face image")
|
||||
parser.add_argument("output_path", help="Output JSON path")
|
||||
parser.add_argument("name", help="Name for the registered face")
|
||||
parser.add_argument(
|
||||
"--metadata", "-m", help="Path to metadata JSON file", default=""
|
||||
)
|
||||
parser.add_argument(
|
||||
"--database", "-d", help="Path to face database", default="face_database.json"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--list", "-l", help="List registered faces", action="store_true"
|
||||
)
|
||||
parser.add_argument("--delete", help="Delete face by ID", default="")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Initialize registration
|
||||
registration = FaceRegistration()
|
||||
|
||||
# Load database
|
||||
registration.load_database(args.database)
|
||||
|
||||
# Handle list command
|
||||
if args.list:
|
||||
result = registration.list_faces()
|
||||
with open(args.output_path, "w") as f:
|
||||
json.dump(result, f, indent=2)
|
||||
print(result["message"])
|
||||
return
|
||||
|
||||
# Handle delete command
|
||||
if args.delete:
|
||||
result = registration.delete_face(args.delete)
|
||||
with open(args.output_path, "w") as f:
|
||||
json.dump(result, f, indent=2)
|
||||
print(result["message"])
|
||||
return
|
||||
|
||||
# Load metadata if provided
|
||||
metadata = {}
|
||||
if args.metadata and os.path.exists(args.metadata):
|
||||
try:
|
||||
with open(args.metadata, "r") as f:
|
||||
metadata = json.load(f)
|
||||
except Exception as e:
|
||||
print(f"[WARNING] Failed to load metadata: {e}")
|
||||
|
||||
# Register face
|
||||
result = registration.register_face(
|
||||
image_path=args.image_path, name=args.name, metadata=metadata
|
||||
)
|
||||
|
||||
# Save result
|
||||
with open(args.output_path, "w") as f:
|
||||
json.dump(result, f, indent=2)
|
||||
|
||||
print(result["message"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user