det_transforms.py 54 KB

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