operators.py 50 KB

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