det_transforms.py 54 KB

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