aligners/__init__.py

from .base import BaseAligner
from .none import NoneAligner
from .differentiable_face_aligner import DifferentiableFaceAligner
from .retinaface_aligner import RetinaFaceAligner

def get_aligner(aligner_cfg):

    if aligner_cfg.name == 'none':
        aligner = NoneAligner.from_config(aligner_cfg)
    elif aligner_cfg.name == 'retinaface_aligner':
        aligner = RetinaFaceAligner.from_config(aligner_cfg)
    elif aligner_cfg.name == 'differentiable_face_aligner':
        aligner = DifferentiableFaceAligner.from_config(aligner_cfg)
    else:
        raise ValueError(f"Unknown classifier: {aligner_cfg.name}")

    if aligner_cfg.start_from:
        aligner.load_state_dict_from_path(aligner_cfg.start_from)

    if aligner_cfg.freeze:
        for param in aligner.parameters():
            param.requires_grad = False
    return aligner