det_transforms.py 55 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186
  1. # copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from .ops import *
  15. from .box_utils import *
  16. import random
  17. import os.path as osp
  18. import numpy as np
  19. from PIL import Image, ImageEnhance
  20. import cv2
  21. class Compose:
  22. """根据数据预处理/增强列表对输入数据进行操作。
  23. 所有操作的输入图像流形状均是[H, W, C],其中H为图像高,W为图像宽,C为图像通道数。
  24. Args:
  25. transforms (list): 数据预处理/增强列表。
  26. Raises:
  27. TypeError: 形参数据类型不满足需求。
  28. ValueError: 数据长度不匹配。
  29. """
  30. def __init__(self, transforms):
  31. if not isinstance(transforms, list):
  32. raise TypeError('The transforms must be a list!')
  33. if len(transforms) < 1:
  34. raise ValueError('The length of transforms ' + \
  35. 'must be equal or larger than 1!')
  36. self.transforms = transforms
  37. self.use_mixup = False
  38. for t in self.transforms:
  39. if t.__class__.__name__ == 'MixupImage':
  40. self.use_mixup = True
  41. def __call__(self, im, im_info=None, label_info=None):
  42. """
  43. Args:
  44. im (str/np.ndarray): 图像路径/图像np.ndarray数据。
  45. im_info (dict): 存储与图像相关的信息,dict中的字段如下:
  46. - im_id (np.ndarray): 图像序列号,形状为(1,)。
  47. - origin_shape (np.ndarray): 图像原始大小,形状为(2,),
  48. origin_shape[0]为高,origin_shape[1]为宽。
  49. - mixup (list): list为[im, im_info, label_info],分别对应
  50. 与当前图像进行mixup的图像np.ndarray数据、图像相关信息、标注框相关信息;
  51. 注意,当前epoch若无需进行mixup,则无该字段。
  52. label_info (dict): 存储与标注框相关的信息,dict中的字段如下:
  53. - gt_bbox (np.ndarray): 真实标注框坐标[x1, y1, x2, y2],形状为(n, 4),
  54. 其中n代表真实标注框的个数。
  55. - gt_class (np.ndarray): 每个真实标注框对应的类别序号,形状为(n, 1),
  56. 其中n代表真实标注框的个数。
  57. - gt_score (np.ndarray): 每个真实标注框对应的混合得分,形状为(n, 1),
  58. 其中n代表真实标注框的个数。
  59. - gt_poly (list): 每个真实标注框内的多边形分割区域,每个分割区域由点的x、y坐标组成,
  60. 长度为n,其中n代表真实标注框的个数。
  61. - is_crowd (np.ndarray): 每个真实标注框中是否是一组对象,形状为(n, 1),
  62. 其中n代表真实标注框的个数。
  63. - difficult (np.ndarray): 每个真实标注框中的对象是否为难识别对象,形状为(n, 1),
  64. 其中n代表真实标注框的个数。
  65. Returns:
  66. tuple: 根据网络所需字段所组成的tuple;
  67. 字段由transforms中的最后一个数据预处理操作决定。
  68. """
  69. def decode_image(im_file, im_info, label_info):
  70. if im_info is None:
  71. im_info = dict()
  72. try:
  73. im = cv2.imread(im_file).astype('float32')
  74. except:
  75. raise TypeError(
  76. 'Can\'t read The image file {}!'.format(im_file))
  77. im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
  78. # make default im_info with [h, w, 1]
  79. im_info['im_resize_info'] = np.array(
  80. [im.shape[0], im.shape[1], 1.], dtype=np.float32)
  81. # copy augment_shape from origin_shape
  82. im_info['augment_shape'] = np.array([im.shape[0],
  83. im.shape[1]]).astype('int32')
  84. if not self.use_mixup:
  85. if 'mixup' in im_info:
  86. del im_info['mixup']
  87. # decode mixup image
  88. if 'mixup' in im_info:
  89. im_info['mixup'] = \
  90. decode_image(im_info['mixup'][0],
  91. im_info['mixup'][1],
  92. im_info['mixup'][2])
  93. if label_info is None:
  94. return (im, im_info)
  95. else:
  96. return (im, im_info, label_info)
  97. outputs = decode_image(im, im_info, label_info)
  98. im = outputs[0]
  99. im_info = outputs[1]
  100. if len(outputs) == 3:
  101. label_info = outputs[2]
  102. for op in self.transforms:
  103. if im is None:
  104. return None
  105. outputs = op(im, im_info, label_info)
  106. im = outputs[0]
  107. return outputs
  108. class ResizeByShort:
  109. """根据图像的短边调整图像大小(resize)。
  110. 1. 获取图像的长边和短边长度。
  111. 2. 根据短边与short_size的比例,计算长边的目标长度,
  112. 此时高、宽的resize比例为short_size/原图短边长度。
  113. 3. 如果max_size>0,调整resize比例:
  114. 如果长边的目标长度>max_size,则高、宽的resize比例为max_size/原图长边长度。
  115. 4. 根据调整大小的比例对图像进行resize。
  116. Args:
  117. target_size (int): 短边目标长度。默认为800。
  118. max_size (int): 长边目标长度的最大限制。默认为1333。
  119. Raises:
  120. TypeError: 形参数据类型不满足需求。
  121. """
  122. def __init__(self, short_size=800, max_size=1333):
  123. self.max_size = int(max_size)
  124. if not isinstance(short_size, int):
  125. raise TypeError(
  126. "Type of short_size is invalid. Must be Integer, now is {}".
  127. format(type(short_size)))
  128. self.short_size = short_size
  129. if not (isinstance(self.max_size, int)):
  130. raise TypeError("max_size: input type is invalid.")
  131. def __call__(self, im, im_info=None, label_info=None):
  132. """
  133. Args:
  134. im (numnp.ndarraypy): 图像np.ndarray数据。
  135. im_info (dict, 可选): 存储与图像相关的信息。
  136. label_info (dict, 可选): 存储与标注框相关的信息。
  137. Returns:
  138. tuple: 当label_info为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  139. 当label_info不为空时,返回的tuple为(im, im_info, label_info),分别对应图像np.ndarray数据、
  140. 存储与标注框相关信息的字典。
  141. 其中,im_info更新字段为:
  142. - im_resize_info (np.ndarray): resize后的图像高、resize后的图像宽、resize后的图像相对原始图的缩放比例
  143. 三者组成的np.ndarray,形状为(3,)。
  144. Raises:
  145. TypeError: 形参数据类型不满足需求。
  146. ValueError: 数据长度不匹配。
  147. """
  148. if im_info is None:
  149. im_info = dict()
  150. if not isinstance(im, np.ndarray):
  151. raise TypeError("ResizeByShort: image type is not numpy.")
  152. if len(im.shape) != 3:
  153. raise ValueError('ResizeByShort: image is not 3-dimensional.')
  154. im_short_size = min(im.shape[0], im.shape[1])
  155. im_long_size = max(im.shape[0], im.shape[1])
  156. scale = float(self.short_size) / im_short_size
  157. if self.max_size > 0 and np.round(
  158. scale * im_long_size) > self.max_size:
  159. scale = float(self.max_size) / float(im_long_size)
  160. resized_width = int(round(im.shape[1] * scale))
  161. resized_height = int(round(im.shape[0] * scale))
  162. im_resize_info = [resized_height, resized_width, scale]
  163. im = cv2.resize(
  164. im, (resized_width, resized_height),
  165. interpolation=cv2.INTER_LINEAR)
  166. im_info['im_resize_info'] = np.array(im_resize_info).astype(np.float32)
  167. if label_info is None:
  168. return (im, im_info)
  169. else:
  170. return (im, im_info, label_info)
  171. class Padding:
  172. """将图像的长和宽padding至coarsest_stride的倍数。如输入图像为[300, 640],
  173. `coarest_stride`为32,则由于300不为32的倍数,因此在图像最右和最下使用0值
  174. 进行padding,最终输出图像为[320, 640]。
  175. 1. 如果coarsest_stride为1则直接返回。
  176. 2. 获取图像的高H、宽W。
  177. 3. 计算填充后图像的高H_new、宽W_new。
  178. 4. 构建大小为(H_new, W_new, 3)像素值为0的np.ndarray,
  179. 并将原图的np.ndarray粘贴于左上角。
  180. Args:
  181. coarsest_stride (int): 填充后的图像长、宽为该参数的倍数,默认为1。
  182. target_size (int|list): 填充后的图像长、宽,默认为1。
  183. """
  184. def __init__(self, coarsest_stride=1, target_size=None):
  185. self.coarsest_stride = coarsest_stride
  186. self.target_size = target_size
  187. def __call__(self, im, im_info=None, label_info=None):
  188. """
  189. Args:
  190. im (numnp.ndarraypy): 图像np.ndarray数据。
  191. im_info (dict, 可选): 存储与图像相关的信息。
  192. label_info (dict, 可选): 存储与标注框相关的信息。
  193. Returns:
  194. tuple: 当label_info为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  195. 当label_info不为空时,返回的tuple为(im, im_info, label_info),分别对应图像np.ndarray数据、
  196. 存储与标注框相关信息的字典。
  197. Raises:
  198. TypeError: 形参数据类型不满足需求。
  199. ValueError: 数据长度不匹配。
  200. ValueError: target_size小于原图的大小。
  201. """
  202. if self.coarsest_stride == 1 and self.target_size is None:
  203. if label_info is None:
  204. return (im, im_info)
  205. else:
  206. return (im, im_info, label_info)
  207. if im_info is None:
  208. im_info = dict()
  209. if not isinstance(im, np.ndarray):
  210. raise TypeError("Padding: image type is not numpy.")
  211. if len(im.shape) != 3:
  212. raise ValueError('Padding: image is not 3-dimensional.')
  213. im_h, im_w, im_c = im.shape[:]
  214. if self.coarsest_stride > 1:
  215. padding_im_h = int(
  216. np.ceil(im_h / self.coarsest_stride) * self.coarsest_stride)
  217. padding_im_w = int(
  218. np.ceil(im_w / self.coarsest_stride) * self.coarsest_stride)
  219. if self.target_size is not None:
  220. if isinstance(self.target_size, int):
  221. padding_im_h = self.target_size
  222. padding_im_w = self.target_size
  223. else:
  224. padding_im_h = self.target_size[0]
  225. padding_im_w = self.target_size[1]
  226. pad_height = padding_im_h - im_h
  227. pad_width = padding_im_w - im_w
  228. if pad_height < 0 or pad_width < 0:
  229. raise ValueError(
  230. 'the size of image should be less than target_size, but the size of image ({}, {}), is larger than target_size ({}, {})'
  231. .format(im_w, im_h, padding_im_w, padding_im_h))
  232. padding_im = np.zeros((padding_im_h, padding_im_w, im_c),
  233. dtype=np.float32)
  234. padding_im[:im_h, :im_w, :] = im
  235. if label_info is None:
  236. return (padding_im, im_info)
  237. else:
  238. return (padding_im, im_info, label_info)
  239. class Resize:
  240. """调整图像大小(resize)。
  241. - 当目标大小(target_size)类型为int时,根据插值方式,
  242. 将图像resize为[target_size, target_size]。
  243. - 当目标大小(target_size)类型为list或tuple时,根据插值方式,
  244. 将图像resize为target_size。
  245. 注意:当插值方式为“RANDOM”时,则随机选取一种插值方式进行resize。
  246. Args:
  247. target_size (int/list/tuple): 短边目标长度。默认为608。
  248. interp (str): resize的插值方式,与opencv的插值方式对应,取值范围为
  249. ['NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4', 'RANDOM']。默认为"LINEAR"。
  250. Raises:
  251. TypeError: 形参数据类型不满足需求。
  252. ValueError: 插值方式不在['NEAREST', 'LINEAR', 'CUBIC',
  253. 'AREA', 'LANCZOS4', 'RANDOM']中。
  254. """
  255. # The interpolation mode
  256. interp_dict = {
  257. 'NEAREST': cv2.INTER_NEAREST,
  258. 'LINEAR': cv2.INTER_LINEAR,
  259. 'CUBIC': cv2.INTER_CUBIC,
  260. 'AREA': cv2.INTER_AREA,
  261. 'LANCZOS4': cv2.INTER_LANCZOS4
  262. }
  263. def __init__(self, target_size=608, interp='LINEAR'):
  264. self.interp = interp
  265. if not (interp == "RANDOM" or interp in self.interp_dict):
  266. raise ValueError("interp should be one of {}".format(
  267. self.interp_dict.keys()))
  268. if isinstance(target_size, list) or isinstance(target_size, tuple):
  269. if len(target_size) != 2:
  270. raise TypeError(
  271. 'when target is list or tuple, it should include 2 elements, but it is {}'
  272. .format(target_size))
  273. elif not isinstance(target_size, int):
  274. raise TypeError(
  275. "Type of target_size is invalid. Must be Integer or List or tuple, now is {}"
  276. .format(type(target_size)))
  277. self.target_size = target_size
  278. def __call__(self, im, im_info=None, label_info=None):
  279. """
  280. Args:
  281. im (np.ndarray): 图像np.ndarray数据。
  282. im_info (dict, 可选): 存储与图像相关的信息。
  283. label_info (dict, 可选): 存储与标注框相关的信息。
  284. Returns:
  285. tuple: 当label_info为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  286. 当label_info不为空时,返回的tuple为(im, im_info, label_info),分别对应图像np.ndarray数据、
  287. 存储与标注框相关信息的字典。
  288. Raises:
  289. TypeError: 形参数据类型不满足需求。
  290. ValueError: 数据长度不匹配。
  291. """
  292. if im_info is None:
  293. im_info = dict()
  294. if not isinstance(im, np.ndarray):
  295. raise TypeError("Resize: image type is not numpy.")
  296. if len(im.shape) != 3:
  297. raise ValueError('Resize: image is not 3-dimensional.')
  298. if self.interp == "RANDOM":
  299. interp = random.choice(list(self.interp_dict.keys()))
  300. else:
  301. interp = self.interp
  302. im = resize(im, self.target_size, self.interp_dict[interp])
  303. if label_info is None:
  304. return (im, im_info)
  305. else:
  306. return (im, im_info, label_info)
  307. class RandomHorizontalFlip:
  308. """随机翻转图像、标注框、分割信息,模型训练时的数据增强操作。
  309. 1. 随机采样一个0-1之间的小数,当小数小于水平翻转概率时,
  310. 执行2-4步操作,否则直接返回。
  311. 2. 水平翻转图像。
  312. 3. 计算翻转后的真实标注框的坐标,更新label_info中的gt_bbox信息。
  313. 4. 计算翻转后的真实分割区域的坐标,更新label_info中的gt_poly信息。
  314. Args:
  315. prob (float): 随机水平翻转的概率。默认为0.5。
  316. Raises:
  317. TypeError: 形参数据类型不满足需求。
  318. """
  319. def __init__(self, prob=0.5):
  320. self.prob = prob
  321. if not isinstance(self.prob, float):
  322. raise TypeError("RandomHorizontalFlip: input type is invalid.")
  323. def __call__(self, im, im_info=None, label_info=None):
  324. """
  325. Args:
  326. im (np.ndarray): 图像np.ndarray数据。
  327. im_info (dict, 可选): 存储与图像相关的信息。
  328. label_info (dict, 可选): 存储与标注框相关的信息。
  329. Returns:
  330. tuple: 当label_info为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  331. 当label_info不为空时,返回的tuple为(im, im_info, label_info),分别对应图像np.ndarray数据、
  332. 存储与标注框相关信息的字典。
  333. 其中,im_info更新字段为:
  334. - gt_bbox (np.ndarray): 水平翻转后的标注框坐标[x1, y1, x2, y2],形状为(n, 4),
  335. 其中n代表真实标注框的个数。
  336. - gt_poly (list): 水平翻转后的多边形分割区域的x、y坐标,长度为n,
  337. 其中n代表真实标注框的个数。
  338. Raises:
  339. TypeError: 形参数据类型不满足需求。
  340. ValueError: 数据长度不匹配。
  341. """
  342. if not isinstance(im, np.ndarray):
  343. raise TypeError(
  344. "RandomHorizontalFlip: image is not a numpy array.")
  345. if len(im.shape) != 3:
  346. raise ValueError(
  347. "RandomHorizontalFlip: image is not 3-dimensional.")
  348. if im_info is None or label_info is None:
  349. raise TypeError(
  350. 'Cannot do RandomHorizontalFlip! ' +
  351. 'Becasuse the im_info and label_info can not be None!')
  352. if 'augment_shape' not in im_info:
  353. raise TypeError('Cannot do RandomHorizontalFlip! ' + \
  354. 'Becasuse augment_shape is not in im_info!')
  355. if 'gt_bbox' not in label_info:
  356. raise TypeError('Cannot do RandomHorizontalFlip! ' + \
  357. 'Becasuse gt_bbox is not in label_info!')
  358. augment_shape = im_info['augment_shape']
  359. gt_bbox = label_info['gt_bbox']
  360. height = augment_shape[0]
  361. width = augment_shape[1]
  362. if np.random.uniform(0, 1) < self.prob:
  363. im = horizontal_flip(im)
  364. if gt_bbox.shape[0] == 0:
  365. if label_info is None:
  366. return (im, im_info)
  367. else:
  368. return (im, im_info, label_info)
  369. label_info['gt_bbox'] = box_horizontal_flip(gt_bbox, width)
  370. if 'gt_poly' in label_info and \
  371. len(label_info['gt_poly']) != 0:
  372. label_info['gt_poly'] = segms_horizontal_flip(
  373. label_info['gt_poly'], height, width)
  374. if label_info is None:
  375. return (im, im_info)
  376. else:
  377. return (im, im_info, label_info)
  378. class Normalize:
  379. """对图像进行标准化。
  380. 1. 归一化图像到到区间[0.0, 1.0]。
  381. 2. 对图像进行减均值除以标准差操作。
  382. Args:
  383. mean (list): 图像数据集的均值。默认为[0.485, 0.456, 0.406]。
  384. std (list): 图像数据集的标准差。默认为[0.229, 0.224, 0.225]。
  385. Raises:
  386. TypeError: 形参数据类型不满足需求。
  387. """
  388. def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
  389. self.mean = mean
  390. self.std = std
  391. if not (isinstance(self.mean, list) and isinstance(self.std, list)):
  392. raise TypeError("NormalizeImage: input type is invalid.")
  393. from functools import reduce
  394. if reduce(lambda x, y: x * y, self.std) == 0:
  395. raise TypeError('NormalizeImage: std is invalid!')
  396. def __call__(self, im, im_info=None, label_info=None):
  397. """
  398. Args:
  399. im (numnp.ndarraypy): 图像np.ndarray数据。
  400. im_info (dict, 可选): 存储与图像相关的信息。
  401. label_info (dict, 可选): 存储与标注框相关的信息。
  402. Returns:
  403. tuple: 当label_info为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  404. 当label_info不为空时,返回的tuple为(im, im_info, label_info),分别对应图像np.ndarray数据、
  405. 存储与标注框相关信息的字典。
  406. """
  407. mean = np.array(self.mean)[np.newaxis, np.newaxis, :]
  408. std = np.array(self.std)[np.newaxis, np.newaxis, :]
  409. im = normalize(im, mean, std)
  410. if label_info is None:
  411. return (im, im_info)
  412. else:
  413. return (im, im_info, label_info)
  414. class RandomDistort:
  415. """以一定的概率对图像进行随机像素内容变换,模型训练时的数据增强操作
  416. 1. 对变换的操作顺序进行随机化操作。
  417. 2. 按照1中的顺序以一定的概率在范围[-range, range]对图像进行随机像素内容变换。
  418. Args:
  419. brightness_range (float): 明亮度因子的范围。默认为0.5。
  420. brightness_prob (float): 随机调整明亮度的概率。默认为0.5。
  421. contrast_range (float): 对比度因子的范围。默认为0.5。
  422. contrast_prob (float): 随机调整对比度的概率。默认为0.5。
  423. saturation_range (float): 饱和度因子的范围。默认为0.5。
  424. saturation_prob (float): 随机调整饱和度的概率。默认为0.5。
  425. hue_range (int): 色调因子的范围。默认为18。
  426. hue_prob (float): 随机调整色调的概率。默认为0.5。
  427. """
  428. def __init__(self,
  429. brightness_range=0.5,
  430. brightness_prob=0.5,
  431. contrast_range=0.5,
  432. contrast_prob=0.5,
  433. saturation_range=0.5,
  434. saturation_prob=0.5,
  435. hue_range=18,
  436. hue_prob=0.5):
  437. self.brightness_range = brightness_range
  438. self.brightness_prob = brightness_prob
  439. self.contrast_range = contrast_range
  440. self.contrast_prob = contrast_prob
  441. self.saturation_range = saturation_range
  442. self.saturation_prob = saturation_prob
  443. self.hue_range = hue_range
  444. self.hue_prob = hue_prob
  445. def __call__(self, im, im_info=None, label_info=None):
  446. """
  447. Args:
  448. im (np.ndarray): 图像np.ndarray数据。
  449. im_info (dict, 可选): 存储与图像相关的信息。
  450. label_info (dict, 可选): 存储与标注框相关的信息。
  451. Returns:
  452. tuple: 当label_info为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  453. 当label_info不为空时,返回的tuple为(im, im_info, label_info),分别对应图像np.ndarray数据、
  454. 存储与标注框相关信息的字典。
  455. """
  456. brightness_lower = 1 - self.brightness_range
  457. brightness_upper = 1 + self.brightness_range
  458. contrast_lower = 1 - self.contrast_range
  459. contrast_upper = 1 + self.contrast_range
  460. saturation_lower = 1 - self.saturation_range
  461. saturation_upper = 1 + self.saturation_range
  462. hue_lower = -self.hue_range
  463. hue_upper = self.hue_range
  464. ops = [brightness, contrast, saturation, hue]
  465. random.shuffle(ops)
  466. params_dict = {
  467. 'brightness': {
  468. 'brightness_lower': brightness_lower,
  469. 'brightness_upper': brightness_upper
  470. },
  471. 'contrast': {
  472. 'contrast_lower': contrast_lower,
  473. 'contrast_upper': contrast_upper
  474. },
  475. 'saturation': {
  476. 'saturation_lower': saturation_lower,
  477. 'saturation_upper': saturation_upper
  478. },
  479. 'hue': {
  480. 'hue_lower': hue_lower,
  481. 'hue_upper': hue_upper
  482. }
  483. }
  484. prob_dict = {
  485. 'brightness': self.brightness_prob,
  486. 'contrast': self.contrast_prob,
  487. 'saturation': self.saturation_prob,
  488. 'hue': self.hue_prob
  489. }
  490. im = im.astype('uint8')
  491. im = Image.fromarray(im)
  492. for id in range(4):
  493. params = params_dict[ops[id].__name__]
  494. prob = prob_dict[ops[id].__name__]
  495. params['im'] = im
  496. if np.random.uniform(0, 1) < prob:
  497. im = ops[id](**params)
  498. im = np.asarray(im).astype('float32')
  499. if label_info is None:
  500. return (im, im_info)
  501. else:
  502. return (im, im_info, label_info)
  503. class MixupImage:
  504. """对图像进行mixup操作,模型训练时的数据增强操作,目前仅YOLOv3模型支持该transform。
  505. 当label_info中不存在mixup字段时,直接返回,否则进行下述操作:
  506. 1. 从随机beta分布中抽取出随机因子factor。
  507. 2.
  508. - 当factor>=1.0时,去除label_info中的mixup字段,直接返回。
  509. - 当factor<=0.0时,直接返回label_info中的mixup字段,并在label_info中去除该字段。
  510. - 其余情况,执行下述操作:
  511. (1)原图像乘以factor,mixup图像乘以(1-factor),叠加2个结果。
  512. (2)拼接原图像标注框和mixup图像标注框。
  513. (3)拼接原图像标注框类别和mixup图像标注框类别。
  514. (4)原图像标注框混合得分乘以factor,mixup图像标注框混合得分乘以(1-factor),叠加2个结果。
  515. 3. 更新im_info中的augment_shape信息。
  516. Args:
  517. alpha (float): 随机beta分布的下限。默认为1.5。
  518. beta (float): 随机beta分布的上限。默认为1.5。
  519. mixup_epoch (int): 在前mixup_epoch轮使用mixup增强操作;当该参数为-1时,该策略不会生效。
  520. 默认为-1。
  521. Raises:
  522. ValueError: 数据长度不匹配。
  523. """
  524. def __init__(self, alpha=1.5, beta=1.5, mixup_epoch=-1):
  525. self.alpha = alpha
  526. self.beta = beta
  527. if self.alpha <= 0.0:
  528. raise ValueError("alpha shold be positive in MixupImage")
  529. if self.beta <= 0.0:
  530. raise ValueError("beta shold be positive in MixupImage")
  531. self.mixup_epoch = mixup_epoch
  532. def _mixup_img(self, img1, img2, factor):
  533. h = max(img1.shape[0], img2.shape[0])
  534. w = max(img1.shape[1], img2.shape[1])
  535. img = np.zeros((h, w, img1.shape[2]), 'float32')
  536. img[:img1.shape[0], :img1.shape[1], :] = \
  537. img1.astype('float32') * factor
  538. img[:img2.shape[0], :img2.shape[1], :] += \
  539. img2.astype('float32') * (1.0 - factor)
  540. return img.astype('uint8')
  541. def __call__(self, im, im_info=None, label_info=None):
  542. """
  543. Args:
  544. im (np.ndarray): 图像np.ndarray数据。
  545. im_info (dict, 可选): 存储与图像相关的信息。
  546. label_info (dict, 可选): 存储与标注框相关的信息。
  547. Returns:
  548. tuple: 当label_info为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  549. 当label_info不为空时,返回的tuple为(im, im_info, label_info),分别对应图像np.ndarray数据、
  550. 存储与标注框相关信息的字典。
  551. 其中,im_info更新字段为:
  552. - augment_shape (np.ndarray): mixup后的图像高、宽二者组成的np.ndarray,形状为(2,)。
  553. im_info删除的字段:
  554. - mixup (list): 与当前字段进行mixup的图像相关信息。
  555. label_info更新字段为:
  556. - gt_bbox (np.ndarray): mixup后真实标注框坐标,形状为(n, 4),
  557. 其中n代表真实标注框的个数。
  558. - gt_class (np.ndarray): mixup后每个真实标注框对应的类别序号,形状为(n, 1),
  559. 其中n代表真实标注框的个数。
  560. - gt_score (np.ndarray): mixup后每个真实标注框对应的混合得分,形状为(n, 1),
  561. 其中n代表真实标注框的个数。
  562. Raises:
  563. TypeError: 形参数据类型不满足需求。
  564. """
  565. if im_info is None:
  566. raise TypeError('Cannot do MixupImage! ' +
  567. 'Becasuse the im_info can not be None!')
  568. if 'mixup' not in im_info:
  569. if label_info is None:
  570. return (im, im_info)
  571. else:
  572. return (im, im_info, label_info)
  573. factor = np.random.beta(self.alpha, self.beta)
  574. factor = max(0.0, min(1.0, factor))
  575. if im_info['epoch'] > self.mixup_epoch \
  576. or factor >= 1.0:
  577. im_info.pop('mixup')
  578. if label_info is None:
  579. return (im, im_info)
  580. else:
  581. return (im, im_info, label_info)
  582. if factor <= 0.0:
  583. return im_info.pop('mixup')
  584. im = self._mixup_img(im, im_info['mixup'][0], factor)
  585. if label_info is None:
  586. raise TypeError('Cannot do MixupImage! ' +
  587. 'Becasuse the label_info can not be None!')
  588. if 'gt_bbox' not in label_info or \
  589. 'gt_class' not in label_info or \
  590. 'gt_score' not in label_info:
  591. raise TypeError('Cannot do MixupImage! ' + \
  592. 'Becasuse gt_bbox/gt_class/gt_score is not in label_info!')
  593. gt_bbox1 = label_info['gt_bbox']
  594. gt_bbox2 = im_info['mixup'][2]['gt_bbox']
  595. gt_bbox = np.concatenate((gt_bbox1, gt_bbox2), axis=0)
  596. gt_class1 = label_info['gt_class']
  597. gt_class2 = im_info['mixup'][2]['gt_class']
  598. gt_class = np.concatenate((gt_class1, gt_class2), axis=0)
  599. gt_score1 = label_info['gt_score']
  600. gt_score2 = im_info['mixup'][2]['gt_score']
  601. gt_score = np.concatenate(
  602. (gt_score1 * factor, gt_score2 * (1. - factor)), axis=0)
  603. label_info['gt_bbox'] = gt_bbox
  604. label_info['gt_score'] = gt_score
  605. label_info['gt_class'] = gt_class
  606. im_info['augment_shape'] = np.array([im.shape[0],
  607. im.shape[1]]).astype('int32')
  608. im_info.pop('mixup')
  609. if label_info is None:
  610. return (im, im_info)
  611. else:
  612. return (im, im_info, label_info)
  613. class RandomExpand:
  614. """随机扩张图像,模型训练时的数据增强操作。
  615. 1. 随机选取扩张比例(扩张比例大于1时才进行扩张)。
  616. 2. 计算扩张后图像大小。
  617. 3. 初始化像素值为数据集均值的图像,并将原图像随机粘贴于该图像上。
  618. 4. 根据原图像粘贴位置换算出扩张后真实标注框的位置坐标。
  619. Args:
  620. max_ratio (float): 图像扩张的最大比例。默认为4.0。
  621. prob (float): 随机扩张的概率。默认为0.5。
  622. mean (list): 图像数据集的均值(0-255)。默认为[127.5, 127.5, 127.5]。
  623. """
  624. def __init__(self, max_ratio=4., prob=0.5, mean=[127.5, 127.5, 127.5]):
  625. self.max_ratio = max_ratio
  626. self.mean = mean
  627. self.prob = prob
  628. def __call__(self, im, im_info=None, label_info=None):
  629. """
  630. Args:
  631. im (np.ndarray): 图像np.ndarray数据。
  632. im_info (dict, 可选): 存储与图像相关的信息。
  633. label_info (dict, 可选): 存储与标注框相关的信息。
  634. Returns:
  635. tuple: 当label_info为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  636. 当label_info不为空时,返回的tuple为(im, im_info, label_info),分别对应图像np.ndarray数据、
  637. 存储与标注框相关信息的字典。
  638. 其中,im_info更新字段为:
  639. - augment_shape (np.ndarray): 扩张后的图像高、宽二者组成的np.ndarray,形状为(2,)。
  640. label_info更新字段为:
  641. - gt_bbox (np.ndarray): 随机扩张后真实标注框坐标,形状为(n, 4),
  642. 其中n代表真实标注框的个数。
  643. - gt_class (np.ndarray): 随机扩张后每个真实标注框对应的类别序号,形状为(n, 1),
  644. 其中n代表真实标注框的个数。
  645. Raises:
  646. TypeError: 形参数据类型不满足需求。
  647. """
  648. if im_info is None or label_info is None:
  649. raise TypeError(
  650. 'Cannot do RandomExpand! ' +
  651. 'Becasuse the im_info and label_info can not be None!')
  652. if 'augment_shape' not in im_info:
  653. raise TypeError('Cannot do RandomExpand! ' + \
  654. 'Becasuse augment_shape is not in im_info!')
  655. if 'gt_bbox' not in label_info or \
  656. 'gt_class' not in label_info:
  657. raise TypeError('Cannot do RandomExpand! ' + \
  658. 'Becasuse gt_bbox/gt_class is not in label_info!')
  659. prob = np.random.uniform(0, 1)
  660. augment_shape = im_info['augment_shape']
  661. im_width = augment_shape[1]
  662. im_height = augment_shape[0]
  663. gt_bbox = label_info['gt_bbox']
  664. gt_class = label_info['gt_class']
  665. if prob < self.prob:
  666. if self.max_ratio - 1 >= 0.01:
  667. expand_ratio = np.random.uniform(1, self.max_ratio)
  668. height = int(im_height * expand_ratio)
  669. width = int(im_width * expand_ratio)
  670. h_off = math.floor(np.random.uniform(0, height - im_height))
  671. w_off = math.floor(np.random.uniform(0, width - im_width))
  672. expand_bbox = [
  673. -w_off / im_width, -h_off / im_height,
  674. (width - w_off) / im_width, (height - h_off) / im_height
  675. ]
  676. expand_im = np.ones((height, width, 3))
  677. expand_im = np.uint8(expand_im * np.squeeze(self.mean))
  678. expand_im = Image.fromarray(expand_im)
  679. im = im.astype('uint8')
  680. im = Image.fromarray(im)
  681. expand_im.paste(im, (int(w_off), int(h_off)))
  682. expand_im = np.asarray(expand_im)
  683. for i in range(gt_bbox.shape[0]):
  684. gt_bbox[i][0] = gt_bbox[i][0] / im_width
  685. gt_bbox[i][1] = gt_bbox[i][1] / im_height
  686. gt_bbox[i][2] = gt_bbox[i][2] / im_width
  687. gt_bbox[i][3] = gt_bbox[i][3] / im_height
  688. gt_bbox, gt_class, _ = filter_and_process(
  689. expand_bbox, gt_bbox, gt_class)
  690. for i in range(gt_bbox.shape[0]):
  691. gt_bbox[i][0] = gt_bbox[i][0] * width
  692. gt_bbox[i][1] = gt_bbox[i][1] * height
  693. gt_bbox[i][2] = gt_bbox[i][2] * width
  694. gt_bbox[i][3] = gt_bbox[i][3] * height
  695. im = expand_im.astype('float32')
  696. label_info['gt_bbox'] = gt_bbox
  697. label_info['gt_class'] = gt_class
  698. im_info['augment_shape'] = np.array([height,
  699. width]).astype('int32')
  700. if label_info is None:
  701. return (im, im_info)
  702. else:
  703. return (im, im_info, label_info)
  704. class RandomCrop:
  705. """随机裁剪图像。
  706. 1. 根据batch_sampler计算获取裁剪候选区域的位置。
  707. (1) 根据min scale、max scale、min aspect ratio、max aspect ratio计算随机剪裁的高、宽。
  708. (2) 根据随机剪裁的高、宽随机选取剪裁的起始点。
  709. (3) 筛选出裁剪候选区域:
  710. - 当satisfy_all为True时,需所有真实标注框与裁剪候选区域的重叠度满足需求时,该裁剪候选区域才可保留。
  711. - 当satisfy_all为False时,当有一个真实标注框与裁剪候选区域的重叠度满足需求时,该裁剪候选区域就可保留。
  712. 2. 遍历所有裁剪候选区域:
  713. (1) 若真实标注框与候选裁剪区域不重叠,或其中心点不在候选裁剪区域,
  714. 则将该真实标注框去除。
  715. (2) 计算相对于该候选裁剪区域,真实标注框的位置,并筛选出对应的类别、混合得分。
  716. (3) 若avoid_no_bbox为False,返回当前裁剪后的信息即可;
  717. 反之,要找到一个裁剪区域中真实标注框个数不为0的区域,才返回裁剪后的信息。
  718. Args:
  719. batch_sampler (list): 随机裁剪参数的多种组合,每种组合包含8个值,如下:
  720. - max sample (int):满足当前组合的裁剪区域的个数上限。
  721. - max trial (int): 查找满足当前组合的次数。
  722. - min scale (float): 裁剪面积相对原面积,每条边缩短比例的最小限制。
  723. - max scale (float): 裁剪面积相对原面积,每条边缩短比例的最大限制。
  724. - min aspect ratio (float): 裁剪后短边缩放比例的最小限制。
  725. - max aspect ratio (float): 裁剪后短边缩放比例的最大限制。
  726. - min overlap (float): 真实标注框与裁剪图像重叠面积的最小限制。
  727. - max overlap (float): 真实标注框与裁剪图像重叠面积的最大限制。
  728. 默认值为None,当为None时采用如下设置:
  729. [[1, 1, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0],
  730. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.1, 1.0],
  731. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.3, 1.0],
  732. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.5, 1.0],
  733. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.7, 1.0],
  734. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.9, 1.0],
  735. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.0, 1.0]]
  736. satisfy_all (bool): 是否需要所有标注框满足条件,裁剪候选区域才保留。默认为False。
  737. avoid_no_bbox (bool): 是否对裁剪图像不存在标注框的图像进行保留。默认为True。
  738. """
  739. def __init__(self,
  740. batch_sampler=None,
  741. satisfy_all=False,
  742. avoid_no_bbox=True):
  743. if batch_sampler is None:
  744. batch_sampler = [[1, 1, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
  745. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.1, 1.0],
  746. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.3, 1.0],
  747. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.5, 1.0],
  748. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.7, 1.0],
  749. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.9, 1.0],
  750. [1, 50, 0.3, 1.0, 0.5, 2.0, 0.0, 1.0]]
  751. self.batch_sampler = batch_sampler
  752. self.satisfy_all = satisfy_all
  753. self.avoid_no_bbox = avoid_no_bbox
  754. def __call__(self, im, im_info=None, label_info=None):
  755. """
  756. Args:
  757. im (np.ndarray): 图像np.ndarray数据。
  758. im_info (dict, 可选): 存储与图像相关的信息。
  759. label_info (dict, 可选): 存储与标注框相关的信息。
  760. Returns:
  761. tuple: 当label_info为空时,返回的tuple为(im, im_info),分别对应图像np.ndarray数据、存储与图像相关信息的字典;
  762. 当label_info不为空时,返回的tuple为(im, im_info, label_info),分别对应图像np.ndarray数据、
  763. 存储与标注框相关信息的字典。
  764. 其中,label_info更新字段为:
  765. - gt_bbox (np.ndarray): 随机裁剪后真实标注框坐标,形状为(n, 4),
  766. 其中n代表真实标注框的个数。
  767. - gt_class (np.ndarray): 随机裁剪后每个真实标注框对应的类别序号,形状为(n, 1),
  768. 其中n代表真实标注框的个数。
  769. - gt_score (np.ndarray): 随机裁剪后每个真实标注框对应的混合得分,形状为(n, 1),
  770. 其中n代表真实标注框的个数。
  771. Raises:
  772. TypeError: 形参数据类型不满足需求。
  773. """
  774. if im_info is None or label_info is None:
  775. raise TypeError(
  776. 'Cannot do RandomCrop! ' +
  777. 'Becasuse the im_info and label_info can not be None!')
  778. if 'augment_shape' not in im_info:
  779. raise TypeError('Cannot do RandomCrop! ' + \
  780. 'Becasuse augment_shape is not in im_info!')
  781. if 'gt_bbox' not in label_info or \
  782. 'gt_class' not in label_info:
  783. raise TypeError('Cannot do RandomCrop! ' + \
  784. 'Becasuse gt_bbox/gt_class is not in label_info!')
  785. augment_shape = im_info['augment_shape']
  786. im_width = augment_shape[1]
  787. im_height = augment_shape[0]
  788. gt_bbox = label_info['gt_bbox']
  789. gt_bbox_tmp = gt_bbox.copy()
  790. for i in range(gt_bbox_tmp.shape[0]):
  791. gt_bbox_tmp[i][0] = gt_bbox[i][0] / im_width
  792. gt_bbox_tmp[i][1] = gt_bbox[i][1] / im_height
  793. gt_bbox_tmp[i][2] = gt_bbox[i][2] / im_width
  794. gt_bbox_tmp[i][3] = gt_bbox[i][3] / im_height
  795. gt_class = label_info['gt_class']
  796. gt_score = None
  797. if 'gt_score' in label_info:
  798. gt_score = label_info['gt_score']
  799. sampled_bbox = []
  800. gt_bbox_tmp = gt_bbox_tmp.tolist()
  801. for sampler in self.batch_sampler:
  802. found = 0
  803. for i in range(sampler[1]):
  804. if found >= sampler[0]:
  805. break
  806. sample_bbox = generate_sample_bbox(sampler)
  807. if satisfy_sample_constraint(sampler, sample_bbox, gt_bbox_tmp,
  808. self.satisfy_all):
  809. sampled_bbox.append(sample_bbox)
  810. found = found + 1
  811. im = np.array(im)
  812. while sampled_bbox:
  813. idx = int(np.random.uniform(0, len(sampled_bbox)))
  814. sample_bbox = sampled_bbox.pop(idx)
  815. sample_bbox = clip_bbox(sample_bbox)
  816. crop_bbox, crop_class, crop_score = \
  817. filter_and_process(sample_bbox, gt_bbox_tmp, gt_class, gt_score)
  818. if self.avoid_no_bbox:
  819. if len(crop_bbox) < 1:
  820. continue
  821. xmin = int(sample_bbox[0] * im_width)
  822. xmax = int(sample_bbox[2] * im_width)
  823. ymin = int(sample_bbox[1] * im_height)
  824. ymax = int(sample_bbox[3] * im_height)
  825. im = im[ymin:ymax, xmin:xmax]
  826. for i in range(crop_bbox.shape[0]):
  827. crop_bbox[i][0] = crop_bbox[i][0] * (xmax - xmin)
  828. crop_bbox[i][1] = crop_bbox[i][1] * (ymax - ymin)
  829. crop_bbox[i][2] = crop_bbox[i][2] * (xmax - xmin)
  830. crop_bbox[i][3] = crop_bbox[i][3] * (ymax - ymin)
  831. label_info['gt_bbox'] = crop_bbox
  832. label_info['gt_class'] = crop_class
  833. label_info['gt_score'] = crop_score
  834. im_info['augment_shape'] = np.array([ymax - ymin,
  835. xmax - xmin]).astype('int32')
  836. if label_info is None:
  837. return (im, im_info)
  838. else:
  839. return (im, im_info, label_info)
  840. if label_info is None:
  841. return (im, im_info)
  842. else:
  843. return (im, im_info, label_info)
  844. class ArrangeFasterRCNN:
  845. """获取FasterRCNN模型训练/验证/预测所需信息。
  846. Args:
  847. mode (str): 指定数据用于何种用途,取值范围为['train', 'eval', 'test', 'quant']。
  848. Raises:
  849. ValueError: mode的取值不在['train', 'eval', 'test', 'quant']之内。
  850. """
  851. def __init__(self, mode=None):
  852. if mode not in ['train', 'eval', 'test', 'quant']:
  853. raise ValueError(
  854. "mode must be in ['train', 'eval', 'test', 'quant']!")
  855. self.mode = mode
  856. def __call__(self, im, im_info=None, label_info=None):
  857. """
  858. Args:
  859. im (np.ndarray): 图像np.ndarray数据。
  860. im_info (dict, 可选): 存储与图像相关的信息。
  861. label_info (dict, 可选): 存储与标注框相关的信息。
  862. Returns:
  863. tuple: 当mode为'train'时,返回(im, im_resize_info, gt_bbox, gt_class, is_crowd),分别对应
  864. 图像np.ndarray数据、图像相当对于原图的resize信息、真实标注框、真实标注框对应的类别、真实标注框内是否是一组对象;
  865. 当mode为'eval'时,返回(im, im_resize_info, im_id, im_shape, gt_bbox, gt_class, is_difficult),
  866. 分别对应图像np.ndarray数据、图像相当对于原图的resize信息、图像id、图像大小信息、真实标注框、真实标注框对应的类别、
  867. 真实标注框是否为难识别对象;当mode为'test'或'quant'时,返回(im, im_resize_info, im_shape),分别对应图像np.ndarray数据、
  868. 图像相当对于原图的resize信息、图像大小信息。
  869. Raises:
  870. TypeError: 形参数据类型不满足需求。
  871. ValueError: 数据长度不匹配。
  872. """
  873. im = permute(im, False)
  874. if self.mode == 'train':
  875. if im_info is None or label_info is None:
  876. raise TypeError(
  877. 'Cannot do ArrangeFasterRCNN! ' +
  878. 'Becasuse the im_info and label_info can not be None!')
  879. if len(label_info['gt_bbox']) != len(label_info['gt_class']):
  880. raise ValueError("gt num mismatch: bbox and class.")
  881. im_resize_info = im_info['im_resize_info']
  882. gt_bbox = label_info['gt_bbox']
  883. gt_class = label_info['gt_class']
  884. is_crowd = label_info['is_crowd']
  885. outputs = (im, im_resize_info, gt_bbox, gt_class, is_crowd)
  886. elif self.mode == 'eval':
  887. if im_info is None or label_info is None:
  888. raise TypeError(
  889. 'Cannot do ArrangeFasterRCNN! ' +
  890. 'Becasuse the im_info and label_info can not be None!')
  891. im_resize_info = im_info['im_resize_info']
  892. im_id = im_info['im_id']
  893. im_shape = np.array(
  894. (im_info['augment_shape'][0], im_info['augment_shape'][1], 1),
  895. dtype=np.float32)
  896. gt_bbox = label_info['gt_bbox']
  897. gt_class = label_info['gt_class']
  898. is_difficult = label_info['difficult']
  899. outputs = (im, im_resize_info, im_id, im_shape, gt_bbox, gt_class,
  900. is_difficult)
  901. else:
  902. if im_info is None:
  903. raise TypeError('Cannot do ArrangeFasterRCNN! ' +
  904. 'Becasuse the im_info can not be None!')
  905. im_resize_info = im_info['im_resize_info']
  906. im_shape = np.array(
  907. (im_info['augment_shape'][0], im_info['augment_shape'][1], 1),
  908. dtype=np.float32)
  909. outputs = (im, im_resize_info, im_shape)
  910. return outputs
  911. class ArrangeMaskRCNN:
  912. """获取MaskRCNN模型训练/验证/预测所需信息。
  913. Args:
  914. mode (str): 指定数据用于何种用途,取值范围为['train', 'eval', 'test', 'quant']。
  915. Raises:
  916. ValueError: mode的取值不在['train', 'eval', 'test', 'quant']之内。
  917. """
  918. def __init__(self, mode=None):
  919. if mode not in ['train', 'eval', 'test', 'quant']:
  920. raise ValueError(
  921. "mode must be in ['train', 'eval', 'test', 'quant']!")
  922. self.mode = mode
  923. def __call__(self, im, im_info=None, label_info=None):
  924. """
  925. Args:
  926. im (np.ndarray): 图像np.ndarray数据。
  927. im_info (dict, 可选): 存储与图像相关的信息。
  928. label_info (dict, 可选): 存储与标注框相关的信息。
  929. Returns:
  930. tuple: 当mode为'train'时,返回(im, im_resize_info, gt_bbox, gt_class, is_crowd, gt_masks),分别对应
  931. 图像np.ndarray数据、图像相当对于原图的resize信息、真实标注框、真实标注框对应的类别、真实标注框内是否是一组对象、
  932. 真实分割区域;当mode为'eval'时,返回(im, im_resize_info, im_id, im_shape),分别对应图像np.ndarray数据、
  933. 图像相当对于原图的resize信息、图像id、图像大小信息;当mode为'test'或'quant'时,返回(im, im_resize_info, im_shape),
  934. 分别对应图像np.ndarray数据、图像相当对于原图的resize信息、图像大小信息。
  935. Raises:
  936. TypeError: 形参数据类型不满足需求。
  937. ValueError: 数据长度不匹配。
  938. """
  939. im = permute(im, False)
  940. if self.mode == 'train':
  941. if im_info is None or label_info is None:
  942. raise TypeError(
  943. 'Cannot do ArrangeTrainMaskRCNN! ' +
  944. 'Becasuse the im_info and label_info can not be None!')
  945. if len(label_info['gt_bbox']) != len(label_info['gt_class']):
  946. raise ValueError("gt num mismatch: bbox and class.")
  947. im_resize_info = im_info['im_resize_info']
  948. gt_bbox = label_info['gt_bbox']
  949. gt_class = label_info['gt_class']
  950. is_crowd = label_info['is_crowd']
  951. assert 'gt_poly' in label_info
  952. segms = label_info['gt_poly']
  953. if len(segms) != 0:
  954. assert len(segms) == is_crowd.shape[0]
  955. gt_masks = []
  956. valid = True
  957. for i in range(len(segms)):
  958. segm = segms[i]
  959. gt_segm = []
  960. if is_crowd[i]:
  961. gt_segm.append([[0, 0]])
  962. else:
  963. for poly in segm:
  964. if len(poly) == 0:
  965. valid = False
  966. break
  967. gt_segm.append(np.array(poly).reshape(-1, 2))
  968. if (not valid) or len(gt_segm) == 0:
  969. break
  970. gt_masks.append(gt_segm)
  971. outputs = (im, im_resize_info, gt_bbox, gt_class, is_crowd,
  972. gt_masks)
  973. else:
  974. if im_info is None:
  975. raise TypeError('Cannot do ArrangeMaskRCNN! ' +
  976. 'Becasuse the im_info can not be None!')
  977. im_resize_info = im_info['im_resize_info']
  978. im_shape = np.array(
  979. (im_info['augment_shape'][0], im_info['augment_shape'][1], 1),
  980. dtype=np.float32)
  981. if self.mode == 'eval':
  982. im_id = im_info['im_id']
  983. outputs = (im, im_resize_info, im_id, im_shape)
  984. else:
  985. outputs = (im, im_resize_info, im_shape)
  986. return outputs
  987. class ArrangeYOLOv3:
  988. """获取YOLOv3模型训练/验证/预测所需信息。
  989. Args:
  990. mode (str): 指定数据用于何种用途,取值范围为['train', 'eval', 'test', 'quant']。
  991. Raises:
  992. ValueError: mode的取值不在['train', 'eval', 'test', 'quant']之内。
  993. """
  994. def __init__(self, mode=None):
  995. if mode not in ['train', 'eval', 'test', 'quant']:
  996. raise ValueError(
  997. "mode must be in ['train', 'eval', 'test', 'quant']!")
  998. self.mode = mode
  999. def __call__(self, im, im_info=None, label_info=None):
  1000. """
  1001. Args:
  1002. im (np.ndarray): 图像np.ndarray数据。
  1003. im_info (dict, 可选): 存储与图像相关的信息。
  1004. label_info (dict, 可选): 存储与标注框相关的信息。
  1005. Returns:
  1006. tuple: 当mode为'train'时,返回(im, gt_bbox, gt_class, gt_score, im_shape),分别对应
  1007. 图像np.ndarray数据、真实标注框、真实标注框对应的类别、真实标注框混合得分、图像大小信息;
  1008. 当mode为'eval'时,返回(im, im_shape, im_id, gt_bbox, gt_class, difficult),
  1009. 分别对应图像np.ndarray数据、图像大小信息、图像id、真实标注框、真实标注框对应的类别、
  1010. 真实标注框是否为难识别对象;当mode为'test'或'quant'时,返回(im, im_shape),
  1011. 分别对应图像np.ndarray数据、图像大小信息。
  1012. Raises:
  1013. TypeError: 形参数据类型不满足需求。
  1014. ValueError: 数据长度不匹配。
  1015. """
  1016. im = permute(im, False)
  1017. if self.mode == 'train':
  1018. if im_info is None or label_info is None:
  1019. raise TypeError(
  1020. 'Cannot do ArrangeYolov3! ' +
  1021. 'Becasuse the im_info and label_info can not be None!')
  1022. im_shape = im_info['augment_shape']
  1023. if len(label_info['gt_bbox']) != len(label_info['gt_class']):
  1024. raise ValueError("gt num mismatch: bbox and class.")
  1025. if len(label_info['gt_bbox']) != len(label_info['gt_score']):
  1026. raise ValueError("gt num mismatch: bbox and score.")
  1027. gt_bbox = np.zeros((50, 4), dtype=im.dtype)
  1028. gt_class = np.zeros((50, ), dtype=np.int32)
  1029. gt_score = np.zeros((50, ), dtype=im.dtype)
  1030. gt_num = min(50, len(label_info['gt_bbox']))
  1031. if gt_num > 0:
  1032. label_info['gt_class'][:gt_num, 0] = label_info[
  1033. 'gt_class'][:gt_num, 0] - 1
  1034. gt_bbox[:gt_num, :] = label_info['gt_bbox'][:gt_num, :]
  1035. gt_class[:gt_num] = label_info['gt_class'][:gt_num, 0]
  1036. gt_score[:gt_num] = label_info['gt_score'][:gt_num, 0]
  1037. # parse [x1, y1, x2, y2] to [x, y, w, h]
  1038. gt_bbox[:, 2:4] = gt_bbox[:, 2:4] - gt_bbox[:, :2]
  1039. gt_bbox[:, :2] = gt_bbox[:, :2] + gt_bbox[:, 2:4] / 2.
  1040. outputs = (im, gt_bbox, gt_class, gt_score, im_shape)
  1041. elif self.mode == 'eval':
  1042. if im_info is None or label_info is None:
  1043. raise TypeError(
  1044. 'Cannot do ArrangeYolov3! ' +
  1045. 'Becasuse the im_info and label_info can not be None!')
  1046. im_shape = im_info['augment_shape']
  1047. if len(label_info['gt_bbox']) != len(label_info['gt_class']):
  1048. raise ValueError("gt num mismatch: bbox and class.")
  1049. im_id = im_info['im_id']
  1050. gt_bbox = np.zeros((50, 4), dtype=im.dtype)
  1051. gt_class = np.zeros((50, ), dtype=np.int32)
  1052. difficult = np.zeros((50, ), dtype=np.int32)
  1053. gt_num = min(50, len(label_info['gt_bbox']))
  1054. if gt_num > 0:
  1055. label_info['gt_class'][:gt_num, 0] = label_info[
  1056. 'gt_class'][:gt_num, 0] - 1
  1057. gt_bbox[:gt_num, :] = label_info['gt_bbox'][:gt_num, :]
  1058. gt_class[:gt_num] = label_info['gt_class'][:gt_num, 0]
  1059. difficult[:gt_num] = label_info['difficult'][:gt_num, 0]
  1060. outputs = (im, im_shape, im_id, gt_bbox, gt_class, difficult)
  1061. else:
  1062. if im_info is None:
  1063. raise TypeError('Cannot do ArrangeYolov3! ' +
  1064. 'Becasuse the im_info can not be None!')
  1065. im_shape = im_info['augment_shape']
  1066. outputs = (im, im_shape)
  1067. return outputs