from .common import count_params
from .detector import Detector
from .extractor import Extractor
from .orientation import OrientationNet
from .recognizer import Recognizer

BUDGET = {"detector": 1_200_000, "recognizer": 3_200_000, "extractor": 400_000, "orientation": 200_000}

BUILDERS = {"detector": Detector, "recognizer": Recognizer, "extractor": Extractor, "orientation": OrientationNet}


def build(name, **kwargs):
    return BUILDERS[name](**kwargs)


def parameter_report():
    rows = {name: count_params(build(name)) for name in BUILDERS}
    rows["total"] = sum(rows.values())
    return rows