transforms.py 2.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. import os
  2. import os.path as osp
  3. import imghdr
  4. import gdal
  5. gdal.UseExceptions()
  6. gdal.PushErrorHandler('CPLQuietErrorHandler')
  7. import numpy as np
  8. from PIL import Image
  9. from paddlex.seg import transforms
  10. import paddlex.utils.logging as logging
  11. def read_img(img_path):
  12. img_format = imghdr.what(img_path)
  13. name, ext = osp.splitext(img_path)
  14. if img_format == 'tiff' or ext == '.img':
  15. try:
  16. dataset = gdal.Open(img_path)
  17. except:
  18. logging.error(gdal.GetLastErrorMsg())
  19. if dataset == None:
  20. raise Exception('Can not open', img_path)
  21. im_data = dataset.ReadAsArray()
  22. return im_data.transpose((1, 2, 0))
  23. elif img_format == 'png':
  24. return np.asarray(Image.open(img_path))
  25. elif ext == '.npy':
  26. return np.load(img_path)
  27. else:
  28. raise Exception('Image format {} is not supported!'.format(ext))
  29. def decode_image(im, label):
  30. if isinstance(im, np.ndarray):
  31. if len(im.shape) != 3:
  32. raise Exception(
  33. "im should be 3-dimensions, but now is {}-dimensions".format(
  34. len(im.shape)))
  35. else:
  36. try:
  37. im = read_img(im)
  38. except:
  39. raise ValueError('Can\'t read The image file {}!'.format(im))
  40. im = im.astype('float32')
  41. if label is not None:
  42. if isinstance(label, np.ndarray):
  43. if len(label.shape) != 2:
  44. raise Exception(
  45. "label should be 2-dimensions, but now is {}-dimensions".
  46. format(len(label.shape)))
  47. else:
  48. try:
  49. label = np.asarray(Image.open(label))
  50. except:
  51. ValueError('Can\'t read The label file {}!'.format(label))
  52. im_height, im_width, _ = im.shape
  53. label_height, label_width = label.shape
  54. if im_height != label_height or im_width != label_width:
  55. raise Exception(
  56. "The height or width of the image is not same as the label")
  57. return (im, label)
  58. class Clip(transforms.SegTransform):
  59. """
  60. 对图像上超出一定范围的数据进行裁剪。
  61. Args:
  62. min_val (list): 裁剪的下限,小于min_val的数值均设为min_val. 默认值0.
  63. max_val (list): 裁剪的上限,大于max_val的数值均设为max_val. 默认值255.0.
  64. """
  65. def __init__(self, min_val=[0, 0, 0], max_val=[255.0, 255.0, 255.0]):
  66. self.min_val = min_val
  67. self.max_val = max_val
  68. if not (isinstance(self.min_val, list) and isinstance(self.max_val,
  69. list)):
  70. raise ValueError("{}: input type is invalid.".format(self))
  71. def __call__(self, im, im_info=None, label=None):
  72. for k in range(im.shape[2]):
  73. np.clip(
  74. im[:, :, k], self.min_val[k], self.max_val[k], out=im[:, :, k])
  75. if label is None:
  76. return (im, im_info)
  77. else:
  78. return (im, im_info, label)