utils/registry.py

from typing import Callable, Dict, Any, List, Optional

class ModelEntry:
    def __init__(self, model: Any, preprocess: Callable, postprocess: Callable, class_names: List[str],
                 display_name: Optional[str] = None, contributor: Optional[str] = None, model_path: Optional[str] = None,
                 architecture: Optional[str] = None, dataset: Optional[str] = None):
        self.model = model
        self.preprocess = preprocess
        self.postprocess = postprocess
        self.class_names = class_names
        self.display_name = display_name
        self.contributor = contributor
        self.model_path = model_path
        self.architecture = architecture
        self.dataset = dataset

MODEL_REGISTRY: Dict[str, ModelEntry] = {}

def register_model(model_id: str, model: Any, preprocess: Callable, postprocess: Callable, class_names: List[str], architecture: Optional[str] = None, dataset: Optional[str] = None):
    MODEL_REGISTRY[model_id] = ModelEntry(model, preprocess, postprocess, class_names, architecture=architecture, dataset=dataset)