image.py 713 B

123456789101112131415161718192021222324
  1. import numpy as np
  2. from PIL.Image import Image as PIL_Image
  3. def dimensions(im):
  4. if isinstance(im, np.ndarray):
  5. if im.ndim != 3:
  6. import pdb; pdb.set_trace()
  7. assert im.ndim == 3, "Only RGB images are currently supported!"
  8. return im.shape
  9. elif isinstance(im, PIL_Image):
  10. w, h = im.size
  11. c = len(im.getbands())
  12. # assert c == 3, "Only RGB images are currently supported!"
  13. return h, w, c
  14. else:
  15. raise ValueError("Unknown image instance ({})!".format(type(im)))
  16. def asarray(im, dtype=np.uint8):
  17. if isinstance(im, np.ndarray):
  18. return im.astype(dtype)
  19. elif isinstance(im, PIL_Image):
  20. return np.asarray(im, dtype=dtype)
  21. else:
  22. raise ValueError("Unknown image instance ({})!".format(type(im)))