registry.py

import logging

# Configure Logger
logger = logging.getLogger(__name__)

def get_model(task: str, model_key: str, device="cpu"):
    """
    Dynamically retrieves the model instance based on the task and model_key.

    Args:
        task (str): One of "detection", "segmentation", or "depth".
        model_key (str): Model identifier or variant.
        device (str): Device to run inference on ("cpu" or "cuda").

    Returns:
        object: Initialized model ready for inference.

    Raises:
        ValueError: If task is unsupported or model loading fails.
    """
    logger.info(f"Request received to load model '{model_key}' for task '{task}' on device '{device}'")

    try:
        if task == "detection":
            from models.detection.detector import ObjectDetector
            return ObjectDetector(model_key=model_key, device=device)

        elif task == "segmentation":
            from models.segmentation.segmenter import Segmenter
            return Segmenter(model_key=model_key, device=device)

        elif task == "depth":
            from models.depth.depth_estimator import DepthEstimator
            return DepthEstimator(model_key=model_key, device=device)

        else:
            error_msg = f"Unsupported task '{task}'. Valid options are: 'detection', 'segmentation', 'depth'."
            logger.error(error_msg)
            raise ValueError(error_msg)

    except Exception as e:
        logger.error(f"Error while loading model '{model_key}' for task '{task}': {e}")
        raise