io_utils.py 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. import os
  2. import os.path as osp
  3. import glob
  4. from pathlib import Path
  5. import cv2
  6. import numpy as np
  7. import json
  8. IMG_EXT = ['.bmp', '.jpg', '.png', '.jpeg']
  9. NP_BOOL_TYPES = (np.bool_, np.bool8)
  10. NP_FLOAT_TYPES = (np.float_, np.float16, np.float32, np.float64)
  11. NP_INT_TYPES = (np.int_, np.int8, np.int16, np.int32, np.int64, np.uint, np.uint8, np.uint16, np.uint32, np.uint64)
  12. # https://stackoverflow.com/questions/26646362/numpy-array-is-not-json-serializable
  13. class NumpyEncoder(json.JSONEncoder):
  14. def default(self, obj):
  15. if isinstance(obj, np.ndarray):
  16. return obj.tolist()
  17. elif isinstance(obj, np.ScalarType):
  18. if isinstance(obj, NP_BOOL_TYPES):
  19. return bool(obj)
  20. elif isinstance(obj, NP_FLOAT_TYPES):
  21. return float(obj)
  22. elif isinstance(obj, NP_INT_TYPES):
  23. return int(obj)
  24. return json.JSONEncoder.default(self, obj)
  25. def find_all_imgs(img_dir, abs_path=False):
  26. imglist = list()
  27. for filep in glob.glob(osp.join(img_dir, "*")):
  28. filename = osp.basename(filep)
  29. file_suffix = Path(filename).suffix
  30. if file_suffix.lower() not in IMG_EXT:
  31. continue
  32. if abs_path:
  33. imglist.append(filep)
  34. else:
  35. imglist.append(filename)
  36. return imglist
  37. def imread(imgpath, read_type=cv2.IMREAD_COLOR):
  38. # img = cv2.imread(imgpath, read_type)
  39. # if img is None:
  40. img = cv2.imdecode(np.fromfile(imgpath, dtype=np.uint8), read_type)
  41. return img
  42. def imwrite(img_path, img, ext='.png'):
  43. suffix = Path(img_path).suffix
  44. if suffix != '':
  45. img_path = img_path.replace(suffix, ext)
  46. else:
  47. img_path += ext
  48. cv2.imencode(ext, img)[1].tofile(img_path)