import torch


def transform_model_input_data_type(model_input, input_type: str):
    """
    Default model input follow RGB format. Based the model input specified,
    change the type. Supported types: RGB, BGR, LAB
    """
    model_output = model_input
    # In case the model takes BGR input type, we convert the RGB to BGR
    if input_type == "bgr":
        model_output = model_input[:, [2, 1, 0], :, :]
    # In case of LAB image, we take only "L" channel as input. Split the data
    # along the channel dimension into [L, AB] and keep only L channel.
    if input_type == "lab":
        model_output = torch.split(model_input, [1, 2], dim=1)[0]
    return model_output