uniface/tools/analyze.py

# Copyright 2025-2026 Yakhyokhuja Valikhujaev
# Author: Yakhyokhuja Valikhujaev
# GitHub: https://github.com/yakhyo

"""Face analysis using FaceAnalyzer.

Usage:
    python tools/analyze.py --source path/to/image.jpg
    python tools/analyze.py --source path/to/video.mp4
    python tools/analyze.py --source 0  # webcam
"""

from __future__ import annotations

import argparse
import os
from pathlib import Path

from _common import get_source_type
import cv2
import numpy as np

from uniface.analyzer import FaceAnalyzer
from uniface.attribute import AgeGender
from uniface.detection import RetinaFace
from uniface.draw import draw_detections
from uniface.recognition import ArcFace


def draw_face_info(image, face):
    """Draw face attributes above bounding box."""
    x1, y1, _x2, y2 = map(int, face.bbox)
    lines = []
    if face.age is not None and face.sex is not None:
        lines.append(f'{face.sex}, {face.age}y')
    if face.emotion is not None:
        lines.append(face.emotion)

    if not lines:
        return

    for i, line in enumerate(lines):
        y_pos = y1 - 10 - (len(lines) - 1 - i) * 25
        if y_pos < 20:
            y_pos = y2 + 20 + i * 25
        (tw, th), _ = cv2.getTextSize(line, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
        cv2.rectangle(image, (x1, y_pos - th - 5), (x1 + tw + 10, y_pos + 5), (0, 255, 0), -1)
        cv2.putText(image, line, (x1 + 5, y_pos), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2)


def process_image(analyzer, image_path: str, save_dir: str = 'outputs', show_similarity: bool = True):
    """Process a single image."""
    image = cv2.imread(image_path)
    if image is None:
        print(f"Error: Failed to load image from '{image_path}'")
        return

    faces = analyzer.analyze(image)
    print(f'Detected {len(faces)} face(s)')

    if not faces:
        return

    for i, face in enumerate(faces, 1):
        info = f'  Face {i}: {face.sex}, {face.age}y' if face.age and face.sex else f'  Face {i}'
        if face.embedding is not None:
            info += f' (embedding: {face.embedding.shape})'
        print(info)

    if show_similarity and len(faces) >= 2:
        print('\nSimilarity Matrix:')
        n = len(faces)
        sim_matrix = np.zeros((n, n))

        for i in range(n):
            for j in range(i, n):
                if i == j:
                    sim_matrix[i][j] = 1.0
                else:
                    sim = faces[i].compute_similarity(faces[j])
                    sim_matrix[i][j] = sim
                    sim_matrix[j][i] = sim

        print('     ', end='')
        for i in range(n):
            print(f'  F{i + 1:2d}  ', end='')
        print('\n     ' + '-' * (7 * n))

        for i in range(n):
            print(f'F{i + 1:2d} | ', end='')
            for j in range(n):
                print(f'{sim_matrix[i][j]:6.3f} ', end='')
            print()

        pairs = [(i, j, sim_matrix[i][j]) for i in range(n) for j in range(i + 1, n)]
        pairs.sort(key=lambda x: x[2], reverse=True)

        print('\nTop matches (>0.4 = same person):')
        for i, j, sim in pairs[:3]:
            status = 'Same' if sim > 0.4 else 'Different'
            print(f'  Face {i + 1} ↔ Face {j + 1}: {sim:.3f} ({status})')

    draw_detections(image=image, faces=faces, corner_bbox=True)

    for face in faces:
        draw_face_info(image, face)

    os.makedirs(save_dir, exist_ok=True)
    output_path = os.path.join(save_dir, f'{Path(image_path).stem}_analysis.jpg')
    cv2.imwrite(output_path, image)
    print(f'Output saved: {output_path}')


def process_video(analyzer, video_path: str, save_dir: str = 'outputs'):
    """Process a video file."""
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error: Cannot open video file '{video_path}'")
        return

    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))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    os.makedirs(save_dir, exist_ok=True)
    output_path = os.path.join(save_dir, f'{Path(video_path).stem}_analysis.mp4')
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    print(f'Processing video: {video_path} ({total_frames} frames)')
    frame_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        frame_count += 1
        faces = analyzer.analyze(frame)

        draw_detections(image=frame, faces=faces, corner_bbox=True)

        for face in faces:
            draw_face_info(frame, face)

        cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        out.write(frame)

        if frame_count % 100 == 0:
            print(f'  Processed {frame_count}/{total_frames} frames...')

    cap.release()
    out.release()
    print(f'Done! Output saved: {output_path}')


def run_camera(analyzer, camera_id: int = 0):
    """Run real-time analysis on webcam."""
    cap = cv2.VideoCapture(camera_id)
    if not cap.isOpened():
        print(f'Cannot open camera {camera_id}')
        return

    print("Press 'q' to quit")

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame = cv2.flip(frame, 1)

        faces = analyzer.analyze(frame)

        draw_detections(image=frame, faces=faces, corner_bbox=True)

        for face in faces:
            draw_face_info(frame, face)

        cv2.putText(frame, f'Faces: {len(faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.imshow('Face Analyzer', frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


def main():
    parser = argparse.ArgumentParser(description='Face analysis with detection, recognition, and attributes')
    parser.add_argument('--source', type=str, required=True, help='Image/video path or camera ID (0, 1, ...)')
    parser.add_argument('--save-dir', type=str, default='outputs', help='Output directory')
    parser.add_argument('--no-similarity', action='store_true', help='Skip similarity matrix computation')
    args = parser.parse_args()

    detector = RetinaFace()
    recognizer = ArcFace()
    age_gender = AgeGender()
    analyzer = FaceAnalyzer(detector, recognizer=recognizer, attributes=[age_gender])

    source_type = get_source_type(args.source)

    if source_type == 'camera':
        run_camera(analyzer, int(args.source))
    elif source_type == 'image':
        if not os.path.exists(args.source):
            print(f'Error: Image not found: {args.source}')
            return
        process_image(analyzer, args.source, args.save_dir, show_similarity=not args.no_similarity)
    elif source_type == 'video':
        if not os.path.exists(args.source):
            print(f'Error: Video not found: {args.source}')
            return
        process_video(analyzer, args.source, args.save_dir)
    else:
        print(f"Error: Unknown source type for '{args.source}'")
        print('Supported formats: images (.jpg, .png, ...), videos (.mp4, .avi, ...), or camera ID (0, 1, ...)')


if __name__ == '__main__':
    main()