import cv2
import numpy as np
from typing import Union

from .exceptions import InvalidInputException


def load_image(img_data_or_path: Union[str, bytes]) -> np.ndarray:
    if isinstance(img_data_or_path, str):
        img = cv2.imread(img_data_or_path, cv2.IMREAD_COLOR)
    else:
        img = cv2.imdecode(
            np.frombuffer(img_data_or_path.read(), dtype=np.uint8), cv2.IMREAD_COLOR
        )

    if img is None:
        raise InvalidInputException("Couldn't load the invalid input image")

    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)