src/models/__init__.py

from __future__ import annotations

from .segformer import SegFormerStrokeSegmenter
from .unet import UnetStrokeSegmenter

__all__ = ["SegFormerStrokeSegmenter", "UnetStrokeSegmenter", "make_model"]


def make_model(name: str, **kwargs):
    if name == "segformer":
        return SegFormerStrokeSegmenter(**kwargs)
    if name == "unet":
        return UnetStrokeSegmenter(**kwargs)
    raise ValueError(f"Unknown model: {name}")