| 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 |
|
|
|
|
