landmarkdiff 0.2.3__py3-none-any.whl

landmarkdiff/landmarks.py

@@ -0,0 +1,432 @@
+"""Facial landmark extraction using MediaPipe Face Mesh v2."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+
+import cv2
+import mediapipe as mp
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# Region color map for visualization (BGR)
+REGION_COLORS: dict[str, tuple[int, int, int]] = {
+    "jawline": (255, 255, 255),  # white
+    "eyebrow_left": (0, 255, 0),  # green
+    "eyebrow_right": (0, 255, 0),
+    "eye_left": (255, 255, 0),  # cyan
+    "eye_right": (255, 255, 0),
+    "nose": (0, 255, 255),  # yellow
+    "lips": (0, 0, 255),  # red
+    "iris_left": (255, 0, 255),  # magenta
+    "iris_right": (255, 0, 255),
+}
+
+# MediaPipe landmark index groups by anatomical region
+LANDMARK_REGIONS: dict[str, list[int]] = {
+    "jawline": [
+        10,
+        338,
+        297,
+        332,
+        284,
+        251,
+        389,
+        356,
+        454,
+        323,
+        361,
+        288,
+        397,
+        365,
+        379,
+        378,
+        400,
+        377,
+        152,
+        148,
+        176,
+        149,
+        150,
+        136,
+        172,
+        58,
+        132,
+        93,
+        234,
+        127,
+        162,
+        21,
+        54,
+        103,
+        67,
+        109,
+    ],
+    "eye_left": [
+        33,
+        7,
+        163,
+        144,
+        145,
+        153,
+        154,
+        155,
+        133,
+        173,
+        157,
+        158,
+        159,
+        160,
+        161,
+        246,
+    ],
+    "eye_right": [
+        362,
+        382,
+        381,
+        380,
+        374,
+        373,
+        390,
+        249,
+        263,
+        466,
+        388,
+        387,
+        386,
+        385,
+        384,
+        398,
+    ],
+    "eyebrow_left": [70, 63, 105, 66, 107, 55, 65, 52, 53, 46],
+    "eyebrow_right": [300, 293, 334, 296, 336, 285, 295, 282, 283, 276],
+    "nose": [
+        1,
+        2,
+        4,
+        5,
+        6,
+        19,
+        94,
+        141,
+        168,
+        195,
+        197,
+        236,
+        240,
+        274,
+        275,
+        278,
+        279,
+        294,
+        326,
+        327,
+        360,
+        363,
+        370,
+        456,
+        460,
+    ],
+    "lips": [
+        61,
+        146,
+        91,
+        181,
+        84,
+        17,
+        314,
+        405,
+        321,
+        375,
+        291,
+        308,
+        324,
+        318,
+        402,
+        317,
+        14,
+        87,
+        178,
+        88,
+        95,
+        78,
+    ],
+    "iris_left": [468, 469, 470, 471, 472],
+    "iris_right": [473, 474, 475, 476, 477],
+}
+
+
+@dataclass(frozen=True)
+class FaceLandmarks:
+    """Extracted facial landmarks with metadata."""
+
+    landmarks: np.ndarray  # (478, 3) normalized (x, y, z)
+    image_width: int
+    image_height: int
+    confidence: float
+
+    @property
+    def pixel_coords(self) -> np.ndarray:
+        """Convert normalized landmarks to pixel coordinates (478, 2)."""
+        coords = self.landmarks[:, :2].copy()
+        coords[:, 0] *= self.image_width
+        coords[:, 1] *= self.image_height
+        return coords
+
+    def pixel_coords_at(self, width: int, height: int) -> np.ndarray:
+        """Convert normalized landmarks to pixel coordinates at a given size.
+
+        Use this when the image has been resized after landmark extraction.
+        """
+        coords = self.landmarks[:, :2].copy()
+        coords[:, 0] *= width
+        coords[:, 1] *= height
+        return coords
+
+    def rescale(self, width: int, height: int) -> FaceLandmarks:
+        """Return a copy with updated image dimensions.
+
+        Landmarks stay in normalized [0,1] space; only the stored
+        width/height change, so ``pixel_coords`` returns values at
+        the new resolution.
+        """
+        return FaceLandmarks(
+            landmarks=self.landmarks.copy(),
+            image_width=width,
+            image_height=height,
+            confidence=self.confidence,
+        )
+
+    def get_region(self, region: str) -> np.ndarray:
+        """Get landmark indices for a named region."""
+        indices = LANDMARK_REGIONS.get(region, [])
+        return self.landmarks[indices]
+
+
+def extract_landmarks(
+    image: np.ndarray,
+    min_detection_confidence: float = 0.5,
+    min_tracking_confidence: float = 0.5,
+) -> FaceLandmarks | None:
+    """Extract 478 facial landmarks from an image using MediaPipe Face Mesh.
+
+    Args:
+        image: BGR image as numpy array.
+        min_detection_confidence: Minimum face detection confidence.
+        min_tracking_confidence: Minimum landmark tracking confidence.
+
+    Returns:
+        FaceLandmarks if a face is detected, None otherwise.
+    """
+    h, w = image.shape[:2]
+    rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    # Try new Tasks API first (mediapipe >= 0.10.20), fall back to legacy solutions API
+    try:
+        landmarks, confidence = _extract_tasks_api(rgb, min_detection_confidence)
+    except Exception:
+        logger.debug("Tasks API unavailable, trying Solutions API", exc_info=True)
+        try:
+            landmarks, confidence = _extract_solutions_api(
+                rgb, min_detection_confidence, min_tracking_confidence
+            )
+        except Exception:
+            logger.debug("Both MediaPipe APIs failed", exc_info=True)
+            return None
+
+    if landmarks is None:
+        return None
+
+    return FaceLandmarks(
+        landmarks=landmarks,
+        image_width=w,
+        image_height=h,
+        confidence=confidence,
+    )
+
+
+def _extract_tasks_api(
+    rgb: np.ndarray,
+    min_confidence: float,
+) -> tuple[np.ndarray | None, float]:
+    """Extract landmarks using MediaPipe Tasks API (>= 0.10.20)."""
+    FaceLandmarker = mp.tasks.vision.FaceLandmarker
+    FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions
+    RunningMode = mp.tasks.vision.RunningMode
+    BaseOptions = mp.tasks.BaseOptions
+    import tempfile
+    import urllib.request
+
+    # Download model if not cached
+    model_path = Path(tempfile.gettempdir()) / "face_landmarker_v2_with_blendshapes.task"
+    if not model_path.exists():
+        url = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task"
+        urllib.request.urlretrieve(url, str(model_path))
+
+    options = FaceLandmarkerOptions(
+        base_options=BaseOptions(model_asset_path=str(model_path)),
+        running_mode=RunningMode.IMAGE,
+        num_faces=1,
+        min_face_detection_confidence=min_confidence,
+        output_face_blendshapes=False,
+        output_facial_transformation_matrixes=False,
+    )
+
+    with FaceLandmarker.create_from_options(options) as landmarker:
+        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
+        result = landmarker.detect(mp_image)
+
+    if not result.face_landmarks:
+        return None, 0.0
+
+    face_lms = result.face_landmarks[0]
+    landmarks = np.array(
+        [(lm.x, lm.y, lm.z) for lm in face_lms],
+        dtype=np.float32,
+    )
+
+    # MediaPipe Tasks API doesn't expose per-landmark detection confidence;
+    # return 1.0 to indicate successful detection
+    return landmarks, 1.0
+
+
+def _extract_solutions_api(
+    rgb: np.ndarray,
+    min_detection_confidence: float,
+    min_tracking_confidence: float,
+) -> tuple[np.ndarray | None, float]:
+    """Extract landmarks using legacy MediaPipe Solutions API."""
+    with mp.solutions.face_mesh.FaceMesh(
+        static_image_mode=True,
+        max_num_faces=1,
+        refine_landmarks=True,
+        min_detection_confidence=min_detection_confidence,
+        min_tracking_confidence=min_tracking_confidence,
+    ) as face_mesh:
+        results = face_mesh.process(rgb)
+
+    if not results.multi_face_landmarks:
+        return None, 0.0
+
+    face = results.multi_face_landmarks[0]
+    landmarks = np.array(
+        [(lm.x, lm.y, lm.z) for lm in face.landmark],
+        dtype=np.float32,
+    )
+    # Legacy API doesn't expose detection confidence; return 1.0 for success
+    return landmarks, 1.0
+
+
+def visualize_landmarks(
+    image: np.ndarray,
+    face: FaceLandmarks,
+    radius: int = 1,
+    draw_regions: bool = True,
+) -> np.ndarray:
+    """Draw colored landmark dots on image by anatomical region.
+
+    Args:
+        image: BGR image to draw on (will be copied).
+        face: Extracted face landmarks.
+        radius: Dot radius in pixels.
+        draw_regions: If True, color by region. Otherwise all white.
+
+    Returns:
+        Annotated image copy.
+    """
+    canvas = image.copy()
+    coords = face.pixel_coords
+
+    if draw_regions:
+        # Build index -> color mapping
+        idx_to_color: dict[int, tuple[int, int, int]] = {}
+        for region, indices in LANDMARK_REGIONS.items():
+            color = REGION_COLORS.get(region, (255, 255, 255))
+            for idx in indices:
+                idx_to_color[idx] = color
+
+        for i, (x, y) in enumerate(coords):
+            color = idx_to_color.get(i, (128, 128, 128))
+            cv2.circle(canvas, (int(x), int(y)), radius, color, -1)
+    else:
+        for x, y in coords:
+            cv2.circle(canvas, (int(x), int(y)), radius, (255, 255, 255), -1)
+
+    return canvas
+
+
+def render_landmark_image(
+    face: FaceLandmarks,
+    width: int | None = None,
+    height: int | None = None,
+    radius: int = 2,
+) -> np.ndarray:
+    """Render MediaPipe face mesh tessellation on black canvas.
+
+    Draws the full 2556-edge tessellation mesh that CrucibleAI/ControlNetMediaPipeFace
+    was pre-trained on. This is critical -- the ControlNet expects dense triangulated
+    wireframes, not sparse dots.
+
+    Falls back to colored dots if tessellation connections aren't available.
+
+    Args:
+        face: Extracted face landmarks.
+        width: Canvas width (defaults to face.image_width).
+        height: Canvas height (defaults to face.image_height).
+        radius: Dot radius (used for key landmark dots overlay).
+
+    Returns:
+        BGR image with face mesh on black background.
+    """
+    w = width or face.image_width
+    h = height or face.image_height
+    canvas = np.zeros((h, w, 3), dtype=np.uint8)
+
+    coords = face.landmarks[:, :2].copy()
+    coords[:, 0] *= w
+    coords[:, 1] *= h
+    pts = coords.astype(np.int32)
+
+    # Draw tessellation mesh (what CrucibleAI ControlNet expects)
+    try:
+        from mediapipe.tasks.python.vision.face_landmarker import FaceLandmarksConnections
+
+        tessellation = FaceLandmarksConnections.FACE_LANDMARKS_TESSELATION
+        contours = FaceLandmarksConnections.FACE_LANDMARKS_CONTOURS
+
+        # Draw tessellation edges (thin, gray-white)
+        for conn in tessellation:
+            p1 = tuple(pts[conn.start])
+            p2 = tuple(pts[conn.end])
+            cv2.line(canvas, p1, p2, (192, 192, 192), 1, cv2.LINE_AA)
+
+        # Draw contour edges on top (brighter, key features)
+        for conn in contours:
+            p1 = tuple(pts[conn.start])
+            p2 = tuple(pts[conn.end])
+            cv2.line(canvas, p1, p2, (255, 255, 255), 1, cv2.LINE_AA)
+
+    except (ImportError, AttributeError):
+        # Fallback: draw colored dots if tessellation not available
+        idx_to_color: dict[int, tuple[int, int, int]] = {}
+        for region, indices in LANDMARK_REGIONS.items():
+            color = REGION_COLORS.get(region, (128, 128, 128))
+            for idx in indices:
+                idx_to_color[idx] = color
+
+        for i, (x, y) in enumerate(coords):
+            color = idx_to_color.get(i, (128, 128, 128))
+            cv2.circle(canvas, (int(x), int(y)), radius, color, -1)
+
+    return canvas
+
+
+def load_image(path: str | Path) -> np.ndarray:
+    """Load an image from disk as BGR numpy array."""
+    img = cv2.imread(str(path))
+    if img is None:
+        raise FileNotFoundError(f"Could not load image: {path}")
+    return img

landmarkdiff/log.py

@@ -0,0 +1,90 @@
+"""Centralized logging configuration for LandmarkDiff.
+
+Usage:
+    from landmarkdiff.log import get_logger
+    logger = get_logger(__name__)
+    logger.info("Training started")
+
+Configure globally:
+    from landmarkdiff.log import setup_logging
+    setup_logging(level="DEBUG")  # affects all LandmarkDiff loggers
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+
+_CONFIGURED = False
+
+# Default format
+LOG_FORMAT = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+LOG_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
+
+
+def setup_logging(
+    level: str | int = "INFO",
+    fmt: str | None = None,
+    stream: object = None,
+    log_file: str | None = None,
+) -> None:
+    """Configure logging for the landmarkdiff package.
+
+    Call once at application startup. Subsequent calls update the level.
+
+    Args:
+        level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL).
+        fmt: Custom format string. None uses the default.
+        stream: Output stream. None defaults to stderr.
+        log_file: Optional file path for logging to a file.
+    """
+    global _CONFIGURED
+
+    if isinstance(level, str):
+        level = getattr(logging, level.upper(), logging.INFO)
+
+    root_logger = logging.getLogger("landmarkdiff")
+    root_logger.setLevel(level)
+
+    if not _CONFIGURED:
+        formatter = logging.Formatter(
+            fmt or LOG_FORMAT,
+            datefmt=LOG_DATE_FORMAT,
+        )
+
+        # Add stream handler
+        stream_handler = logging.StreamHandler(stream or sys.stderr)
+        stream_handler.setFormatter(formatter)
+        root_logger.addHandler(stream_handler)
+
+        # Add file handler if log_file is specified
+        if log_file:
+            file_handler = logging.FileHandler(log_file)
+            file_handler.setFormatter(formatter)
+            root_logger.addHandler(file_handler)
+
+        # Prevent propagation to root logger to avoid duplicate messages
+        root_logger.propagate = False
+        _CONFIGURED = True
+    else:
+        # Just update the level
+        root_logger.setLevel(level)
+
+
+def get_logger(name: str) -> logging.Logger:
+    """Get a logger for a LandmarkDiff module.
+
+    The returned logger is a child of the 'landmarkdiff' root logger,
+    so setup_logging() controls its level and formatting.
+
+    Args:
+        name: Module name (typically __name__).
+
+    Returns:
+        Configured logging.Logger instance.
+    """
+    # Ensure base configuration exists
+    if not _CONFIGURED:
+        setup_logging()
+
+    return logging.getLogger(name)