main.py

"""
main.py — Face Mask Detection System entry point.

Orchestrates the threaded video capture pipeline, face detection,
mask classification, and UI rendering in a single, clean main loop.

Usage
-----
    python main.py

Controls
--------
    Q / ESC  — quit
    Close window button — quit
"""

import logging
import time

import cv2

from src import config, detector, ui
from src.camera import ThreadedCamera


def main() -> None:
    logging.basicConfig(format=config.LOG_FORMAT, level=logging.INFO)
    logger = logging.getLogger(__name__)
    logger.info("Starting Face Mask Detection System...")

    # ------------------------------------------------------------------
    # 1. Initialise AI components
    # ------------------------------------------------------------------
    try:
        face_tracker = detector.YuNetFaceDetector()
        mask_classifier = detector.MaskDetector()
        ui_manager = ui.UIManager()
    except Exception as exc:
        logger.critical(f"Initialisation failed: {exc}")
        return

    # ------------------------------------------------------------------
    # 2. Initialise threaded camera
    # ------------------------------------------------------------------
    logger.info(
        f"Starting threaded video stream (camera index {config.CAMERA_INDEX})..."
    )
    try:
        cam = ThreadedCamera(
            src=config.CAMERA_INDEX,
            width=config.CAMERA_WIDTH,
            height=config.CAMERA_HEIGHT,
        ).start()
    except RuntimeError as exc:
        logger.critical(f"Camera could not be opened: {exc}")
        return

    window_name = config.WINDOW_NAME
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
    logger.info("System running. Press Q or ESC to quit.")

    # ------------------------------------------------------------------
    # 3. Main processing loop
    # ------------------------------------------------------------------
    try:
        while True:
            # --- 3a. Capture ---
            display_frame = cam.read()
            if display_frame is None:
                time.sleep(0.01)
                continue

            # Mirror for a natural selfie-camera feel
            display_frame = cv2.flip(display_frame, 1)

            # --- 3b. Downscale for faster inference ---
            proc_frame = UIManager_resize(display_frame, ui_manager)
            h_disp, w_disp = display_frame.shape[:2]
            h_proc, w_proc = proc_frame.shape[:2]
            ratio_x = w_disp / float(w_proc)
            ratio_y = h_disp / float(h_proc)

            # --- 3c. Face detection ---
            detected_faces = face_tracker.detect(proc_frame)

            batch_faces = []
            batch_locs = []

            for face in detected_faces:
                x, y, w, h = face[0:4].astype(int)

                # Clamp to frame boundaries
                x = max(0, x)
                y = max(0, y)
                end_x = min(w_proc, x + w)
                end_y = min(h_proc, y + h)

                # Reject tiny detections (noise / tracking artefacts)
                if (end_x - x) < config.MIN_FACE_SIZE_PX or (
                    end_y - y
                ) < config.MIN_FACE_SIZE_PX:
                    continue

                face_roi = proc_frame[y:end_y, x:end_x]
                if face_roi.size > 0:
                    batch_faces.append(face_roi)
                    batch_locs.append((x, y, w, h))

            mask_count = 0
            no_mask_count = 0

            # --- 3d. Batch mask classification ---
            if batch_faces:
                predictions = mask_classifier.predict_batch(batch_faces)

                for (x, y, w, h), (mask_prob, no_mask_prob) in zip(
                    batch_locs, predictions
                ):
                    if mask_prob >= no_mask_prob:
                        label = "Mask"
                        color = config.COLOR_MASK
                        confidence = mask_prob * 100
                        mask_count += 1
                    else:
                        label = "No Mask"
                        color = config.COLOR_NO_MASK
                        confidence = no_mask_prob * 100
                        no_mask_count += 1

                    # Map inference coordinates back to display resolution
                    real_x = int(x * ratio_x)
                    real_y = int(y * ratio_y)
                    real_end_x = min(w_disp, real_x + int(w * ratio_x))
                    real_end_y = min(h_disp, real_y + int(h * ratio_y))

                    ui_manager.draw_pill_label(
                        display_frame,
                        (real_x, real_y, real_end_x, real_end_y),
                        label,
                        confidence,
                        color,
                    )

            # --- 3e. Dashboard ---
            ui_manager.draw_floating_dashboard(display_frame, mask_count, no_mask_count)
            cv2.imshow(window_name, display_frame)

            # --- 3f. Exit conditions ---
            key = cv2.waitKey(1) & 0xFF
            if key in (ord("q"), ord("Q"), 27):  # Q or ESC
                break
            try:
                if cv2.getWindowProperty(window_name, cv2.WND_PROP_VISIBLE) < 1:
                    break
            except cv2.error:
                break

    except KeyboardInterrupt:
        logger.info("Interrupted by user (Ctrl+C).")
    except Exception as exc:
        logger.error(f"Unexpected runtime error: {exc}", exc_info=True)
    finally:
        # Guaranteed cleanup regardless of how the loop exited
        cam.stop()
        cv2.destroyAllWindows()
        logger.info("Application exited cleanly.")


def UIManager_resize(display_frame, ui_manager: ui.UIManager):
    """Thin delegation so main() stays clean."""
    return ui_manager.resize_for_inference(display_frame, width=config.INFERENCE_WIDTH)


if __name__ == "__main__":
    main()