image_util.py 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232
  1. # Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from io import StringIO
  15. import numpy as np
  16. from PIL import Image
  17. __all__ = []
  18. def resize_image(img, target_size):
  19. """
  20. Resize an image so that the shorter edge has length target_size.
  21. img: the input image to be resized.
  22. target_size: the target resized image size.
  23. """
  24. percent = target_size / float(min(img.size[0], img.size[1]))
  25. resized_size = int(round(img.size[0] * percent)), int(
  26. round(img.size[1] * percent)
  27. )
  28. img = img.resize(resized_size, Image.ANTIALIAS)
  29. return img
  30. def flip(im):
  31. """
  32. Return the flipped image.
  33. Flip an image along the horizontal direction.
  34. im: input image, (K x H x W) ndarrays
  35. """
  36. if len(im.shape) == 3:
  37. return im[:, :, ::-1]
  38. else:
  39. return im[:, ::-1]
  40. def crop_img(im, inner_size, color=True, test=True):
  41. """
  42. Return cropped image.
  43. The size of the cropped image is inner_size * inner_size.
  44. im: (K x H x W) ndarrays
  45. inner_size: the cropped image size.
  46. color: whether it is color image.
  47. test: whether in test mode.
  48. If False, does random cropping and flipping.
  49. If True, crop the center of images.
  50. """
  51. if color:
  52. height, width = max(inner_size, im.shape[1]), max(
  53. inner_size, im.shape[2]
  54. )
  55. padded_im = np.zeros((3, height, width))
  56. startY = (height - im.shape[1]) / 2
  57. startX = (width - im.shape[2]) / 2
  58. endY, endX = startY + im.shape[1], startX + im.shape[2]
  59. padded_im[:, startY:endY, startX:endX] = im
  60. else:
  61. im = im.astype('float32')
  62. height, width = max(inner_size, im.shape[0]), max(
  63. inner_size, im.shape[1]
  64. )
  65. padded_im = np.zeros((height, width))
  66. startY = (height - im.shape[0]) / 2
  67. startX = (width - im.shape[1]) / 2
  68. endY, endX = startY + im.shape[0], startX + im.shape[1]
  69. padded_im[startY:endY, startX:endX] = im
  70. if test:
  71. startY = (height - inner_size) / 2
  72. startX = (width - inner_size) / 2
  73. else:
  74. startY = np.random.randint(0, height - inner_size + 1)
  75. startX = np.random.randint(0, width - inner_size + 1)
  76. endY, endX = startY + inner_size, startX + inner_size
  77. if color:
  78. pic = padded_im[:, startY:endY, startX:endX]
  79. else:
  80. pic = padded_im[startY:endY, startX:endX]
  81. if (not test) and (np.random.randint(2) == 0):
  82. pic = flip(pic)
  83. return pic
  84. def decode_jpeg(jpeg_string):
  85. np_array = np.array(Image.open(StringIO(jpeg_string)))
  86. if len(np_array.shape) == 3:
  87. np_array = np.transpose(np_array, (2, 0, 1))
  88. return np_array
  89. def preprocess_img(im, img_mean, crop_size, is_train, color=True):
  90. """
  91. Does data augmentation for images.
  92. If is_train is false, cropping the center region from the image.
  93. If is_train is true, randomly crop a region from the image,
  94. and random does flipping.
  95. im: (K x H x W) ndarrays
  96. """
  97. im = im.astype('float32')
  98. test = not is_train
  99. pic = crop_img(im, crop_size, color, test)
  100. pic -= img_mean
  101. return pic.flatten()
  102. def load_meta(meta_path, mean_img_size, crop_size, color=True):
  103. """
  104. Return the loaded meta file.
  105. Load the meta image, which is the mean of the images in the dataset.
  106. The mean image is subtracted from every input image so that the expected mean
  107. of each input image is zero.
  108. """
  109. mean = np.load(meta_path)['data_mean']
  110. border = (mean_img_size - crop_size) / 2
  111. if color:
  112. assert mean_img_size * mean_img_size * 3 == mean.shape[0]
  113. mean = mean.reshape(3, mean_img_size, mean_img_size)
  114. mean = mean[
  115. :, border : border + crop_size, border : border + crop_size
  116. ].astype('float32')
  117. else:
  118. assert mean_img_size * mean_img_size == mean.shape[0]
  119. mean = mean.reshape(mean_img_size, mean_img_size)
  120. mean = mean[
  121. border : border + crop_size, border : border + crop_size
  122. ].astype('float32')
  123. return mean
  124. def load_image(img_path, is_color=True):
  125. """
  126. Load image and return.
  127. img_path: image path.
  128. is_color: is color image or not.
  129. """
  130. img = Image.open(img_path)
  131. img.load()
  132. return img
  133. def oversample(img, crop_dims):
  134. """
  135. image : iterable of (H x W x K) ndarrays
  136. crop_dims: (height, width) tuple for the crops.
  137. Returned data contains ten crops of input image, namely,
  138. four corner patches and the center patch as well as their
  139. horizontal reflections.
  140. """
  141. # Dimensions and center.
  142. im_shape = np.array(img[0].shape)
  143. crop_dims = np.array(crop_dims)
  144. im_center = im_shape[:2] / 2.0
  145. # Make crop coordinates
  146. h_indices = (0, im_shape[0] - crop_dims[0])
  147. w_indices = (0, im_shape[1] - crop_dims[1])
  148. crops_ix = np.empty((5, 4), dtype=int)
  149. curr = 0
  150. for i in h_indices:
  151. for j in w_indices:
  152. crops_ix[curr] = (i, j, i + crop_dims[0], j + crop_dims[1])
  153. curr += 1
  154. crops_ix[4] = np.tile(im_center, (1, 2)) + np.concatenate(
  155. [-crop_dims / 2.0, crop_dims / 2.0]
  156. )
  157. crops_ix = np.tile(crops_ix, (2, 1))
  158. # Extract crops
  159. crops = np.empty(
  160. (10 * len(img), crop_dims[0], crop_dims[1], im_shape[-1]),
  161. dtype=np.float32,
  162. )
  163. ix = 0
  164. for im in img:
  165. for crop in crops_ix:
  166. crops[ix] = im[crop[0] : crop[2], crop[1] : crop[3], :]
  167. ix += 1
  168. crops[ix - 5 : ix] = crops[ix - 5 : ix, :, ::-1, :] # flip for mirrors
  169. return crops
  170. class ImageTransformer:
  171. def __init__(
  172. self, transpose=None, channel_swap=None, mean=None, is_color=True
  173. ):
  174. self.is_color = is_color
  175. self.set_transpose(transpose)
  176. self.set_channel_swap(channel_swap)
  177. self.set_mean(mean)
  178. def set_transpose(self, order):
  179. if order is not None:
  180. if self.is_color:
  181. assert 3 == len(order)
  182. self.transpose = order
  183. def set_channel_swap(self, order):
  184. if order is not None:
  185. if self.is_color:
  186. assert 3 == len(order)
  187. self.channel_swap = order
  188. def set_mean(self, mean):
  189. if mean is not None:
  190. # mean value, may be one value per channel
  191. if mean.ndim == 1:
  192. mean = mean[:, np.newaxis, np.newaxis]
  193. else:
  194. # elementwise mean
  195. if self.is_color:
  196. assert len(mean.shape) == 3
  197. self.mean = mean
  198. def transformer(self, data):
  199. if self.transpose is not None:
  200. data = data.transpose(self.transpose)
  201. if self.channel_swap is not None:
  202. data = data[self.channel_swap, :, :]
  203. if self.mean is not None:
  204. data -= self.mean
  205. return data