dump a single batch of images :dump_dict: a dictionary containing elements to dump out 'input_images': source image 'gt_images': label 'img_names': img_names 'assets': dict with keys: 'predictions': final prediction 'p
(self, dump_dict, val_idx)
| 277 | self.imgs_to_webpage = [] |
| 278 | |
| 279 | def dump(self, dump_dict, val_idx): |
| 280 | """ |
| 281 | dump a single batch of images |
| 282 | |
| 283 | :dump_dict: a dictionary containing elements to dump out |
| 284 | 'input_images': source image |
| 285 | 'gt_images': label |
| 286 | 'img_names': img_names |
| 287 | 'assets': dict with keys: |
| 288 | 'predictions': final prediction |
| 289 | 'pred_*': different scales of predictions |
| 290 | 'attn_*': different scales of attn |
| 291 | 'err_mask': err_mask |
| 292 | """ |
| 293 | if self.dump_for_auto_labelling or self.dump_for_submission: |
| 294 | pass |
| 295 | elif (val_idx % self.dump_frequency or cfg.GLOBAL_RANK != 0): |
| 296 | return |
| 297 | else: |
| 298 | pass |
| 299 | |
| 300 | colorize_mask_fn = cfg.DATASET_INST.colorize_mask |
| 301 | idx = 0 # only use first element of batch |
| 302 | |
| 303 | input_image = dump_dict['input_images'][idx] |
| 304 | prob_image = dump_dict['assets']['prob_mask'][idx] |
| 305 | gt_image = dump_dict['gt_images'][idx] |
| 306 | prediction = dump_dict['assets']['predictions'][idx] |
| 307 | del dump_dict['assets']['predictions'] |
| 308 | img_name = dump_dict['img_names'][idx] |
| 309 | |
| 310 | if self.dump_for_auto_labelling: |
| 311 | # Dump Prob |
| 312 | prob_fn = '{}_prob.png'.format(img_name) |
| 313 | prob_fn = os.path.join(self.save_dir, prob_fn) |
| 314 | cv2.imwrite(prob_fn, (prob_image.cpu().numpy()*255).astype(np.uint8)) |
| 315 | |
| 316 | if self.dump_for_auto_labelling or self.dump_for_submission: |
| 317 | # Dump Predictions |
| 318 | prediction_cpu = np.array(prediction) |
| 319 | label_out = np.zeros_like(prediction) |
| 320 | submit_fn = '{}.png'.format(img_name) |
| 321 | for label_id, train_id in cfg.DATASET_INST.id_to_trainid.items(): |
| 322 | label_out[np.where(prediction_cpu == train_id)] = label_id |
| 323 | cv2.imwrite(os.path.join(self.save_dir, submit_fn), label_out) |
| 324 | return |
| 325 | |
| 326 | input_image = self.inv_normalize(input_image) |
| 327 | input_image = input_image.cpu() |
| 328 | input_image = standard_transforms.ToPILImage()(input_image) |
| 329 | input_image = input_image.convert("RGB") |
| 330 | input_image_fn = f'{img_name}_input.png' |
| 331 | input_image.save(os.path.join(self.save_dir, input_image_fn)) |
| 332 | |
| 333 | gt_fn = '{}_gt.png'.format(img_name) |
| 334 | gt_pil = colorize_mask_fn(gt_image.cpu().numpy()) |
| 335 | gt_pil.save(os.path.join(self.save_dir, gt_fn)) |
| 336 |
no outgoing calls
no test coverage detected