|
| 1 | +"""Landmark extraction module providing a unified interface for multiple extraction models. |
| 2 | +
|
| 3 | +This module provides a factory class for creating landmark extractors with different |
| 4 | +underlying implementations. It supports multiple extraction models through a consistent |
| 5 | +interface. |
| 6 | +
|
| 7 | +Example: |
| 8 | + Basic usage with default settings: |
| 9 | +
|
| 10 | + >>> from mukh.landmarks import LandmarkExtractor |
| 11 | + >>> extractor = LandmarkExtractor.create("blazeface") |
| 12 | + >>> landmarks = extractor.extract("image.jpg") |
| 13 | +
|
| 14 | + List available models: |
| 15 | +
|
| 16 | + >>> LandmarkExtractor.list_available_models() |
| 17 | + ['blazeface', 'mediapipe'] |
| 18 | +""" |
| 19 | + |
| 20 | +from typing import List, Literal |
| 21 | + |
| 22 | +from .models.base_extractor import BaseLandmarkExtractor |
| 23 | +from .models.blazeface import BlazeFaceLandmarkExtractor |
| 24 | + |
| 25 | +ExtractorType = Literal["blazeface", "mediapipe"] |
| 26 | + |
| 27 | + |
| 28 | +class LandmarkExtractor: |
| 29 | + """Factory class for creating landmark extraction model instances. |
| 30 | +
|
| 31 | + This class provides a unified interface to create and use different landmark extraction |
| 32 | + models through a consistent API. |
| 33 | + """ |
| 34 | + |
| 35 | + @staticmethod |
| 36 | + def create(model: ExtractorType) -> BaseLandmarkExtractor: |
| 37 | + """Creates a landmark extractor instance of the specified type. |
| 38 | +
|
| 39 | + Args: |
| 40 | + model: The type of extractor to create. Must be one of: "blazeface" or |
| 41 | + "mediapipe". |
| 42 | +
|
| 43 | + Returns: |
| 44 | + A BaseLandmarkExtractor instance of the requested type. |
| 45 | +
|
| 46 | + Raises: |
| 47 | + ValueError: If the specified model type is not supported. |
| 48 | + """ |
| 49 | + detectors = { |
| 50 | + "blazeface": BlazeFaceLandmarkExtractor, |
| 51 | + } |
| 52 | + |
| 53 | + if model not in detectors: |
| 54 | + raise ValueError( |
| 55 | + f"Unknown detector model: {model}. " |
| 56 | + f"Available models: {list(detectors.keys())}" |
| 57 | + ) |
| 58 | + |
| 59 | + return detectors[model]() |
| 60 | + |
| 61 | + @staticmethod |
| 62 | + def list_available_models() -> List[str]: |
| 63 | + """Returns a list of available face detection model names. |
| 64 | +
|
| 65 | + Returns: |
| 66 | + List of strings containing supported model names. |
| 67 | + """ |
| 68 | + return ["blazeface", "mediapipe"] |
0 commit comments