det_transforms.py 55 KB

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