seg_transforms.py 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943
  1. # coding: utf8
  2. # copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. from .ops import *
  16. import random
  17. import os.path as osp
  18. import numpy as np
  19. from PIL import Image
  20. import cv2
  21. from collections import OrderedDict
  22. class Compose:
  23. """根据数据预处理/增强算子对输入数据进行操作。
  24. 所有操作的输入图像流形状均是[H, W, C],其中H为图像高,W为图像宽,C为图像通道数。
  25. Args:
  26. transforms (list): 数据预处理/增强算子。
  27. Raises:
  28. TypeError: transforms不是list对象
  29. ValueError: transforms元素个数小于1。
  30. """
  31. def __init__(self, transforms):
  32. if not isinstance(transforms, list):
  33. raise TypeError('The transforms must be a list!')
  34. if len(transforms) < 1:
  35. raise ValueError('The length of transforms ' + \
  36. 'must be equal or larger than 1!')
  37. self.transforms = transforms
  38. self.to_rgb = False
  39. def __call__(self, im, im_info=None, label=None):
  40. """
  41. Args:
  42. im (str/np.ndarray): 图像路径/图像np.ndarray数据。
  43. im_info (dict): 存储与图像相关的信息,dict中的字段如下:
  44. - shape_before_resize (tuple): 图像resize之前的大小(h, w)。
  45. - shape_before_padding (tuple): 图像padding之前的大小(h, w)。
  46. label (str/np.ndarray): 标注图像路径/标注图像np.ndarray数据。
  47. Returns:
  48. tuple: 根据网络所需字段所组成的tuple;字段由transforms中的最后一个数据预处理操作决定。
  49. """
  50. if im_info is None:
  51. im_info = dict()
  52. try:
  53. im = cv2.imread(im).astype('float32')
  54. except:
  55. raise ValueError('Can\'t read The image file {}!'.format(im))
  56. if self.to_rgb:
  57. im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
  58. if label is not None:
  59. label = np.asarray(Image.open(label))
  60. for op in self.transforms:
  61. outputs = op(im, im_info, label)
  62. im = outputs[0]
  63. if len(outputs) >= 2:
  64. im_info = outputs[1]
  65. if len(outputs) == 3:
  66. label = outputs[2]
  67. return outputs
  68. class RandomHorizontalFlip:
  69. """以一定的概率对图像进行水平翻转。当存在标注图像时,则同步进行翻转。
  70. Args:
  71. prob (float): 随机水平翻转的概率。默认值为0.5。
  72. """
  73. def __init__(self, prob=0.5):
  74. self.prob = prob
  75. def __call__(self, im, im_info=None, label=None):
  76. """
  77. Args:
  78. im (np.ndarray): 图像np.ndarray数据。
  79. im_info (dict): 存储与图像相关的信息。
  80. label (np.ndarray): 标注图像np.ndarray数据。
  81. Returns:
  82. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  83. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  84. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  85. """
  86. if random.random() < self.prob:
  87. im = horizontal_flip(im)
  88. if label is not None:
  89. label = horizontal_flip(label)
  90. if label is None:
  91. return (im, im_info)
  92. else:
  93. return (im, im_info, label)
  94. class RandomVerticalFlip:
  95. """以一定的概率对图像进行垂直翻转。当存在标注图像时,则同步进行翻转。
  96. Args:
  97. prob (float): 随机垂直翻转的概率。默认值为0.1。
  98. """
  99. def __init__(self, prob=0.1):
  100. self.prob = prob
  101. def __call__(self, im, im_info=None, label=None):
  102. """
  103. Args:
  104. im (np.ndarray): 图像np.ndarray数据。
  105. im_info (dict): 存储与图像相关的信息。
  106. label (np.ndarray): 标注图像np.ndarray数据。
  107. Returns:
  108. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  109. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  110. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  111. """
  112. if random.random() < self.prob:
  113. im = vertical_flip(im)
  114. if label is not None:
  115. label = vertical_flip(label)
  116. if label is None:
  117. return (im, im_info)
  118. else:
  119. return (im, im_info, label)
  120. class Resize:
  121. """调整图像大小(resize),当存在标注图像时,则同步进行处理。
  122. - 当目标大小(target_size)类型为int时,根据插值方式,
  123. 将图像resize为[target_size, target_size]。
  124. - 当目标大小(target_size)类型为list或tuple时,根据插值方式,
  125. 将图像resize为target_size, target_size的输入应为[w, h]或(w, h)。
  126. Args:
  127. target_size (int|list|tuple): 目标大小。
  128. interp (str): resize的插值方式,与opencv的插值方式对应,
  129. 可选的值为['NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4'],默认为"LINEAR"。
  130. Raises:
  131. TypeError: target_size不是int/list/tuple。
  132. ValueError: target_size为list/tuple时元素个数不等于2。
  133. AssertionError: interp的取值不在['NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4']之内。
  134. """
  135. # The interpolation mode
  136. interp_dict = {
  137. 'NEAREST': cv2.INTER_NEAREST,
  138. 'LINEAR': cv2.INTER_LINEAR,
  139. 'CUBIC': cv2.INTER_CUBIC,
  140. 'AREA': cv2.INTER_AREA,
  141. 'LANCZOS4': cv2.INTER_LANCZOS4
  142. }
  143. def __init__(self, target_size, interp='LINEAR'):
  144. self.interp = interp
  145. assert interp in self.interp_dict, "interp should be one of {}".format(
  146. interp_dict.keys())
  147. if isinstance(target_size, list) or isinstance(target_size, tuple):
  148. if len(target_size) != 2:
  149. raise ValueError(
  150. 'when target is list or tuple, it should include 2 elements, but it is {}'
  151. .format(target_size))
  152. elif not isinstance(target_size, int):
  153. raise TypeError(
  154. "Type of target_size is invalid. Must be Integer or List or tuple, now is {}"
  155. .format(type(target_size)))
  156. self.target_size = target_size
  157. def __call__(self, im, im_info=None, label=None):
  158. """
  159. Args:
  160. im (np.ndarray): 图像np.ndarray数据。
  161. im_info (dict): 存储与图像相关的信息。
  162. label (np.ndarray): 标注图像np.ndarray数据。
  163. Returns:
  164. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  165. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  166. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  167. 其中,im_info跟新字段为:
  168. -shape_before_resize (tuple): 保存resize之前图像的形状(h, w)。
  169. Raises:
  170. ZeroDivisionError: im的短边为0。
  171. TypeError: im不是np.ndarray数据。
  172. ValueError: im不是3维nd.ndarray。
  173. """
  174. if im_info is None:
  175. im_info = OrderedDict()
  176. im_info['shape_before_resize'] = im.shape[:2]
  177. if not isinstance(im, np.ndarray):
  178. raise TypeError("ResizeImage: image type is not np.ndarray.")
  179. if len(im.shape) != 3:
  180. raise ValueError('ResizeImage: image is not 3-dimensional.')
  181. im_shape = im.shape
  182. im_size_min = np.min(im_shape[0:2])
  183. im_size_max = np.max(im_shape[0:2])
  184. if float(im_size_min) == 0:
  185. raise ZeroDivisionError('ResizeImage: min size of image is 0')
  186. if isinstance(self.target_size, int):
  187. resize_w = self.target_size
  188. resize_h = self.target_size
  189. else:
  190. resize_w = self.target_size[0]
  191. resize_h = self.target_size[1]
  192. im_scale_x = float(resize_w) / float(im_shape[1])
  193. im_scale_y = float(resize_h) / float(im_shape[0])
  194. im = cv2.resize(
  195. im,
  196. None,
  197. None,
  198. fx=im_scale_x,
  199. fy=im_scale_y,
  200. interpolation=self.interp_dict[self.interp])
  201. if label is not None:
  202. label = cv2.resize(
  203. label,
  204. None,
  205. None,
  206. fx=im_scale_x,
  207. fy=im_scale_y,
  208. interpolation=self.interp_dict['NEAREST'])
  209. if label is None:
  210. return (im, im_info)
  211. else:
  212. return (im, im_info, label)
  213. class ResizeByLong:
  214. """对图像长边resize到固定值,短边按比例进行缩放。当存在标注图像时,则同步进行处理。
  215. Args:
  216. long_size (int): resize后图像的长边大小。
  217. """
  218. def __init__(self, long_size):
  219. self.long_size = long_size
  220. def __call__(self, im, im_info=None, label=None):
  221. """
  222. Args:
  223. im (np.ndarray): 图像np.ndarray数据。
  224. im_info (dict): 存储与图像相关的信息。
  225. label (np.ndarray): 标注图像np.ndarray数据。
  226. Returns:
  227. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  228. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  229. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  230. 其中,im_info新增字段为:
  231. -shape_before_resize (tuple): 保存resize之前图像的形状(h, w)。
  232. """
  233. if im_info is None:
  234. im_info = OrderedDict()
  235. im_info['shape_before_resize'] = im.shape[:2]
  236. im = resize_long(im, self.long_size)
  237. if label is not None:
  238. label = resize_long(label, self.long_size, cv2.INTER_NEAREST)
  239. if label is None:
  240. return (im, im_info)
  241. else:
  242. return (im, im_info, label)
  243. class ResizeRangeScaling:
  244. """对图像长边随机resize到指定范围内,短边按比例进行缩放。当存在标注图像时,则同步进行处理。
  245. Args:
  246. min_value (int): 图像长边resize后的最小值。默认值400。
  247. max_value (int): 图像长边resize后的最大值。默认值600。
  248. Raises:
  249. ValueError: min_value大于max_value
  250. """
  251. def __init__(self, min_value=400, max_value=600):
  252. if min_value > max_value:
  253. raise ValueError('min_value must be less than max_value, '
  254. 'but they are {} and {}.'.format(
  255. min_value, max_value))
  256. self.min_value = min_value
  257. self.max_value = max_value
  258. def __call__(self, im, im_info=None, label=None):
  259. """
  260. Args:
  261. im (np.ndarray): 图像np.ndarray数据。
  262. im_info (dict): 存储与图像相关的信息。
  263. label (np.ndarray): 标注图像np.ndarray数据。
  264. Returns:
  265. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  266. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  267. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  268. """
  269. if self.min_value == self.max_value:
  270. random_size = self.max_value
  271. else:
  272. random_size = int(
  273. np.random.uniform(self.min_value, self.max_value) + 0.5)
  274. im = resize_long(im, random_size, cv2.INTER_LINEAR)
  275. if label is not None:
  276. label = resize_long(label, random_size, cv2.INTER_NEAREST)
  277. if label is None:
  278. return (im, im_info)
  279. else:
  280. return (im, im_info, label)
  281. class ResizeStepScaling:
  282. """对图像按照某一个比例resize,这个比例以scale_step_size为步长
  283. 在[min_scale_factor, max_scale_factor]随机变动。当存在标注图像时,则同步进行处理。
  284. Args:
  285. min_scale_factor(float), resize最小尺度。默认值0.75。
  286. max_scale_factor (float), resize最大尺度。默认值1.25。
  287. scale_step_size (float), resize尺度范围间隔。默认值0.25。
  288. Raises:
  289. ValueError: min_scale_factor大于max_scale_factor
  290. """
  291. def __init__(self,
  292. min_scale_factor=0.75,
  293. max_scale_factor=1.25,
  294. scale_step_size=0.25):
  295. if min_scale_factor > max_scale_factor:
  296. raise ValueError(
  297. 'min_scale_factor must be less than max_scale_factor, '
  298. 'but they are {} and {}.'.format(min_scale_factor,
  299. max_scale_factor))
  300. self.min_scale_factor = min_scale_factor
  301. self.max_scale_factor = max_scale_factor
  302. self.scale_step_size = scale_step_size
  303. def __call__(self, im, im_info=None, label=None):
  304. """
  305. Args:
  306. im (np.ndarray): 图像np.ndarray数据。
  307. im_info (dict): 存储与图像相关的信息。
  308. label (np.ndarray): 标注图像np.ndarray数据。
  309. Returns:
  310. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  311. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  312. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  313. """
  314. if self.min_scale_factor == self.max_scale_factor:
  315. scale_factor = self.min_scale_factor
  316. elif self.scale_step_size == 0:
  317. scale_factor = np.random.uniform(self.min_scale_factor,
  318. self.max_scale_factor)
  319. else:
  320. num_steps = int((self.max_scale_factor - self.min_scale_factor) /
  321. self.scale_step_size + 1)
  322. scale_factors = np.linspace(self.min_scale_factor,
  323. self.max_scale_factor,
  324. num_steps).tolist()
  325. np.random.shuffle(scale_factors)
  326. scale_factor = scale_factors[0]
  327. im = cv2.resize(
  328. im, (0, 0),
  329. fx=scale_factor,
  330. fy=scale_factor,
  331. interpolation=cv2.INTER_LINEAR)
  332. if label is not None:
  333. label = cv2.resize(
  334. label, (0, 0),
  335. fx=scale_factor,
  336. fy=scale_factor,
  337. interpolation=cv2.INTER_NEAREST)
  338. if label is None:
  339. return (im, im_info)
  340. else:
  341. return (im, im_info, label)
  342. class Normalize:
  343. """对图像进行标准化。
  344. 1.尺度缩放到 [0,1]。
  345. 2.对图像进行减均值除以标准差操作。
  346. Args:
  347. mean (list): 图像数据集的均值。默认值[0.5, 0.5, 0.5]。
  348. std (list): 图像数据集的标准差。默认值[0.5, 0.5, 0.5]。
  349. Raises:
  350. ValueError: mean或std不是list对象。std包含0。
  351. """
  352. def __init__(self, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]):
  353. self.mean = mean
  354. self.std = std
  355. if not (isinstance(self.mean, list) and isinstance(self.std, list)):
  356. raise ValueError("{}: input type is invalid.".format(self))
  357. from functools import reduce
  358. if reduce(lambda x, y: x * y, self.std) == 0:
  359. raise ValueError('{}: std is invalid!'.format(self))
  360. def __call__(self, im, im_info=None, label=None):
  361. """
  362. Args:
  363. im (np.ndarray): 图像np.ndarray数据。
  364. im_info (dict): 存储与图像相关的信息。
  365. label (np.ndarray): 标注图像np.ndarray数据。
  366. Returns:
  367. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  368. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  369. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  370. """
  371. mean = np.array(self.mean)[np.newaxis, np.newaxis, :]
  372. std = np.array(self.std)[np.newaxis, np.newaxis, :]
  373. im = normalize(im, mean, std)
  374. if label is None:
  375. return (im, im_info)
  376. else:
  377. return (im, im_info, label)
  378. class Padding:
  379. """对图像或标注图像进行padding,padding方向为右和下。
  380. 根据提供的值对图像或标注图像进行padding操作。
  381. Args:
  382. target_size (int|list|tuple): padding后图像的大小。
  383. im_padding_value (list): 图像padding的值。默认为[127.5, 127.5, 127.5]。
  384. label_padding_value (int): 标注图像padding的值。默认值为255。
  385. Raises:
  386. TypeError: target_size不是int|list|tuple。
  387. ValueError: target_size为list|tuple时元素个数不等于2。
  388. """
  389. def __init__(self,
  390. target_size,
  391. im_padding_value=[127.5, 127.5, 127.5],
  392. label_padding_value=255):
  393. if isinstance(target_size, list) or isinstance(target_size, tuple):
  394. if len(target_size) != 2:
  395. raise ValueError(
  396. 'when target is list or tuple, it should include 2 elements, but it is {}'
  397. .format(target_size))
  398. elif not isinstance(target_size, int):
  399. raise TypeError(
  400. "Type of target_size is invalid. Must be Integer or List or tuple, now is {}"
  401. .format(type(target_size)))
  402. self.target_size = target_size
  403. self.im_padding_value = im_padding_value
  404. self.label_padding_value = label_padding_value
  405. def __call__(self, im, im_info=None, label=None):
  406. """
  407. Args:
  408. im (np.ndarray): 图像np.ndarray数据。
  409. im_info (dict): 存储与图像相关的信息。
  410. label (np.ndarray): 标注图像np.ndarray数据。
  411. Returns:
  412. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  413. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  414. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  415. 其中,im_info新增字段为:
  416. -shape_before_padding (tuple): 保存padding之前图像的形状(h, w)。
  417. Raises:
  418. ValueError: 输入图像im或label的形状大于目标值
  419. """
  420. if im_info is None:
  421. im_info = OrderedDict()
  422. im_info['shape_before_padding'] = im.shape[:2]
  423. im_height, im_width = im.shape[0], im.shape[1]
  424. if isinstance(self.target_size, int):
  425. target_height = self.target_size
  426. target_width = self.target_size
  427. else:
  428. target_height = self.target_size[1]
  429. target_width = self.target_size[0]
  430. pad_height = target_height - im_height
  431. pad_width = target_width - im_width
  432. if pad_height < 0 or pad_width < 0:
  433. raise ValueError(
  434. 'the size of image should be less than target_size, but the size of image ({}, {}), is larger than target_size ({}, {})'
  435. .format(im_width, im_height, target_width, target_height))
  436. else:
  437. im = cv2.copyMakeBorder(
  438. im,
  439. 0,
  440. pad_height,
  441. 0,
  442. pad_width,
  443. cv2.BORDER_CONSTANT,
  444. value=self.im_padding_value)
  445. if label is not None:
  446. label = cv2.copyMakeBorder(
  447. label,
  448. 0,
  449. pad_height,
  450. 0,
  451. pad_width,
  452. cv2.BORDER_CONSTANT,
  453. value=self.label_padding_value)
  454. if label is None:
  455. return (im, im_info)
  456. else:
  457. return (im, im_info, label)
  458. class RandomPaddingCrop:
  459. """对图像和标注图进行随机裁剪,当所需要的裁剪尺寸大于原图时,则进行padding操作。
  460. Args:
  461. crop_size (int|list|tuple): 裁剪图像大小。默认为512。
  462. im_padding_value (list): 图像padding的值。默认为[127.5, 127.5, 127.5]。
  463. label_padding_value (int): 标注图像padding的值。默认值为255。
  464. Raises:
  465. TypeError: crop_size不是int/list/tuple。
  466. ValueError: target_size为list/tuple时元素个数不等于2。
  467. """
  468. def __init__(self,
  469. crop_size=512,
  470. im_padding_value=[127.5, 127.5, 127.5],
  471. label_padding_value=255):
  472. if isinstance(crop_size, list) or isinstance(crop_size, tuple):
  473. if len(crop_size) != 2:
  474. raise ValueError(
  475. 'when crop_size is list or tuple, it should include 2 elements, but it is {}'
  476. .format(crop_size))
  477. elif not isinstance(crop_size, int):
  478. raise TypeError(
  479. "Type of crop_size is invalid. Must be Integer or List or tuple, now is {}"
  480. .format(type(crop_size)))
  481. self.crop_size = crop_size
  482. self.im_padding_value = im_padding_value
  483. self.label_padding_value = label_padding_value
  484. def __call__(self, im, im_info=None, label=None):
  485. """
  486. Args:
  487. im (np.ndarray): 图像np.ndarray数据。
  488. im_info (dict): 存储与图像相关的信息。
  489. label (np.ndarray): 标注图像np.ndarray数据。
  490. Returns:
  491. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  492. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  493. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  494. """
  495. if isinstance(self.crop_size, int):
  496. crop_width = self.crop_size
  497. crop_height = self.crop_size
  498. else:
  499. crop_width = self.crop_size[0]
  500. crop_height = self.crop_size[1]
  501. img_height = im.shape[0]
  502. img_width = im.shape[1]
  503. if img_height == crop_height and img_width == crop_width:
  504. if label is None:
  505. return (im, im_info)
  506. else:
  507. return (im, im_info, label)
  508. else:
  509. pad_height = max(crop_height - img_height, 0)
  510. pad_width = max(crop_width - img_width, 0)
  511. if (pad_height > 0 or pad_width > 0):
  512. im = cv2.copyMakeBorder(
  513. im,
  514. 0,
  515. pad_height,
  516. 0,
  517. pad_width,
  518. cv2.BORDER_CONSTANT,
  519. value=self.im_padding_value)
  520. if label is not None:
  521. label = cv2.copyMakeBorder(
  522. label,
  523. 0,
  524. pad_height,
  525. 0,
  526. pad_width,
  527. cv2.BORDER_CONSTANT,
  528. value=self.label_padding_value)
  529. img_height = im.shape[0]
  530. img_width = im.shape[1]
  531. if crop_height > 0 and crop_width > 0:
  532. h_off = np.random.randint(img_height - crop_height + 1)
  533. w_off = np.random.randint(img_width - crop_width + 1)
  534. im = im[h_off:(crop_height + h_off), w_off:(
  535. w_off + crop_width), :]
  536. if label is not None:
  537. label = label[h_off:(crop_height + h_off), w_off:(
  538. w_off + crop_width)]
  539. if label is None:
  540. return (im, im_info)
  541. else:
  542. return (im, im_info, label)
  543. class RandomBlur:
  544. """以一定的概率对图像进行高斯模糊。
  545. Args:
  546. prob (float): 图像模糊概率。默认为0.1。
  547. """
  548. def __init__(self, prob=0.1):
  549. self.prob = prob
  550. def __call__(self, im, im_info=None, label=None):
  551. """
  552. Args:
  553. im (np.ndarray): 图像np.ndarray数据。
  554. im_info (dict): 存储与图像相关的信息。
  555. label (np.ndarray): 标注图像np.ndarray数据。
  556. Returns:
  557. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  558. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  559. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  560. """
  561. if self.prob <= 0:
  562. n = 0
  563. elif self.prob >= 1:
  564. n = 1
  565. else:
  566. n = int(1.0 / self.prob)
  567. if n > 0:
  568. if np.random.randint(0, n) == 0:
  569. radius = np.random.randint(3, 10)
  570. if radius % 2 != 1:
  571. radius = radius + 1
  572. if radius > 9:
  573. radius = 9
  574. im = cv2.GaussianBlur(im, (radius, radius), 0, 0)
  575. if label is None:
  576. return (im, im_info)
  577. else:
  578. return (im, im_info, label)
  579. class RandomRotate:
  580. """对图像进行随机旋转, 模型训练时的数据增强操作。
  581. 在旋转区间[-rotate_range, rotate_range]内,对图像进行随机旋转,当存在标注图像时,同步进行,
  582. 并对旋转后的图像和标注图像进行相应的padding。
  583. Args:
  584. rotate_range (float): 最大旋转角度。默认为15度。
  585. im_padding_value (list): 图像padding的值。默认为[127.5, 127.5, 127.5]。
  586. label_padding_value (int): 标注图像padding的值。默认为255。
  587. """
  588. def __init__(self,
  589. rotate_range=15,
  590. im_padding_value=[127.5, 127.5, 127.5],
  591. label_padding_value=255):
  592. self.rotate_range = rotate_range
  593. self.im_padding_value = im_padding_value
  594. self.label_padding_value = label_padding_value
  595. def __call__(self, im, im_info=None, label=None):
  596. """
  597. Args:
  598. im (np.ndarray): 图像np.ndarray数据。
  599. im_info (dict): 存储与图像相关的信息。
  600. label (np.ndarray): 标注图像np.ndarray数据。
  601. Returns:
  602. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  603. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  604. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  605. """
  606. if self.rotate_range > 0:
  607. (h, w) = im.shape[:2]
  608. do_rotation = np.random.uniform(-self.rotate_range,
  609. self.rotate_range)
  610. pc = (w // 2, h // 2)
  611. r = cv2.getRotationMatrix2D(pc, do_rotation, 1.0)
  612. cos = np.abs(r[0, 0])
  613. sin = np.abs(r[0, 1])
  614. nw = int((h * sin) + (w * cos))
  615. nh = int((h * cos) + (w * sin))
  616. (cx, cy) = pc
  617. r[0, 2] += (nw / 2) - cx
  618. r[1, 2] += (nh / 2) - cy
  619. dsize = (nw, nh)
  620. im = cv2.warpAffine(
  621. im,
  622. r,
  623. dsize=dsize,
  624. flags=cv2.INTER_LINEAR,
  625. borderMode=cv2.BORDER_CONSTANT,
  626. borderValue=self.im_padding_value)
  627. label = cv2.warpAffine(
  628. label,
  629. r,
  630. dsize=dsize,
  631. flags=cv2.INTER_NEAREST,
  632. borderMode=cv2.BORDER_CONSTANT,
  633. borderValue=self.label_padding_value)
  634. if label is None:
  635. return (im, im_info)
  636. else:
  637. return (im, im_info, label)
  638. class RandomScaleAspect:
  639. """裁剪并resize回原始尺寸的图像和标注图像。
  640. 按照一定的面积比和宽高比对图像进行裁剪,并reszie回原始图像的图像,当存在标注图时,同步进行。
  641. Args:
  642. min_scale (float):裁取图像占原始图像的面积比,取值[0,1],为0时则返回原图。默认为0.5。
  643. aspect_ratio (float): 裁取图像的宽高比范围,非负值,为0时返回原图。默认为0.33。
  644. """
  645. def __init__(self, min_scale=0.5, aspect_ratio=0.33):
  646. self.min_scale = min_scale
  647. self.aspect_ratio = aspect_ratio
  648. def __call__(self, im, im_info=None, label=None):
  649. """
  650. Args:
  651. im (np.ndarray): 图像np.ndarray数据。
  652. im_info (dict): 存储与图像相关的信息。
  653. label (np.ndarray): 标注图像np.ndarray数据。
  654. Returns:
  655. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  656. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  657. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  658. """
  659. if self.min_scale != 0 and self.aspect_ratio != 0:
  660. img_height = im.shape[0]
  661. img_width = im.shape[1]
  662. for i in range(0, 10):
  663. area = img_height * img_width
  664. target_area = area * np.random.uniform(self.min_scale, 1.0)
  665. aspectRatio = np.random.uniform(self.aspect_ratio,
  666. 1.0 / self.aspect_ratio)
  667. dw = int(np.sqrt(target_area * 1.0 * aspectRatio))
  668. dh = int(np.sqrt(target_area * 1.0 / aspectRatio))
  669. if (np.random.randint(10) < 5):
  670. tmp = dw
  671. dw = dh
  672. dh = tmp
  673. if (dh < img_height and dw < img_width):
  674. h1 = np.random.randint(0, img_height - dh)
  675. w1 = np.random.randint(0, img_width - dw)
  676. im = im[h1:(h1 + dh), w1:(w1 + dw), :]
  677. label = label[h1:(h1 + dh), w1:(w1 + dw)]
  678. im = cv2.resize(
  679. im, (img_width, img_height),
  680. interpolation=cv2.INTER_LINEAR)
  681. label = cv2.resize(
  682. label, (img_width, img_height),
  683. interpolation=cv2.INTER_NEAREST)
  684. break
  685. if label is None:
  686. return (im, im_info)
  687. else:
  688. return (im, im_info, label)
  689. class RandomDistort:
  690. """对图像进行随机失真。
  691. 1. 对变换的操作顺序进行随机化操作。
  692. 2. 按照1中的顺序以一定的概率对图像进行随机像素内容变换。
  693. Args:
  694. brightness_range (float): 明亮度因子的范围。默认为0.5。
  695. brightness_prob (float): 随机调整明亮度的概率。默认为0.5。
  696. contrast_range (float): 对比度因子的范围。默认为0.5。
  697. contrast_prob (float): 随机调整对比度的概率。默认为0.5。
  698. saturation_range (float): 饱和度因子的范围。默认为0.5。
  699. saturation_prob (float): 随机调整饱和度的概率。默认为0.5。
  700. hue_range (int): 色调因子的范围。默认为18。
  701. hue_prob (float): 随机调整色调的概率。默认为0.5。
  702. """
  703. def __init__(self,
  704. brightness_range=0.5,
  705. brightness_prob=0.5,
  706. contrast_range=0.5,
  707. contrast_prob=0.5,
  708. saturation_range=0.5,
  709. saturation_prob=0.5,
  710. hue_range=18,
  711. hue_prob=0.5):
  712. self.brightness_range = brightness_range
  713. self.brightness_prob = brightness_prob
  714. self.contrast_range = contrast_range
  715. self.contrast_prob = contrast_prob
  716. self.saturation_range = saturation_range
  717. self.saturation_prob = saturation_prob
  718. self.hue_range = hue_range
  719. self.hue_prob = hue_prob
  720. def __call__(self, im, im_info=None, label=None):
  721. """
  722. Args:
  723. im (np.ndarray): 图像np.ndarray数据。
  724. im_info (dict): 存储与图像相关的信息。
  725. label (np.ndarray): 标注图像np.ndarray数据。
  726. Returns:
  727. tuple: 当label为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  728. 当label不为空时,返回的tuple为(im, im_info, label),分别对应图像np.ndarray数据、
  729. 存储与图像相关信息的字典和标注图像np.ndarray数据。
  730. """
  731. brightness_lower = 1 - self.brightness_range
  732. brightness_upper = 1 + self.brightness_range
  733. contrast_lower = 1 - self.contrast_range
  734. contrast_upper = 1 + self.contrast_range
  735. saturation_lower = 1 - self.saturation_range
  736. saturation_upper = 1 + self.saturation_range
  737. hue_lower = -self.hue_range
  738. hue_upper = self.hue_range
  739. ops = [brightness, contrast, saturation, hue]
  740. random.shuffle(ops)
  741. params_dict = {
  742. 'brightness': {
  743. 'brightness_lower': brightness_lower,
  744. 'brightness_upper': brightness_upper
  745. },
  746. 'contrast': {
  747. 'contrast_lower': contrast_lower,
  748. 'contrast_upper': contrast_upper
  749. },
  750. 'saturation': {
  751. 'saturation_lower': saturation_lower,
  752. 'saturation_upper': saturation_upper
  753. },
  754. 'hue': {
  755. 'hue_lower': hue_lower,
  756. 'hue_upper': hue_upper
  757. }
  758. }
  759. prob_dict = {
  760. 'brightness': self.brightness_prob,
  761. 'contrast': self.contrast_prob,
  762. 'saturation': self.saturation_prob,
  763. 'hue': self.hue_prob
  764. }
  765. im = im.astype('uint8')
  766. im = Image.fromarray(im)
  767. for id in range(4):
  768. params = params_dict[ops[id].__name__]
  769. prob = prob_dict[ops[id].__name__]
  770. params['im'] = im
  771. if np.random.uniform(0, 1) < prob:
  772. im = ops[id](**params)
  773. im = np.asarray(im).astype('float32')
  774. if label is None:
  775. return (im, im_info)
  776. else:
  777. return (im, im_info, label)
  778. class ArrangeSegmenter:
  779. """获取训练/验证/预测所需的信息。
  780. Args:
  781. mode (str): 指定数据用于何种用途,取值范围为['train', 'eval', 'test', 'quant']。
  782. Raises:
  783. ValueError: mode的取值不在['train', 'eval', 'test', 'quant']之内
  784. """
  785. def __init__(self, mode):
  786. if mode not in ['train', 'eval', 'test', 'quant']:
  787. raise ValueError(
  788. "mode should be defined as one of ['train', 'eval', 'test', 'quant']!"
  789. )
  790. self.mode = mode
  791. def __call__(self, im, im_info, label=None):
  792. """
  793. Args:
  794. im (np.ndarray): 图像np.ndarray数据。
  795. im_info (dict): 存储与图像相关的信息。
  796. label (np.ndarray): 标注图像np.ndarray数据。
  797. Returns:
  798. tuple: 当mode为'train'或'eval'时,返回的tuple为(im, label),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  799. 当mode为'test'时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;当mode为
  800. 'quant'时,返回的tuple为(im,),为图像np.ndarray数据。
  801. """
  802. im = permute(im, False)
  803. if self.mode == 'train' or self.mode == 'eval':
  804. label = label[np.newaxis, :, :]
  805. return (im, label)
  806. elif self.mode == 'test':
  807. return (im, im_info)
  808. else:
  809. return (im, )