src/model/encoder/__init__.py

from typing import Optional, Union

from .encoder import Encoder
from .visualization.encoder_visualizer import EncoderVisualizer
from .of3gs import EncoderOF3GS, EncoderOF3GSCfg

ENCODERS = {
    "of3gs": (EncoderOF3GS, None),
}

EncoderCfg = Union[EncoderOF3GSCfg]


def get_encoder(cfg: EncoderCfg) -> tuple[Encoder, Optional[EncoderVisualizer]]:
    encoder, visualizer = ENCODERS[cfg.name]
    encoder = encoder(cfg)
    if visualizer is not None:
        visualizer = visualizer(cfg.visualizer, encoder)
    return encoder, visualizer