import cv2 import os import os.path as osp from pathlib import Path class ImageCache: data = {} def get_image(name, to_rgb=False): key = (name, to_rgb) if key in ImageCache.data: return ImageCache.data[key] images_dir = osp.join(Path(__file__).parent.absolute(), 'images') ext_names = ['.jpg', '.png', '.jpeg'] image_file = None for ext_name in ext_names: _image_file = osp.join(images_dir, "%s%s"%(name, ext_name)) if osp.exists(_image_file): image_file = _image_file break assert image_file is not None, '%s not found'%name img = cv2.imread(image_file) if to_rgb: img = img[:,:,::-1] ImageCache.data[key] = img return img