yoloserv/modules/paravision/recognition/utils.py

20 lines
554 B
Python

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)