operators.py 65 KB

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