operators.py 57 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550
  1. # Copyright (c) 2022 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. import os
  15. import copy
  16. import random
  17. from numbers import Number
  18. from functools import partial
  19. from operator import methodcaller
  20. try:
  21. from collections.abc import Sequence
  22. except Exception:
  23. from collections import Sequence
  24. import numpy as np
  25. import cv2
  26. import imghdr
  27. from PIL import Image
  28. import paddlers
  29. from .functions import normalize, horizontal_flip, permute, vertical_flip, center_crop, is_poly, \
  30. horizontal_flip_poly, horizontal_flip_rle, vertical_flip_poly, vertical_flip_rle, crop_poly, \
  31. crop_rle, expand_poly, expand_rle, resize_poly, resize_rle
  32. __all__ = [
  33. "Compose", "ImgDecoder", "Resize", "RandomResize", "ResizeByShort",
  34. "RandomResizeByShort", "ResizeByLong", "RandomHorizontalFlip",
  35. "RandomVerticalFlip", "Normalize", "CenterCrop", "RandomCrop",
  36. "RandomScaleAspect", "RandomExpand", "Padding", "MixupImage",
  37. "RandomDistort", "RandomBlur", "ArrangeSegmenter", "ArrangeChangeDetector",
  38. "ArrangeClassifier", "ArrangeDetector"
  39. ]
  40. interp_dict = {
  41. 'NEAREST': cv2.INTER_NEAREST,
  42. 'LINEAR': cv2.INTER_LINEAR,
  43. 'CUBIC': cv2.INTER_CUBIC,
  44. 'AREA': cv2.INTER_AREA,
  45. 'LANCZOS4': cv2.INTER_LANCZOS4
  46. }
  47. class Transform(object):
  48. """
  49. Parent class of all data augmentation operations
  50. """
  51. def __init__(self):
  52. pass
  53. def apply_im(self, image):
  54. pass
  55. def apply_mask(self, mask):
  56. pass
  57. def apply_bbox(self, bbox):
  58. pass
  59. def apply_segm(self, segms):
  60. pass
  61. def apply(self, sample):
  62. if 'image' in sample:
  63. sample['image'] = self.apply_im(sample['image'])
  64. else: # image_tx
  65. sample['image'] = self.apply_im(sample['image_t1'])
  66. sample['image2'] = self.apply_im(sample['image_t2'])
  67. if 'mask' in sample:
  68. sample['mask'] = self.apply_mask(sample['mask'])
  69. if 'gt_bbox' in sample:
  70. sample['gt_bbox'] = self.apply_bbox(sample['gt_bbox'])
  71. if 'aux_masks' in sample:
  72. sample['aux_masks'] = list(map(self.apply_mask, sample['aux_masks']))
  73. return sample
  74. def __call__(self, sample):
  75. if isinstance(sample, Sequence):
  76. sample = [self.apply(s) for s in sample]
  77. else:
  78. sample = self.apply(sample)
  79. return sample
  80. class ImgDecoder(Transform):
  81. """
  82. Decode image(s) in input.
  83. Args:
  84. to_rgb (bool, optional): If True, convert input images from BGR format to RGB format. Defaults to True.
  85. """
  86. def __init__(self, to_rgb=True):
  87. super(ImgDecoder, self).__init__()
  88. self.to_rgb = to_rgb
  89. def read_img(self, img_path, input_channel=3):
  90. img_format = imghdr.what(img_path)
  91. name, ext = os.path.splitext(img_path)
  92. if img_format == 'tiff' or ext == '.img':
  93. try:
  94. import gdal
  95. except:
  96. try:
  97. from osgeo import gdal
  98. except:
  99. raise Exception(
  100. "Failed to import gdal! You can try use conda to install gdal"
  101. )
  102. six.reraise(*sys.exc_info())
  103. dataset = gdal.Open(img_path)
  104. if dataset == None:
  105. raise Exception('Can not open', img_path)
  106. im_data = dataset.ReadAsArray()
  107. if im_data.ndim == 2:
  108. im_data = im_data[:, :, np.newaxis]
  109. elif im_data.ndim == 3:
  110. im_data = im_data.transpose((1, 2, 0))
  111. return im_data
  112. elif img_format in ['jpeg', 'bmp', 'png', 'jpg']:
  113. if input_channel == 3:
  114. return cv2.imread(img_path, cv2.IMREAD_ANYDEPTH |
  115. cv2.IMREAD_ANYCOLOR | cv2.IMREAD_COLOR)
  116. else:
  117. return cv2.imread(img_path, cv2.IMREAD_ANYDEPTH |
  118. cv2.IMREAD_ANYCOLOR)
  119. elif ext == '.npy':
  120. return np.load(img_path)
  121. else:
  122. raise Exception('Image format {} is not supported!'.format(ext))
  123. def apply_im(self, im_path):
  124. if isinstance(im_path, str):
  125. try:
  126. image = self.read_img(im_path)
  127. except:
  128. raise ValueError('Cannot read the image file {}!'.format(
  129. im_path))
  130. else:
  131. image = im_path
  132. if self.to_rgb and image.shape[-1] == 3:
  133. image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  134. return image
  135. def apply_mask(self, mask):
  136. try:
  137. mask = np.asarray(Image.open(mask))
  138. except:
  139. raise ValueError("Cannot read the mask file {}!".format(mask))
  140. if len(mask.shape) != 2:
  141. raise Exception(
  142. "Mask should be a 1-channel image, but recevied is a {}-channel image.".
  143. format(mask.shape[2]))
  144. return mask
  145. def apply(self, sample):
  146. """
  147. Args:
  148. sample (dict): Input sample, containing 'image' at least.
  149. Returns:
  150. dict: Decoded sample.
  151. """
  152. if 'image' in sample:
  153. sample['image'] = self.apply_im(sample['image'])
  154. else: # image_tx
  155. sample['image'] = self.apply_im(sample['image_t1'])
  156. sample['image2'] = self.apply_im(sample['image_t2'])
  157. if 'mask' in sample:
  158. sample['mask'] = self.apply_mask(sample['mask'])
  159. im_height, im_width, _ = sample['image'].shape
  160. se_height, se_width = sample['mask'].shape
  161. if im_height != se_height or im_width != se_width:
  162. raise Exception(
  163. "The height or width of the im is not same as the mask")
  164. if 'aux_masks' in sample:
  165. sample['aux_masks'] = list(map(self.apply_mask, sample['aux_masks']))
  166. # TODO: check the shape of auxiliary masks
  167. sample['im_shape'] = np.array(
  168. sample['image'].shape[:2], dtype=np.float32)
  169. sample['scale_factor'] = np.array([1., 1.], dtype=np.float32)
  170. return sample
  171. class Compose(Transform):
  172. """
  173. Apply a series of data augmentation to the input.
  174. All input images are in Height-Width-Channel ([H, W, C]) format.
  175. Args:
  176. transforms (List[paddlers.transforms.Transform]): List of data preprocess or augmentations.
  177. Raises:
  178. TypeError: Invalid type of transforms.
  179. ValueError: Invalid length of transforms.
  180. """
  181. def __init__(self, transforms):
  182. super(Compose, self).__init__()
  183. if not isinstance(transforms, list):
  184. raise TypeError(
  185. 'Type of transforms is invalid. Must be List, but received is {}'
  186. .format(type(transforms)))
  187. if len(transforms) < 1:
  188. raise ValueError(
  189. 'Length of transforms must not be less than 1, but received is {}'
  190. .format(len(transforms)))
  191. self.transforms = transforms
  192. self.decode_image = ImgDecoder()
  193. self.arrange_outputs = None
  194. self.apply_im_only = False
  195. def __call__(self, sample):
  196. if self.apply_im_only:
  197. if 'mask' in sample:
  198. mask_backup = copy.deepcopy(sample['mask'])
  199. del sample['mask']
  200. if 'aux_masks' in sample:
  201. aux_masks = copy.deepcopy(sample['aux_masks'])
  202. sample = self.decode_image(sample)
  203. for op in self.transforms:
  204. # skip batch transforms amd mixup
  205. if isinstance(op, (paddlers.transforms.BatchRandomResize,
  206. paddlers.transforms.BatchRandomResizeByShort,
  207. MixupImage)):
  208. continue
  209. sample = op(sample)
  210. if self.arrange_outputs is not None:
  211. if self.apply_im_only:
  212. sample['mask'] = mask_backup
  213. if 'aux_masks' in locals():
  214. sample['aux_masks'] = aux_masks
  215. sample = self.arrange_outputs(sample)
  216. return sample
  217. class Resize(Transform):
  218. """
  219. Resize input.
  220. - If target_size is an int, resize the image(s) to (target_size, target_size).
  221. - If target_size is a list or tuple, resize the image(s) to target_size.
  222. Attention: If interp is 'RANDOM', the interpolation method will be chose randomly.
  223. Args:
  224. target_size (int, List[int] or Tuple[int]): Target size. If int, the height and width share the same target_size.
  225. Otherwise, target_size represents [target height, target width].
  226. interp ({'NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4', 'RANDOM'}, optional):
  227. Interpolation method of resize. Defaults to 'LINEAR'.
  228. keep_ratio (bool): the resize scale of width/height is same and width/height after resized is not greater
  229. than target width/height. Defaults to False.
  230. Raises:
  231. TypeError: Invalid type of target_size.
  232. ValueError: Invalid interpolation method.
  233. """
  234. def __init__(self, target_size, interp='LINEAR', keep_ratio=False):
  235. super(Resize, self).__init__()
  236. if not (interp == "RANDOM" or interp in interp_dict):
  237. raise ValueError("interp should be one of {}".format(
  238. interp_dict.keys()))
  239. if isinstance(target_size, int):
  240. target_size = (target_size, target_size)
  241. else:
  242. if not (isinstance(target_size,
  243. (list, tuple)) and len(target_size) == 2):
  244. raise TypeError(
  245. "target_size should be an int or a list of length 2, but received {}".
  246. format(target_size))
  247. # (height, width)
  248. self.target_size = target_size
  249. self.interp = interp
  250. self.keep_ratio = keep_ratio
  251. def apply_im(self, image, interp, target_size):
  252. flag = image.shape[2] == 1
  253. image = cv2.resize(image, target_size, interpolation=interp)
  254. if flag:
  255. image = image[:, :, np.newaxis]
  256. return image
  257. def apply_mask(self, mask, target_size):
  258. mask = cv2.resize(mask, target_size, interpolation=cv2.INTER_NEAREST)
  259. return mask
  260. def apply_bbox(self, bbox, scale, target_size):
  261. im_scale_x, im_scale_y = scale
  262. bbox[:, 0::2] *= im_scale_x
  263. bbox[:, 1::2] *= im_scale_y
  264. bbox[:, 0::2] = np.clip(bbox[:, 0::2], 0, target_size[0])
  265. bbox[:, 1::2] = np.clip(bbox[:, 1::2], 0, target_size[1])
  266. return bbox
  267. def apply_segm(self, segms, im_size, scale):
  268. im_h, im_w = im_size
  269. im_scale_x, im_scale_y = scale
  270. resized_segms = []
  271. for segm in segms:
  272. if is_poly(segm):
  273. # Polygon format
  274. resized_segms.append([
  275. resize_poly(poly, im_scale_x, im_scale_y) for poly in segm
  276. ])
  277. else:
  278. # RLE format
  279. resized_segms.append(
  280. resize_rle(segm, im_h, im_w, im_scale_x, im_scale_y))
  281. return resized_segms
  282. def apply(self, sample):
  283. if self.interp == "RANDOM":
  284. interp = random.choice(list(interp_dict.values()))
  285. else:
  286. interp = interp_dict[self.interp]
  287. im_h, im_w = sample['image'].shape[:2]
  288. im_scale_y = self.target_size[0] / im_h
  289. im_scale_x = self.target_size[1] / im_w
  290. target_size = (self.target_size[1], self.target_size[0])
  291. if self.keep_ratio:
  292. scale = min(im_scale_y, im_scale_x)
  293. target_w = int(round(im_w * scale))
  294. target_h = int(round(im_h * scale))
  295. target_size = (target_w, target_h)
  296. im_scale_y = target_h / im_h
  297. im_scale_x = target_w / im_w
  298. sample['image'] = self.apply_im(sample['image'], interp, target_size)
  299. if 'image2' in sample:
  300. sample['image2'] = self.apply_im(sample['image2'], interp, target_size)
  301. if 'mask' in sample:
  302. sample['mask'] = self.apply_mask(sample['mask'], target_size)
  303. if 'aux_masks' in sample:
  304. sample['aux_masks'] = list(map(partial(self.apply_mask, target_size=target_size), sample['aux_masks']))
  305. if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
  306. sample['gt_bbox'] = self.apply_bbox(
  307. sample['gt_bbox'], [im_scale_x, im_scale_y], target_size)
  308. if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
  309. sample['gt_poly'] = self.apply_segm(
  310. sample['gt_poly'], [im_h, im_w], [im_scale_x, im_scale_y])
  311. sample['im_shape'] = np.asarray(
  312. sample['image'].shape[:2], dtype=np.float32)
  313. if 'scale_factor' in sample:
  314. scale_factor = sample['scale_factor']
  315. sample['scale_factor'] = np.asarray(
  316. [scale_factor[0] * im_scale_y, scale_factor[1] * im_scale_x],
  317. dtype=np.float32)
  318. return sample
  319. class RandomResize(Transform):
  320. """
  321. Resize input to random sizes.
  322. Attention: If interp is 'RANDOM', the interpolation method will be chose randomly.
  323. Args:
  324. target_sizes (List[int], List[list or tuple] or Tuple[list or tuple]):
  325. Multiple target sizes, each target size is an int or list/tuple.
  326. interp ({'NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4', 'RANDOM'}, optional):
  327. Interpolation method of resize. Defaults to 'LINEAR'.
  328. Raises:
  329. TypeError: Invalid type of target_size.
  330. ValueError: Invalid interpolation method.
  331. See Also:
  332. Resize input to a specific size.
  333. """
  334. def __init__(self, target_sizes, interp='LINEAR'):
  335. super(RandomResize, self).__init__()
  336. if not (interp == "RANDOM" or interp in interp_dict):
  337. raise ValueError("interp should be one of {}".format(
  338. interp_dict.keys()))
  339. self.interp = interp
  340. assert isinstance(target_sizes, list), \
  341. "target_size must be List"
  342. for i, item in enumerate(target_sizes):
  343. if isinstance(item, int):
  344. target_sizes[i] = (item, item)
  345. self.target_size = target_sizes
  346. def apply(self, sample):
  347. height, width = random.choice(self.target_size)
  348. resizer = Resize((height, width), interp=self.interp)
  349. sample = resizer(sample)
  350. return sample
  351. class ResizeByShort(Transform):
  352. """
  353. Resize input with keeping the aspect ratio.
  354. Attention: If interp is 'RANDOM', the interpolation method will be chose randomly.
  355. Args:
  356. short_size (int): Target size of the shorter side of the image(s).
  357. max_size (int, optional): The upper bound of longer side of the image(s). If max_size is -1, no upper bound is applied. Defaults to -1.
  358. interp ({'NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4', 'RANDOM'}, optional): Interpolation method of resize. Defaults to 'LINEAR'.
  359. Raises:
  360. ValueError: Invalid interpolation method.
  361. """
  362. def __init__(self, short_size=256, max_size=-1, interp='LINEAR'):
  363. if not (interp == "RANDOM" or interp in interp_dict):
  364. raise ValueError("interp should be one of {}".format(
  365. interp_dict.keys()))
  366. super(ResizeByShort, self).__init__()
  367. self.short_size = short_size
  368. self.max_size = max_size
  369. self.interp = interp
  370. def apply(self, sample):
  371. im_h, im_w = sample['image'].shape[:2]
  372. im_short_size = min(im_h, im_w)
  373. im_long_size = max(im_h, im_w)
  374. scale = float(self.short_size) / float(im_short_size)
  375. if 0 < self.max_size < np.round(scale * im_long_size):
  376. scale = float(self.max_size) / float(im_long_size)
  377. target_w = int(round(im_w * scale))
  378. target_h = int(round(im_h * scale))
  379. sample = Resize(
  380. target_size=(target_h, target_w), interp=self.interp)(sample)
  381. return sample
  382. class RandomResizeByShort(Transform):
  383. """
  384. Resize input to random sizes with keeping the aspect ratio.
  385. Attention: If interp is 'RANDOM', the interpolation method will be chose randomly.
  386. Args:
  387. short_sizes (List[int]): Target size of the shorter side of the image(s).
  388. max_size (int, optional): The upper bound of longer side of the image(s). If max_size is -1, no upper bound is applied. Defaults to -1.
  389. interp ({'NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4', 'RANDOM'}, optional): Interpolation method of resize. Defaults to 'LINEAR'.
  390. Raises:
  391. TypeError: Invalid type of target_size.
  392. ValueError: Invalid interpolation method.
  393. See Also:
  394. ResizeByShort: Resize image(s) in input with keeping the aspect ratio.
  395. """
  396. def __init__(self, short_sizes, max_size=-1, interp='LINEAR'):
  397. super(RandomResizeByShort, self).__init__()
  398. if not (interp == "RANDOM" or interp in interp_dict):
  399. raise ValueError("interp should be one of {}".format(
  400. interp_dict.keys()))
  401. self.interp = interp
  402. assert isinstance(short_sizes, list), \
  403. "short_sizes must be List"
  404. self.short_sizes = short_sizes
  405. self.max_size = max_size
  406. def apply(self, sample):
  407. short_size = random.choice(self.short_sizes)
  408. resizer = ResizeByShort(
  409. short_size=short_size, max_size=self.max_size, interp=self.interp)
  410. sample = resizer(sample)
  411. return sample
  412. class ResizeByLong(Transform):
  413. def __init__(self, long_size=256, interp='LINEAR'):
  414. super(ResizeByLong, self).__init__()
  415. self.long_size = long_size
  416. self.interp = interp
  417. def apply(self, sample):
  418. im_h, im_w = sample['image'].shape[:2]
  419. im_long_size = max(im_h, im_w)
  420. scale = float(self.long_size) / float(im_long_size)
  421. target_h = int(round(im_h * scale))
  422. target_w = int(round(im_w * scale))
  423. sample = Resize(
  424. target_size=(target_h, target_w), interp=self.interp)(sample)
  425. return sample
  426. class RandomHorizontalFlip(Transform):
  427. """
  428. Randomly flip the input horizontally.
  429. Args:
  430. prob(float, optional): Probability of flipping the input. Defaults to .5.
  431. """
  432. def __init__(self, prob=0.5):
  433. super(RandomHorizontalFlip, self).__init__()
  434. self.prob = prob
  435. def apply_im(self, image):
  436. image = horizontal_flip(image)
  437. return image
  438. def apply_mask(self, mask):
  439. mask = horizontal_flip(mask)
  440. return mask
  441. def apply_bbox(self, bbox, width):
  442. oldx1 = bbox[:, 0].copy()
  443. oldx2 = bbox[:, 2].copy()
  444. bbox[:, 0] = width - oldx2
  445. bbox[:, 2] = width - oldx1
  446. return bbox
  447. def apply_segm(self, segms, height, width):
  448. flipped_segms = []
  449. for segm in segms:
  450. if is_poly(segm):
  451. # Polygon format
  452. flipped_segms.append(
  453. [horizontal_flip_poly(poly, width) for poly in segm])
  454. else:
  455. # RLE format
  456. flipped_segms.append(horizontal_flip_rle(segm, height, width))
  457. return flipped_segms
  458. def apply(self, sample):
  459. if random.random() < self.prob:
  460. im_h, im_w = sample['image'].shape[:2]
  461. sample['image'] = self.apply_im(sample['image'])
  462. if 'image2' in sample:
  463. sample['image2'] = self.apply_im(sample['image2'])
  464. if 'mask' in sample:
  465. sample['mask'] = self.apply_mask(sample['mask'])
  466. if 'aux_masks' in sample:
  467. sample['aux_masks'] = list(map(self.apply_mask, sample['aux_masks']))
  468. if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
  469. sample['gt_bbox'] = self.apply_bbox(sample['gt_bbox'], im_w)
  470. if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
  471. sample['gt_poly'] = self.apply_segm(sample['gt_poly'], im_h,
  472. im_w)
  473. return sample
  474. class RandomVerticalFlip(Transform):
  475. """
  476. Randomly flip the input vertically.
  477. Args:
  478. prob(float, optional): Probability of flipping the input. Defaults to .5.
  479. """
  480. def __init__(self, prob=0.5):
  481. super(RandomVerticalFlip, self).__init__()
  482. self.prob = prob
  483. def apply_im(self, image):
  484. image = vertical_flip(image)
  485. return image
  486. def apply_mask(self, mask):
  487. mask = vertical_flip(mask)
  488. return mask
  489. def apply_bbox(self, bbox, height):
  490. oldy1 = bbox[:, 1].copy()
  491. oldy2 = bbox[:, 3].copy()
  492. bbox[:, 0] = height - oldy2
  493. bbox[:, 2] = height - oldy1
  494. return bbox
  495. def apply_segm(self, segms, height, width):
  496. flipped_segms = []
  497. for segm in segms:
  498. if is_poly(segm):
  499. # Polygon format
  500. flipped_segms.append(
  501. [vertical_flip_poly(poly, height) for poly in segm])
  502. else:
  503. # RLE format
  504. flipped_segms.append(vertical_flip_rle(segm, height, width))
  505. return flipped_segms
  506. def apply(self, sample):
  507. if random.random() < self.prob:
  508. im_h, im_w = sample['image'].shape[:2]
  509. sample['image'] = self.apply_im(sample['image'])
  510. if 'image2' in sample:
  511. sample['image2'] = self.apply_im(sample['image2'])
  512. if 'mask' in sample:
  513. sample['mask'] = self.apply_mask(sample['mask'])
  514. if 'aux_masks' in sample:
  515. sample['aux_masks'] = list(map(self.apply_mask, sample['aux_masks']))
  516. if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
  517. sample['gt_bbox'] = self.apply_bbox(sample['gt_bbox'], im_h)
  518. if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
  519. sample['gt_poly'] = self.apply_segm(sample['gt_poly'], im_h,
  520. im_w)
  521. return sample
  522. class Normalize(Transform):
  523. """
  524. Apply min-max normalization to the image(s) in input.
  525. 1. im = (im - min_value) * 1 / (max_value - min_value)
  526. 2. im = im - mean
  527. 3. im = im / std
  528. Args:
  529. mean(List[float] or Tuple[float], optional): Mean of input image(s). Defaults to [0.485, 0.456, 0.406].
  530. std(List[float] or Tuple[float], optional): Standard deviation of input image(s). Defaults to [0.229, 0.224, 0.225].
  531. min_val(List[float] or Tuple[float], optional): Minimum value of input image(s). Defaults to [0, 0, 0, ].
  532. max_val(List[float] or Tuple[float], optional): Max value of input image(s). Defaults to [255., 255., 255.].
  533. is_scale(bool, optional): If True, the image pixel values will be divided by 255.
  534. """
  535. def __init__(self,
  536. mean=[0.485, 0.456, 0.406],
  537. std=[0.229, 0.224, 0.225],
  538. min_val=None,
  539. max_val=None,
  540. is_scale=True):
  541. super(Normalize, self).__init__()
  542. channel = len(mean)
  543. if min_val is None:
  544. min_val = [0] * channel
  545. if max_val is None:
  546. max_val = [255.] * channel
  547. from functools import reduce
  548. if reduce(lambda x, y: x * y, std) == 0:
  549. raise ValueError(
  550. 'Std should not have 0, but received is {}'.format(std))
  551. if is_scale:
  552. if reduce(lambda x, y: x * y,
  553. [a - b for a, b in zip(max_val, min_val)]) == 0:
  554. raise ValueError(
  555. '(max_val - min_val) should not have 0, but received is {}'.
  556. format((np.asarray(max_val) - np.asarray(min_val)).tolist(
  557. )))
  558. self.mean = mean
  559. self.std = std
  560. self.min_val = min_val
  561. self.max_val = max_val
  562. self.is_scale = is_scale
  563. def apply_im(self, image):
  564. image = image.astype(np.float32)
  565. mean = np.asarray(
  566. self.mean, dtype=np.float32)[np.newaxis, np.newaxis, :]
  567. std = np.asarray(self.std, dtype=np.float32)[np.newaxis, np.newaxis, :]
  568. image = normalize(image, mean, std, self.min_val, self.max_val)
  569. return image
  570. def apply(self, sample):
  571. sample['image'] = self.apply_im(sample['image'])
  572. if 'image2' in sample:
  573. sample['image2'] = self.apply_im(sample['image2'])
  574. return sample
  575. class CenterCrop(Transform):
  576. """
  577. Crop the input at the center.
  578. 1. Locate the center of the image.
  579. 2. Crop the sample.
  580. Args:
  581. crop_size(int, optional): target size of the cropped image(s). Defaults to 224.
  582. """
  583. def __init__(self, crop_size=224):
  584. super(CenterCrop, self).__init__()
  585. self.crop_size = crop_size
  586. def apply_im(self, image):
  587. image = center_crop(image, self.crop_size)
  588. return image
  589. def apply_mask(self, mask):
  590. mask = center_crop(mask, self.crop_size)
  591. return mask
  592. def apply(self, sample):
  593. sample['image'] = self.apply_im(sample['image'])
  594. if 'image2' in sample:
  595. sample['image2'] = self.apply_im(sample['image2'])
  596. if 'mask' in sample:
  597. sample['mask'] = self.apply_mask(sample['mask'])
  598. if 'aux_masks' in sample:
  599. sample['aux_masks'] = list(map(self.apply_mask, sample['aux_masks']))
  600. return sample
  601. class RandomCrop(Transform):
  602. """
  603. Randomly crop the input.
  604. 1. Compute the height and width of cropped area according to aspect_ratio and scaling.
  605. 2. Locate the upper left corner of cropped area randomly.
  606. 3. Crop the image(s).
  607. 4. Resize the cropped area to crop_size by crop_size.
  608. Args:
  609. crop_size(int, List[int] or Tuple[int]): Target size of the cropped area. If None, the cropped area will not be
  610. resized. Defaults to None.
  611. aspect_ratio (List[float], optional): Aspect ratio of cropped region in [min, max] format. Defaults to [.5, 2.].
  612. thresholds (List[float], optional): Iou thresholds to decide a valid bbox crop.
  613. Defaults to [.0, .1, .3, .5, .7, .9].
  614. scaling (List[float], optional): Ratio between the cropped region and the original image in [min, max] format.
  615. Defaults to [.3, 1.].
  616. num_attempts (int, optional): The number of tries before giving up. Defaults to 50.
  617. allow_no_crop (bool, optional): Whether returning without doing crop is allowed. Defaults to True.
  618. cover_all_box (bool, optional): Whether to ensure all bboxes are covered in the final crop. Defaults to False.
  619. """
  620. def __init__(self,
  621. crop_size=None,
  622. aspect_ratio=[.5, 2.],
  623. thresholds=[.0, .1, .3, .5, .7, .9],
  624. scaling=[.3, 1.],
  625. num_attempts=50,
  626. allow_no_crop=True,
  627. cover_all_box=False):
  628. super(RandomCrop, self).__init__()
  629. self.crop_size = crop_size
  630. self.aspect_ratio = aspect_ratio
  631. self.thresholds = thresholds
  632. self.scaling = scaling
  633. self.num_attempts = num_attempts
  634. self.allow_no_crop = allow_no_crop
  635. self.cover_all_box = cover_all_box
  636. def _generate_crop_info(self, sample):
  637. im_h, im_w = sample['image'].shape[:2]
  638. if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
  639. thresholds = self.thresholds
  640. if self.allow_no_crop:
  641. thresholds.append('no_crop')
  642. np.random.shuffle(thresholds)
  643. for thresh in thresholds:
  644. if thresh == 'no_crop':
  645. return None
  646. for i in range(self.num_attempts):
  647. crop_box = self._get_crop_box(im_h, im_w)
  648. if crop_box is None:
  649. continue
  650. iou = self._iou_matrix(
  651. sample['gt_bbox'],
  652. np.array(
  653. [crop_box], dtype=np.float32))
  654. if iou.max() < thresh:
  655. continue
  656. if self.cover_all_box and iou.min() < thresh:
  657. continue
  658. cropped_box, valid_ids = self._crop_box_with_center_constraint(
  659. sample['gt_bbox'],
  660. np.array(
  661. crop_box, dtype=np.float32))
  662. if valid_ids.size > 0:
  663. return crop_box, cropped_box, valid_ids
  664. else:
  665. for i in range(self.num_attempts):
  666. crop_box = self._get_crop_box(im_h, im_w)
  667. if crop_box is None:
  668. continue
  669. return crop_box, None, None
  670. return None
  671. def _get_crop_box(self, im_h, im_w):
  672. scale = np.random.uniform(*self.scaling)
  673. if self.aspect_ratio is not None:
  674. min_ar, max_ar = self.aspect_ratio
  675. aspect_ratio = np.random.uniform(
  676. max(min_ar, scale**2), min(max_ar, scale**-2))
  677. h_scale = scale / np.sqrt(aspect_ratio)
  678. w_scale = scale * np.sqrt(aspect_ratio)
  679. else:
  680. h_scale = np.random.uniform(*self.scaling)
  681. w_scale = np.random.uniform(*self.scaling)
  682. crop_h = im_h * h_scale
  683. crop_w = im_w * w_scale
  684. if self.aspect_ratio is None:
  685. if crop_h / crop_w < 0.5 or crop_h / crop_w > 2.0:
  686. return None
  687. crop_h = int(crop_h)
  688. crop_w = int(crop_w)
  689. crop_y = np.random.randint(0, im_h - crop_h)
  690. crop_x = np.random.randint(0, im_w - crop_w)
  691. return [crop_x, crop_y, crop_x + crop_w, crop_y + crop_h]
  692. def _iou_matrix(self, a, b):
  693. tl_i = np.maximum(a[:, np.newaxis, :2], b[:, :2])
  694. br_i = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
  695. area_i = np.prod(br_i - tl_i, axis=2) * (tl_i < br_i).all(axis=2)
  696. area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
  697. area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)
  698. area_o = (area_a[:, np.newaxis] + area_b - area_i)
  699. return area_i / (area_o + 1e-10)
  700. def _crop_box_with_center_constraint(self, box, crop):
  701. cropped_box = box.copy()
  702. cropped_box[:, :2] = np.maximum(box[:, :2], crop[:2])
  703. cropped_box[:, 2:] = np.minimum(box[:, 2:], crop[2:])
  704. cropped_box[:, :2] -= crop[:2]
  705. cropped_box[:, 2:] -= crop[:2]
  706. centers = (box[:, :2] + box[:, 2:]) / 2
  707. valid = np.logical_and(crop[:2] <= centers,
  708. centers < crop[2:]).all(axis=1)
  709. valid = np.logical_and(
  710. valid, (cropped_box[:, :2] < cropped_box[:, 2:]).all(axis=1))
  711. return cropped_box, np.where(valid)[0]
  712. def _crop_segm(self, segms, valid_ids, crop, height, width):
  713. crop_segms = []
  714. for id in valid_ids:
  715. segm = segms[id]
  716. if is_poly(segm):
  717. # Polygon format
  718. crop_segms.append(crop_poly(segm, crop))
  719. else:
  720. # RLE format
  721. crop_segms.append(crop_rle(segm, crop, height, width))
  722. return crop_segms
  723. def apply_im(self, image, crop):
  724. x1, y1, x2, y2 = crop
  725. return image[y1:y2, x1:x2, :]
  726. def apply_mask(self, mask, crop):
  727. x1, y1, x2, y2 = crop
  728. return mask[y1:y2, x1:x2, ...]
  729. def apply(self, sample):
  730. crop_info = self._generate_crop_info(sample)
  731. if crop_info is not None:
  732. crop_box, cropped_box, valid_ids = crop_info
  733. im_h, im_w = sample['image'].shape[:2]
  734. sample['image'] = self.apply_im(sample['image'], crop_box)
  735. if 'image2' in sample:
  736. sample['image2'] = self.apply_im(sample['image2'], crop_box)
  737. if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
  738. crop_polys = self._crop_segm(
  739. sample['gt_poly'],
  740. valid_ids,
  741. np.array(
  742. crop_box, dtype=np.int64),
  743. im_h,
  744. im_w)
  745. if [] in crop_polys:
  746. delete_id = list()
  747. valid_polys = list()
  748. for idx, poly in enumerate(crop_polys):
  749. if not crop_poly:
  750. delete_id.append(idx)
  751. else:
  752. valid_polys.append(poly)
  753. valid_ids = np.delete(valid_ids, delete_id)
  754. if not valid_polys:
  755. return sample
  756. sample['gt_poly'] = valid_polys
  757. else:
  758. sample['gt_poly'] = crop_polys
  759. if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
  760. sample['gt_bbox'] = np.take(cropped_box, valid_ids, axis=0)
  761. sample['gt_class'] = np.take(
  762. sample['gt_class'], valid_ids, axis=0)
  763. if 'gt_score' in sample:
  764. sample['gt_score'] = np.take(
  765. sample['gt_score'], valid_ids, axis=0)
  766. if 'is_crowd' in sample:
  767. sample['is_crowd'] = np.take(
  768. sample['is_crowd'], valid_ids, axis=0)
  769. if 'mask' in sample:
  770. sample['mask'] = self.apply_mask(sample['mask'], crop_box)
  771. if 'aux_masks' in sample:
  772. sample['aux_masks'] = list(map(partial(self.apply_mask, crop=crop_box), sample['aux_masks']))
  773. if self.crop_size is not None:
  774. sample = Resize(self.crop_size)(sample)
  775. return sample
  776. class RandomScaleAspect(Transform):
  777. """
  778. Crop input image(s) and resize back to original sizes.
  779. Args:
  780. min_scale (float): Minimum ratio between the cropped region and the original image.
  781. If 0, image(s) will not be cropped. Defaults to .5.
  782. aspect_ratio (float): Aspect ratio of cropped region. Defaults to .33.
  783. """
  784. def __init__(self, min_scale=0.5, aspect_ratio=0.33):
  785. super(RandomScaleAspect, self).__init__()
  786. self.min_scale = min_scale
  787. self.aspect_ratio = aspect_ratio
  788. def apply(self, sample):
  789. if self.min_scale != 0 and self.aspect_ratio != 0:
  790. img_height, img_width = sample['image'].shape[:2]
  791. sample = RandomCrop(
  792. crop_size=(img_height, img_width),
  793. aspect_ratio=[self.aspect_ratio, 1. / self.aspect_ratio],
  794. scaling=[self.min_scale, 1.],
  795. num_attempts=10,
  796. allow_no_crop=False)(sample)
  797. return sample
  798. class RandomExpand(Transform):
  799. """
  800. Randomly expand the input by padding according to random offsets.
  801. Args:
  802. upper_ratio(float, optional): The maximum ratio to which the original image is expanded. Defaults to 4..
  803. prob(float, optional): The probability of apply expanding. Defaults to .5.
  804. im_padding_value(List[float] or Tuple[float], optional): RGB filling value for the image. Defaults to (127.5, 127.5, 127.5).
  805. label_padding_value(int, optional): Filling value for the mask. Defaults to 255.
  806. See Also:
  807. paddlers.transforms.Padding
  808. """
  809. def __init__(self,
  810. upper_ratio=4.,
  811. prob=.5,
  812. im_padding_value=127.5,
  813. label_padding_value=255):
  814. super(RandomExpand, self).__init__()
  815. assert upper_ratio > 1.01, "expand ratio must be larger than 1.01"
  816. self.upper_ratio = upper_ratio
  817. self.prob = prob
  818. assert isinstance(im_padding_value, (Number, Sequence)), \
  819. "fill value must be either float or sequence"
  820. self.im_padding_value = im_padding_value
  821. self.label_padding_value = label_padding_value
  822. def apply(self, sample):
  823. if random.random() < self.prob:
  824. im_h, im_w = sample['image'].shape[:2]
  825. ratio = np.random.uniform(1., self.upper_ratio)
  826. h = int(im_h * ratio)
  827. w = int(im_w * ratio)
  828. if h > im_h and w > im_w:
  829. y = np.random.randint(0, h - im_h)
  830. x = np.random.randint(0, w - im_w)
  831. target_size = (h, w)
  832. offsets = (x, y)
  833. sample = Padding(
  834. target_size=target_size,
  835. pad_mode=-1,
  836. offsets=offsets,
  837. im_padding_value=self.im_padding_value,
  838. label_padding_value=self.label_padding_value)(sample)
  839. return sample
  840. class Padding(Transform):
  841. def __init__(self,
  842. target_size=None,
  843. pad_mode=0,
  844. offsets=None,
  845. im_padding_value=127.5,
  846. label_padding_value=255,
  847. size_divisor=32):
  848. """
  849. Pad image to a specified size or multiple of size_divisor.
  850. Args:
  851. target_size(int, Sequence, optional): Image target size, if None, pad to multiple of size_divisor. Defaults to None.
  852. pad_mode({-1, 0, 1, 2}, optional): Pad mode, currently only supports four modes [-1, 0, 1, 2]. if -1, use specified offsets
  853. if 0, only pad to right and bottom. If 1, pad according to center. If 2, only pad left and top. Defaults to 0.
  854. im_padding_value(Sequence[float]): RGB value of pad area. Defaults to (127.5, 127.5, 127.5).
  855. label_padding_value(int, optional): Filling value for the mask. Defaults to 255.
  856. size_divisor(int): Image width and height after padding is a multiple of coarsest_stride.
  857. """
  858. super(Padding, self).__init__()
  859. if isinstance(target_size, (list, tuple)):
  860. if len(target_size) != 2:
  861. raise ValueError(
  862. '`target_size` should include 2 elements, but it is {}'.
  863. format(target_size))
  864. if isinstance(target_size, int):
  865. target_size = [target_size] * 2
  866. assert pad_mode in [
  867. -1, 0, 1, 2
  868. ], 'currently only supports four modes [-1, 0, 1, 2]'
  869. if pad_mode == -1:
  870. assert offsets, 'if pad_mode is -1, offsets should not be None'
  871. self.target_size = target_size
  872. self.size_divisor = size_divisor
  873. self.pad_mode = pad_mode
  874. self.offsets = offsets
  875. self.im_padding_value = im_padding_value
  876. self.label_padding_value = label_padding_value
  877. def apply_im(self, image, offsets, target_size):
  878. x, y = offsets
  879. h, w = target_size
  880. im_h, im_w, channel = image.shape[:3]
  881. canvas = np.ones((h, w, channel), dtype=np.float32)
  882. canvas *= np.array(self.im_padding_value, dtype=np.float32)
  883. canvas[y:y + im_h, x:x + im_w, :] = image.astype(np.float32)
  884. return canvas
  885. def apply_mask(self, mask, offsets, target_size):
  886. x, y = offsets
  887. im_h, im_w = mask.shape[:2]
  888. h, w = target_size
  889. canvas = np.ones((h, w), dtype=np.float32)
  890. canvas *= np.array(self.label_padding_value, dtype=np.float32)
  891. canvas[y:y + im_h, x:x + im_w] = mask.astype(np.float32)
  892. return canvas
  893. def apply_bbox(self, bbox, offsets):
  894. return bbox + np.array(offsets * 2, dtype=np.float32)
  895. def apply_segm(self, segms, offsets, im_size, size):
  896. x, y = offsets
  897. height, width = im_size
  898. h, w = size
  899. expanded_segms = []
  900. for segm in segms:
  901. if is_poly(segm):
  902. # Polygon format
  903. expanded_segms.append(
  904. [expand_poly(poly, x, y) for poly in segm])
  905. else:
  906. # RLE format
  907. expanded_segms.append(
  908. expand_rle(segm, x, y, height, width, h, w))
  909. return expanded_segms
  910. def apply(self, sample):
  911. im_h, im_w = sample['image'].shape[:2]
  912. if self.target_size:
  913. h, w = self.target_size
  914. assert (
  915. im_h <= h and im_w <= w
  916. ), 'target size ({}, {}) cannot be less than image size ({}, {})'\
  917. .format(h, w, im_h, im_w)
  918. else:
  919. h = (np.ceil(im_h / self.size_divisor) *
  920. self.size_divisor).astype(int)
  921. w = (np.ceil(im_w / self.size_divisor) *
  922. self.size_divisor).astype(int)
  923. if h == im_h and w == im_w:
  924. return sample
  925. if self.pad_mode == -1:
  926. offsets = self.offsets
  927. elif self.pad_mode == 0:
  928. offsets = [0, 0]
  929. elif self.pad_mode == 1:
  930. offsets = [(w - im_w) // 2, (h - im_h) // 2]
  931. else:
  932. offsets = [w - im_w, h - im_h]
  933. sample['image'] = self.apply_im(sample['image'], offsets, (h, w))
  934. if 'image2' in sample:
  935. sample['image2'] = self.apply_im(sample['image2'], offsets, (h, w))
  936. if 'mask' in sample:
  937. sample['mask'] = self.apply_mask(sample['mask'], offsets, (h, w))
  938. if 'aux_masks' in sample:
  939. sample['aux_masks'] = list(map(partial(self.apply_mask, offsets=offsets, target_size=(h,w)), sample['aux_masks']))
  940. if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
  941. sample['gt_bbox'] = self.apply_bbox(sample['gt_bbox'], offsets)
  942. if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
  943. sample['gt_poly'] = self.apply_segm(
  944. sample['gt_poly'], offsets, im_size=[im_h, im_w], size=[h, w])
  945. return sample
  946. class MixupImage(Transform):
  947. def __init__(self, alpha=1.5, beta=1.5, mixup_epoch=-1):
  948. """
  949. Mixup two images and their gt_bbbox/gt_score.
  950. Args:
  951. alpha (float, optional): Alpha parameter of beta distribution. Defaults to 1.5.
  952. beta (float, optional): Beta parameter of beta distribution. Defaults to 1.5.
  953. """
  954. super(MixupImage, self).__init__()
  955. if alpha <= 0.0:
  956. raise ValueError("alpha should be positive in {}".format(self))
  957. if beta <= 0.0:
  958. raise ValueError("beta should be positive in {}".format(self))
  959. self.alpha = alpha
  960. self.beta = beta
  961. self.mixup_epoch = mixup_epoch
  962. def apply_im(self, image1, image2, factor):
  963. h = max(image1.shape[0], image2.shape[0])
  964. w = max(image1.shape[1], image2.shape[1])
  965. img = np.zeros((h, w, image1.shape[2]), 'float32')
  966. img[:image1.shape[0], :image1.shape[1], :] = \
  967. image1.astype('float32') * factor
  968. img[:image2.shape[0], :image2.shape[1], :] += \
  969. image2.astype('float32') * (1.0 - factor)
  970. return img.astype('uint8')
  971. def __call__(self, sample):
  972. if not isinstance(sample, Sequence):
  973. return sample
  974. assert len(sample) == 2, 'mixup need two samples'
  975. factor = np.random.beta(self.alpha, self.beta)
  976. factor = max(0.0, min(1.0, factor))
  977. if factor >= 1.0:
  978. return sample[0]
  979. if factor <= 0.0:
  980. return sample[1]
  981. image = self.apply_im(sample[0]['image'], sample[1]['image'], factor)
  982. result = copy.deepcopy(sample[0])
  983. result['image'] = image
  984. # apply bbox and score
  985. if 'gt_bbox' in sample[0]:
  986. gt_bbox1 = sample[0]['gt_bbox']
  987. gt_bbox2 = sample[1]['gt_bbox']
  988. gt_bbox = np.concatenate((gt_bbox1, gt_bbox2), axis=0)
  989. result['gt_bbox'] = gt_bbox
  990. if 'gt_poly' in sample[0]:
  991. gt_poly1 = sample[0]['gt_poly']
  992. gt_poly2 = sample[1]['gt_poly']
  993. gt_poly = gt_poly1 + gt_poly2
  994. result['gt_poly'] = gt_poly
  995. if 'gt_class' in sample[0]:
  996. gt_class1 = sample[0]['gt_class']
  997. gt_class2 = sample[1]['gt_class']
  998. gt_class = np.concatenate((gt_class1, gt_class2), axis=0)
  999. result['gt_class'] = gt_class
  1000. gt_score1 = np.ones_like(sample[0]['gt_class'])
  1001. gt_score2 = np.ones_like(sample[1]['gt_class'])
  1002. gt_score = np.concatenate(
  1003. (gt_score1 * factor, gt_score2 * (1. - factor)), axis=0)
  1004. result['gt_score'] = gt_score
  1005. if 'is_crowd' in sample[0]:
  1006. is_crowd1 = sample[0]['is_crowd']
  1007. is_crowd2 = sample[1]['is_crowd']
  1008. is_crowd = np.concatenate((is_crowd1, is_crowd2), axis=0)
  1009. result['is_crowd'] = is_crowd
  1010. if 'difficult' in sample[0]:
  1011. is_difficult1 = sample[0]['difficult']
  1012. is_difficult2 = sample[1]['difficult']
  1013. is_difficult = np.concatenate(
  1014. (is_difficult1, is_difficult2), axis=0)
  1015. result['difficult'] = is_difficult
  1016. return result
  1017. class RandomDistort(Transform):
  1018. """
  1019. Random color distortion.
  1020. Args:
  1021. brightness_range(float, optional): Range of brightness distortion. Defaults to .5.
  1022. brightness_prob(float, optional): Probability of brightness distortion. Defaults to .5.
  1023. contrast_range(float, optional): Range of contrast distortion. Defaults to .5.
  1024. contrast_prob(float, optional): Probability of contrast distortion. Defaults to .5.
  1025. saturation_range(float, optional): Range of saturation distortion. Defaults to .5.
  1026. saturation_prob(float, optional): Probability of saturation distortion. Defaults to .5.
  1027. hue_range(float, optional): Range of hue distortion. Defaults to .5.
  1028. hue_prob(float, optional): Probability of hue distortion. Defaults to .5.
  1029. random_apply (bool, optional): whether to apply in random (yolo) or fixed (SSD)
  1030. order. Defaults to True.
  1031. count (int, optional): the number of doing distortion. Defaults to 4.
  1032. shuffle_channel (bool, optional): whether to swap channels randomly. Defaults to False.
  1033. """
  1034. def __init__(self,
  1035. brightness_range=0.5,
  1036. brightness_prob=0.5,
  1037. contrast_range=0.5,
  1038. contrast_prob=0.5,
  1039. saturation_range=0.5,
  1040. saturation_prob=0.5,
  1041. hue_range=18,
  1042. hue_prob=0.5,
  1043. random_apply=True,
  1044. count=4,
  1045. shuffle_channel=False):
  1046. super(RandomDistort, self).__init__()
  1047. self.brightness_range = [1 - brightness_range, 1 + brightness_range]
  1048. self.brightness_prob = brightness_prob
  1049. self.contrast_range = [1 - contrast_range, 1 + contrast_range]
  1050. self.contrast_prob = contrast_prob
  1051. self.saturation_range = [1 - saturation_range, 1 + saturation_range]
  1052. self.saturation_prob = saturation_prob
  1053. self.hue_range = [1 - hue_range, 1 + hue_range]
  1054. self.hue_prob = hue_prob
  1055. self.random_apply = random_apply
  1056. self.count = count
  1057. self.shuffle_channel = shuffle_channel
  1058. def apply_hue(self, image):
  1059. low, high = self.hue_range
  1060. if np.random.uniform(0., 1.) < self.hue_prob:
  1061. return image
  1062. # it works, but result differ from HSV version
  1063. delta = np.random.uniform(low, high)
  1064. u = np.cos(delta * np.pi)
  1065. w = np.sin(delta * np.pi)
  1066. bt = np.array([[1.0, 0.0, 0.0], [0.0, u, -w], [0.0, w, u]])
  1067. tyiq = np.array([[0.299, 0.587, 0.114], [0.596, -0.274, -0.321],
  1068. [0.211, -0.523, 0.311]])
  1069. ityiq = np.array([[1.0, 0.956, 0.621], [1.0, -0.272, -0.647],
  1070. [1.0, -1.107, 1.705]])
  1071. t = np.dot(np.dot(ityiq, bt), tyiq).T
  1072. res_list = []
  1073. channel = image.shape[2]
  1074. for i in range(channel // 3):
  1075. sub_img = image[:, :, 3*i : 3*(i+1)]
  1076. sub_img = sub_img.astype(np.float32)
  1077. sub_img = np.dot(image, t)
  1078. res_list.append(sub_img)
  1079. if channel % 3 != 0:
  1080. i = channel % 3
  1081. res_list.append(image[:, :, -i:])
  1082. return np.concatenate(res_list, axis=2)
  1083. def apply_saturation(self, image):
  1084. low, high = self.saturation_range
  1085. delta = np.random.uniform(low, high)
  1086. if np.random.uniform(0., 1.) < self.saturation_prob:
  1087. return image
  1088. res_list = []
  1089. channel = image.shape[2]
  1090. for i in range(channel // 3):
  1091. sub_img = image[:, :, 3*i : 3*(i+1)]
  1092. sub_img = sub_img.astype(np.float32)
  1093. # it works, but result differ from HSV version
  1094. gray = sub_img * np.array([[[0.299, 0.587, 0.114]]], dtype=np.float32)
  1095. gray = gray.sum(axis=2, keepdims=True)
  1096. gray *= (1.0 - delta)
  1097. sub_img *= delta
  1098. sub_img += gray
  1099. res_list.append(sub_img)
  1100. if channel % 3 != 0:
  1101. i = channel % 3
  1102. res_list.append(image[:, :, -i:])
  1103. return np.concatenate(res_list, axis=2)
  1104. def apply_contrast(self, image):
  1105. low, high = self.contrast_range
  1106. if np.random.uniform(0., 1.) < self.contrast_prob:
  1107. return image
  1108. delta = np.random.uniform(low, high)
  1109. image = image.astype(np.float32)
  1110. image *= delta
  1111. return image
  1112. def apply_brightness(self, image):
  1113. low, high = self.brightness_range
  1114. if np.random.uniform(0., 1.) < self.brightness_prob:
  1115. return image
  1116. delta = np.random.uniform(low, high)
  1117. image = image.astype(np.float32)
  1118. image += delta
  1119. return image
  1120. def apply(self, sample):
  1121. if self.random_apply:
  1122. functions = [
  1123. self.apply_brightness, self.apply_contrast,
  1124. self.apply_saturation, self.apply_hue
  1125. ]
  1126. distortions = np.random.permutation(functions)[:self.count]
  1127. for func in distortions:
  1128. sample['image'] = func(sample['image'])
  1129. if 'image2' in sample:
  1130. sample['image2'] = func(sample['image2'])
  1131. return sample
  1132. sample['image'] = self.apply_brightness(sample['image'])
  1133. if 'image2' in sample:
  1134. sample['image2'] = self.apply_brightness(sample['image2'])
  1135. mode = np.random.randint(0, 2)
  1136. if mode:
  1137. sample['image'] = self.apply_contrast(sample['image'])
  1138. if 'image2' in sample:
  1139. sample['image2'] = self.apply_contrast(sample['image2'])
  1140. sample['image'] = self.apply_saturation(sample['image'])
  1141. sample['image'] = self.apply_hue(sample['image'])
  1142. if 'image2' in sample:
  1143. sample['image2'] = self.apply_saturation(sample['image2'])
  1144. sample['image2'] = self.apply_hue(sample['image2'])
  1145. if not mode:
  1146. sample['image'] = self.apply_contrast(sample['image'])
  1147. if 'image2' in sample:
  1148. sample['image2'] = self.apply_contrast(sample['image2'])
  1149. if self.shuffle_channel:
  1150. if np.random.randint(0, 2):
  1151. sample['image'] = sample['image'][..., np.random.permutation(3)]
  1152. if 'image2' in sample:
  1153. sample['image2'] = sample['image2'][..., np.random.permutation(3)]
  1154. return sample
  1155. class RandomBlur(Transform):
  1156. """
  1157. Randomly blur input image(s).
  1158. Args:
  1159. prob (float): Probability of blurring.
  1160. """
  1161. def __init__(self, prob=0.1):
  1162. super(RandomBlur, self).__init__()
  1163. self.prob = prob
  1164. def apply_im(self, image, radius):
  1165. image = cv2.GaussianBlur(image, (radius, radius), 0, 0)
  1166. return image
  1167. def apply(self, sample):
  1168. if self.prob <= 0:
  1169. n = 0
  1170. elif self.prob >= 1:
  1171. n = 1
  1172. else:
  1173. n = int(1.0 / self.prob)
  1174. if n > 0:
  1175. if np.random.randint(0, n) == 0:
  1176. radius = np.random.randint(3, 10)
  1177. if radius % 2 != 1:
  1178. radius = radius + 1
  1179. if radius > 9:
  1180. radius = 9
  1181. sample['image'] = self.apply_im(sample['image'], radius)
  1182. if 'image2' in sample:
  1183. sample['image2'] = self.apply_im(sample['image2'], radius)
  1184. return sample
  1185. class _PadBox(Transform):
  1186. def __init__(self, num_max_boxes=50):
  1187. """
  1188. Pad zeros to bboxes if number of bboxes is less than num_max_boxes.
  1189. Args:
  1190. num_max_boxes (int, optional): the max number of bboxes. Defaults to 50.
  1191. """
  1192. self.num_max_boxes = num_max_boxes
  1193. super(_PadBox, self).__init__()
  1194. def apply(self, sample):
  1195. gt_num = min(self.num_max_boxes, len(sample['gt_bbox']))
  1196. num_max = self.num_max_boxes
  1197. pad_bbox = np.zeros((num_max, 4), dtype=np.float32)
  1198. if gt_num > 0:
  1199. pad_bbox[:gt_num, :] = sample['gt_bbox'][:gt_num, :]
  1200. sample['gt_bbox'] = pad_bbox
  1201. if 'gt_class' in sample:
  1202. pad_class = np.zeros((num_max, ), dtype=np.int32)
  1203. if gt_num > 0:
  1204. pad_class[:gt_num] = sample['gt_class'][:gt_num, 0]
  1205. sample['gt_class'] = pad_class
  1206. if 'gt_score' in sample:
  1207. pad_score = np.zeros((num_max, ), dtype=np.float32)
  1208. if gt_num > 0:
  1209. pad_score[:gt_num] = sample['gt_score'][:gt_num, 0]
  1210. sample['gt_score'] = pad_score
  1211. # in training, for example in op ExpandImage,
  1212. # the bbox and gt_class is expanded, but the difficult is not,
  1213. # so, judging by it's length
  1214. if 'difficult' in sample:
  1215. pad_diff = np.zeros((num_max, ), dtype=np.int32)
  1216. if gt_num > 0:
  1217. pad_diff[:gt_num] = sample['difficult'][:gt_num, 0]
  1218. sample['difficult'] = pad_diff
  1219. if 'is_crowd' in sample:
  1220. pad_crowd = np.zeros((num_max, ), dtype=np.int32)
  1221. if gt_num > 0:
  1222. pad_crowd[:gt_num] = sample['is_crowd'][:gt_num, 0]
  1223. sample['is_crowd'] = pad_crowd
  1224. return sample
  1225. class _NormalizeBox(Transform):
  1226. def __init__(self):
  1227. super(_NormalizeBox, self).__init__()
  1228. def apply(self, sample):
  1229. height, width = sample['image'].shape[:2]
  1230. for i in range(sample['gt_bbox'].shape[0]):
  1231. sample['gt_bbox'][i][0] = sample['gt_bbox'][i][0] / width
  1232. sample['gt_bbox'][i][1] = sample['gt_bbox'][i][1] / height
  1233. sample['gt_bbox'][i][2] = sample['gt_bbox'][i][2] / width
  1234. sample['gt_bbox'][i][3] = sample['gt_bbox'][i][3] / height
  1235. return sample
  1236. class _BboxXYXY2XYWH(Transform):
  1237. """
  1238. Convert bbox XYXY format to XYWH format.
  1239. """
  1240. def __init__(self):
  1241. super(_BboxXYXY2XYWH, self).__init__()
  1242. def apply(self, sample):
  1243. bbox = sample['gt_bbox']
  1244. bbox[:, 2:4] = bbox[:, 2:4] - bbox[:, :2]
  1245. bbox[:, :2] = bbox[:, :2] + bbox[:, 2:4] / 2.
  1246. sample['gt_bbox'] = bbox
  1247. return sample
  1248. class _Permute(Transform):
  1249. def __init__(self):
  1250. super(_Permute, self).__init__()
  1251. def apply(self, sample):
  1252. sample['image'] = permute(sample['image'], False)
  1253. if 'image2' in sample:
  1254. sample['image2'] = permute(sample['image2'], False)
  1255. return sample
  1256. class ArrangeSegmenter(Transform):
  1257. def __init__(self, mode):
  1258. super(ArrangeSegmenter, self).__init__()
  1259. if mode not in ['train', 'eval', 'test', 'quant']:
  1260. raise ValueError(
  1261. "mode should be defined as one of ['train', 'eval', 'test', 'quant']!"
  1262. )
  1263. self.mode = mode
  1264. def apply(self, sample):
  1265. if 'mask' in sample:
  1266. mask = sample['mask']
  1267. image = permute(sample['image'], False)
  1268. if self.mode == 'train':
  1269. mask = mask.astype('int64')
  1270. return image, mask
  1271. if self.mode == 'eval':
  1272. mask = np.asarray(Image.open(mask))
  1273. mask = mask[np.newaxis, :, :].astype('int64')
  1274. return image, mask
  1275. if self.mode == 'test':
  1276. return image,
  1277. class ArrangeChangeDetector(Transform):
  1278. def __init__(self, mode):
  1279. super(ArrangeChangeDetector, self).__init__()
  1280. if mode not in ['train', 'eval', 'test', 'quant']:
  1281. raise ValueError(
  1282. "mode should be defined as one of ['train', 'eval', 'test', 'quant']!"
  1283. )
  1284. self.mode = mode
  1285. def apply(self, sample):
  1286. if 'mask' in sample:
  1287. mask = sample['mask']
  1288. image_t1 = permute(sample['image'], False)
  1289. image_t2 = permute(sample['image2'], False)
  1290. if self.mode == 'train':
  1291. mask = mask.astype('int64')
  1292. masks = [mask]
  1293. if 'aux_masks' in sample:
  1294. masks.extend(map(methodcaller('astype', 'int64'), sample['aux_masks']))
  1295. return (image_t1, image_t2,) + tuple(masks)
  1296. if self.mode == 'eval':
  1297. mask = np.asarray(Image.open(mask))
  1298. mask = mask[np.newaxis, :, :].astype('int64')
  1299. return image_t1, image_t2, mask
  1300. if self.mode == 'test':
  1301. return image_t1, image_t2,
  1302. class ArrangeClassifier(Transform):
  1303. def __init__(self, mode):
  1304. super(ArrangeClassifier, self).__init__()
  1305. if mode not in ['train', 'eval', 'test', 'quant']:
  1306. raise ValueError(
  1307. "mode should be defined as one of ['train', 'eval', 'test', 'quant']!"
  1308. )
  1309. self.mode = mode
  1310. def apply(self, sample):
  1311. image = permute(sample['image'], False)
  1312. if self.mode in ['train', 'eval']:
  1313. return image, sample['label']
  1314. else:
  1315. return image
  1316. class ArrangeDetector(Transform):
  1317. def __init__(self, mode):
  1318. super(ArrangeDetector, self).__init__()
  1319. if mode not in ['train', 'eval', 'test', 'quant']:
  1320. raise ValueError(
  1321. "mode should be defined as one of ['train', 'eval', 'test', 'quant']!"
  1322. )
  1323. self.mode = mode
  1324. def apply(self, sample):
  1325. if self.mode == 'eval' and 'gt_poly' in sample:
  1326. del sample['gt_poly']
  1327. return sample