det_transforms.py 54 KB

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