main.py

#!/usr/bin/env python3
"""
Air Paint – draw in mid-air with hand gestures.
Uses MediaPipe for hand tracking and OpenCV for rendering.
"""

import os
import sys
import urllib.request
from collections import deque
from datetime import datetime

import cv2
import numpy as np
import mediapipe as mp
from mediapipe.tasks.python import vision

from config import (
    COLORS, DEFAULT_BRUSH_SIZE, MIN_BRUSH_SIZE, MAX_BRUSH_SIZE,
    SMOOTHING_WINDOW, CAM_WIDTH, CAM_HEIGHT, MODEL_URL,
)
from gestures import classify
from overlay import draw_palette, draw_status, draw_hand, draw_cursor


def _ensure_model(path: str):
    """Download the hand model if missing."""
    if os.path.exists(path):
        return
    print("Downloading hand landmark model ...")
    urllib.request.urlretrieve(MODEL_URL, path)
    print("Done.")


def _build_landmarker(model_path: str):
    opts = vision.HandLandmarkerOptions(
        base_options=mp.tasks.BaseOptions(model_asset_path=model_path),
        running_mode=vision.RunningMode.IMAGE,
        num_hands=1,
        min_hand_detection_confidence=0.6,
        min_hand_presence_confidence=0.6,
        min_tracking_confidence=0.6,
    )
    return vision.HandLandmarker.create_from_options(opts)


def _open_camera():
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("Error: no camera found.")
        sys.exit(1)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAM_WIDTH)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAM_HEIGHT)
    return cap


def main():
    base_dir = os.path.dirname(os.path.abspath(__file__))
    model_path = os.path.join(base_dir, "hand_landmarker.task")
    _ensure_model(model_path)

    landmarker = _build_landmarker(model_path)
    cap = _open_camera()

    # state
    canvas = None
    color_idx = 0
    brush_size = DEFAULT_BRUSH_SIZE
    prev_point = None
    prev_gesture = None
    smooth_buf = deque(maxlen=SMOOTHING_WINDOW)

    # fps tracking
    fps_timer = cv2.getTickCount()
    fps = 0
    frame_count = 0

    print("=" * 46)
    print("  Air Paint")
    print("  Point finger = draw | Fist = pause")
    print("  Pinch = next color   | Q = quit")
    print("=" * 46)

    cv2.namedWindow("Air Paint", cv2.WINDOW_NORMAL)

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

        frame = cv2.flip(frame, 1)  # mirror
        h, w = frame.shape[:2]

        if canvas is None:
            canvas = np.zeros((h, w, 3), dtype=np.uint8)

        # fps
    frame_count += 1
        elapsed = (cv2.getTickCount() - fps_timer) / cv2.getTickFrequency()
        if elapsed >= 1.0:
            fps = int(frame_count / elapsed)
            frame_count = 0
            fps_timer = cv2.getTickCount()

        # detect hand
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        mp_img = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
        results = landmarker.detect(mp_img)

        gesture = None
        finger_pos = None

        if results.hand_landmarks:
            lm = results.hand_landmarks[0]
            hand_label = results.handedness[0][0].category_name
            gesture = classify(lm, hand_label)
            finger_pos = (int(lm[8].x * w), int(lm[8].y * h))
            draw_hand(frame, lm, w, h)

        # gesture -> action
        if gesture == "point" and finger_pos:
            smooth_buf.append(finger_pos)
            sx = int(sum(p[0] for p in smooth_buf) / len(smooth_buf))
            sy = int(sum(p[1] for p in smooth_buf) / len(smooth_buf))
            cur = (sx, sy)

            if prev_point is not None:
                cv2.line(canvas, prev_point, cur,
                         COLORS[color_idx][0], brush_size, cv2.LINE_AA)
            prev_point = cur

        elif gesture == "pinch":
            if prev_gesture != "pinch":          # trigger once per pinch
                color_idx = (color_idx + 1) % len(COLORS)
                print(f"Color: {COLORS[color_idx][1]}")
            prev_point = None
            smooth_buf.clear()

        else:  # fist / other / no hand
            prev_point = None
            smooth_buf.clear()

        prev_gesture = gesture

        # blend canvas on top of camera feed
        gray = cv2.cvtColor(canvas, cv2.COLOR_BGR2GRAY)
        _, mask = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY)
        inv = cv2.bitwise_not(mask)
        bg = cv2.bitwise_and(frame, frame, mask=inv)
        fg = cv2.bitwise_and(canvas, canvas, mask=mask)
        output = cv2.add(bg, fg)

        # ui
        draw_palette(output, color_idx)
        draw_status(output, gesture, brush_size, fps)
        draw_cursor(output, finger_pos, gesture,
                    COLORS[color_idx][0], brush_size)

        cv2.imshow("Air Paint", output)

        # keys
        key = cv2.waitKey(1) & 0xFF
        if key in (ord("q"), ord("Q"), 27):
            break
        elif key in (ord("c"), ord("C")):
            canvas = np.zeros((h, w, 3), dtype=np.uint8)
            print("Canvas cleared.")
        elif key in (ord("s"), ord("S")):
            name = f"air_paint_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
            cv2.imwrite(name, output)
            print(f"Saved: {name}")
        elif key in (ord("+"), ord("=")):
            brush_size = min(brush_size + 2, MAX_BRUSH_SIZE)
        elif key in (ord("-"), ord("_")):
            brush_size = max(brush_size - 2, MIN_BRUSH_SIZE)
        elif ord("1") <= key <= ord("8"):
            color_idx = key - ord("1")
            print(f"Color: {COLORS[color_idx][1]}")

    # cleanup
    landmarker.close()
    cap.release()
    cv2.destroyAllWindows()
    print("Air Paint closed.")


if __name__ == "__main__":
    main()