mirror of
https://github.com/yakhyo/uniface.git
synced 2025-12-30 09:02:25 +00:00
108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
# Face detection on video files
|
|
# Usage: python run_video_detection.py --input video.mp4 --output output.mp4
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
from tqdm import tqdm
|
|
|
|
from uniface import SCRFD, RetinaFace
|
|
from uniface.visualization import draw_detections
|
|
|
|
|
|
def process_video(
|
|
detector,
|
|
input_path: str,
|
|
output_path: str,
|
|
threshold: float = 0.6,
|
|
show_preview: bool = False,
|
|
):
|
|
cap = cv2.VideoCapture(input_path)
|
|
if not cap.isOpened():
|
|
print(f"Error: Cannot open video file '{input_path}'")
|
|
return
|
|
|
|
# get video properties
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
|
|
print(f"Input: {input_path} ({width}x{height}, {fps:.1f} fps, {total_frames} frames)")
|
|
print(f"Output: {output_path}")
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*"mp4v") # codec for .mp4
|
|
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
|
|
|
|
if not out.isOpened():
|
|
print(f"Error: Cannot create output video '{output_path}'")
|
|
cap.release()
|
|
return
|
|
|
|
frame_count = 0
|
|
total_faces = 0
|
|
|
|
for _ in tqdm(range(total_frames), desc="Processing", unit="frames"):
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
frame_count += 1
|
|
faces = detector.detect(frame)
|
|
total_faces += len(faces)
|
|
|
|
bboxes = [f["bbox"] for f in faces]
|
|
scores = [f["confidence"] for f in faces]
|
|
landmarks = [f["landmarks"] for f in faces]
|
|
draw_detections(frame, bboxes, scores, landmarks, vis_threshold=threshold)
|
|
|
|
cv2.putText(
|
|
frame,
|
|
f"Faces: {len(faces)}",
|
|
(10, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
1,
|
|
(0, 255, 0),
|
|
2,
|
|
)
|
|
out.write(frame)
|
|
|
|
if show_preview:
|
|
cv2.imshow("Processing - Press 'q' to cancel", frame)
|
|
if cv2.waitKey(1) & 0xFF == ord("q"):
|
|
print("\nCancelled by user")
|
|
break
|
|
|
|
cap.release()
|
|
out.release()
|
|
if show_preview:
|
|
cv2.destroyAllWindows()
|
|
|
|
avg_faces = total_faces / frame_count if frame_count > 0 else 0
|
|
print(f"\nDone! {frame_count} frames, {total_faces} faces ({avg_faces:.1f} avg/frame)")
|
|
print(f"Saved: {output_path}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Process video with face detection")
|
|
parser.add_argument("--input", type=str, required=True, help="Input video path")
|
|
parser.add_argument("--output", type=str, required=True, help="Output video path")
|
|
parser.add_argument("--detector", type=str, default="retinaface", choices=["retinaface", "scrfd"])
|
|
parser.add_argument("--threshold", type=float, default=0.6, help="Visualization threshold")
|
|
parser.add_argument("--preview", action="store_true", help="Show live preview")
|
|
args = parser.parse_args()
|
|
|
|
if not Path(args.input).exists():
|
|
print(f"Error: Input file '{args.input}' does not exist")
|
|
return
|
|
|
|
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
detector = RetinaFace() if args.detector == "retinaface" else SCRFD()
|
|
process_video(detector, args.input, args.output, args.threshold, args.preview)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|