nail_detection.py

"""
Self-contained nail detection pipeline.

Inlined from hand-ki-model so that the HF Space has no git-clone dependency.
Source: nail_detection/{main,hand_detection,extract_nails}.py + utils/{rotate,polygon,angle,valid_crop,draw_hand}.py
"""

from __future__ import annotations

import math
from typing import List

import cv2
import mediapipe as mp
import numpy as np
from PIL import Image, ImageDraw
from scipy import ndimage


# ---------------------------------------------------------------------------
# Geometry helpers (from utils/)
# ---------------------------------------------------------------------------

def _rotate(origin, point, angle: float):
    """Rotate *point* counter-clockwise by *angle* (radians) around *origin*."""
    ox, oy = origin
    px, py = point
    qx = ox + math.cos(angle) * (px - ox) - math.sin(angle) * (py - oy)
    qy = oy + math.sin(angle) * (px - ox) + math.cos(angle) * (py - oy)
    return qx, qy


def _unit_vector(v):
    return v / np.linalg.norm(v)


def _angle_between(v1, v2) -> float:
    return float(np.arccos(np.clip(np.dot(_unit_vector(v1), _unit_vector(v2)), -1.0, 1.0)))


def _get_polygon_mask(width: int, height: int, polygon_idx) -> np.ndarray:
    polygon_idx = [tuple(xy) for xy in polygon_idx]
    img = Image.new("L", (width, height), 0)
    ImageDraw.Draw(img).polygon(polygon_idx, outline=1, fill=1)
    return np.array(img)


def _valid_crop(image: np.ndarray, mask: np.ndarray, offset: int = 10):
    true_points = np.argwhere(mask)
    top_left = true_points.min(axis=0)
    bottom_right = true_points.max(axis=0)
    x_low = max(top_left[0] - offset, 0)
    x_high = min(bottom_right[0] + offset, image.shape[0])
    y_low = max(top_left[1] - offset, 0)
    y_high = min(bottom_right[1] + offset + 1, image.shape[1])
    return image[x_low:x_high, y_low:y_high], mask[x_low:x_high, y_low:y_high]


# ---------------------------------------------------------------------------
# Hand detection (from nail_detection/hand_detection.py)
# ---------------------------------------------------------------------------

def detect_hand(image: np.ndarray):
    """Return MediaPipe hand landmarks for the first detected hand, or None."""
    mp_hands = mp.solutions.hands
    with mp_hands.Hands(
        static_image_mode=True, max_num_hands=1, min_detection_confidence=0.0
    ) as hands:
        results = hands.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        if results.multi_handedness is None:
            return None
        return results.multi_hand_landmarks[0]


# ---------------------------------------------------------------------------
# Nail extraction (from nail_detection/extract_nails.py)
# ---------------------------------------------------------------------------

def extract_nails(image: np.ndarray, hand_landmarks) -> List[np.ndarray]:
    """Return list of 5 nail crop arrays (thumb → pinky), BGR."""
    mp_hands = mp.solutions.hands
    image_height, image_width, _ = image.shape
    nails: List[np.ndarray] = []

    for tip in [
        mp_hands.HandLandmark.THUMB_TIP,
        mp_hands.HandLandmark.INDEX_FINGER_TIP,
        mp_hands.HandLandmark.MIDDLE_FINGER_TIP,
        mp_hands.HandLandmark.RING_FINGER_TIP,
        mp_hands.HandLandmark.PINKY_TIP,
    ]:
        tip_coords = np.array([
            hand_landmarks.landmark[tip].x * image_width,
            hand_landmarks.landmark[tip].y * image_height,
        ])
        dip_coords = np.array([
            hand_landmarks.landmark[tip - 1].x * image_width,
            hand_landmarks.landmark[tip - 1].y * image_height,
        ])
        dt = tip_coords - dip_coords
        ext = np.array([tip_coords + 3 / 4 * dt, tip_coords - 3 / 4 * dt])

        origin = 0.5 * (ext[0] - ext[1]) + ext[1]
        orth_p1 = _rotate(origin, ext[0], np.deg2rad(90))
        orth_p2 = _rotate(origin, ext[1], np.deg2rad(90))
        orth = np.array([orth_p1, orth_p2])

        half = 0.5 * (ext[0] - ext[1])
        p1 = orth[0] + half
        p2 = orth[0] - half
        p3 = orth[1] + half
        p4 = orth[1] - half

        angle = 90 - np.rad2deg(_angle_between(ext[0] - ext[1], [1, 0]))

        mask = _get_polygon_mask(image_width, image_height, [p2, p1, p3, p4])
        mask3 = np.tile(mask[:, :, None], (1, 1, 3))
        masked = mask3 * image[:, :, ::-1]

        masked, mask3 = _valid_crop(masked, mask3)
        masked = ndimage.rotate(masked, angle)
        mask3 = ndimage.rotate(mask3, angle)
        masked, mask3 = _valid_crop(masked, mask3, offset=0)
        masked = masked[5:-5, 5:-5]
        masked = np.ascontiguousarray(masked)
        nails.append(masked)

    return nails


# ---------------------------------------------------------------------------
# Draw hand skeleton (from utils/draw_hand.py)
# ---------------------------------------------------------------------------

def draw_hand(annotated_image: np.ndarray, hand_landmarks) -> None:
    mp_drawing = mp.solutions.drawing_utils
    mp_drawing_styles = mp.solutions.drawing_styles
    mp_hands = mp.solutions.hands
    mp_drawing.draw_landmarks(
        annotated_image,
        hand_landmarks,
        mp_hands.HAND_CONNECTIONS,
        mp_drawing_styles.get_default_hand_landmarks_style(),
        mp_drawing_styles.get_default_hand_connections_style(),
    )


# ---------------------------------------------------------------------------
# Top-level entry point
# ---------------------------------------------------------------------------

def get_nails_and_landmarks(image_bgr: np.ndarray):
    """
    Detect hand and extract all 5 nail crops.

    Args:
        image_bgr: BGR image as numpy array (as returned by cv2.imread or
                   cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)).

    Returns:
        (nails, hand_landmarks) or (None, None) if no hand detected.
    """
    if image_bgr.shape[0] < image_bgr.shape[1]:  # landscape → rotate
        image_bgr = cv2.rotate(image_bgr, cv2.ROTATE_90_CLOCKWISE)

    landmarks = detect_hand(image_bgr)
    if landmarks is None:
        return None, None

    # Flip upside-down images
    mp_hands = mp.solutions.hands
    if (
        landmarks.landmark[mp_hands.HandLandmark.MIDDLE_FINGER_PIP].y
        < landmarks.landmark[mp_hands.HandLandmark.MIDDLE_FINGER_TIP].y
    ):
        image_bgr = cv2.rotate(image_bgr, cv2.ROTATE_180)
        landmarks = detect_hand(image_bgr)

    nails = extract_nails(image_bgr, landmarks)
    return nails, landmarks