det_transforms.py 54 KB

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