def inf(self, imgs, img_names, gt, inference, net, scales, pbar, base_img): ###################################################################### # Run inference ###################################################################### self.img_name = img_names[0] col_img_name = '{}/{}_color.png'.format(self.rgb_path, self.img_name) pred_img_name = '{}/{}.png'.format(self.pred_path, self.img_name) diff_img_name = '{}/{}_diff.png'.format(self.diff_path, self.img_name) compose_img_name = '{}/{}_compose.png'.format(self.compose_path, self.img_name) to_pil = transforms.ToPILImage() if self.inference_mode == 'pooling': img = imgs pool_base_img = to_pil(base_img[0]) else: img = to_pil(imgs[0]) prediction_pre_argmax_collection = inference(net, img, scales) if self.inference_mode == 'pooling': prediction = prediction_pre_argmax_collection prediction = np.concatenate(prediction, axis=0)[0] else: prediction_pre_argmax = np.mean(prediction_pre_argmax_collection, axis=0) prediction = np.argmax(prediction_pre_argmax, axis=0) if args.dataset == 'kitti' and args.split == 'test': origin_h, origin_w = 375, 1242 pred_pil = Image.fromarray(prediction.astype('uint8')) pred_pil = pred_pil.resize((origin_w, origin_h), Image.NEAREST) small_img = img.copy() small_img = small_img.resize((origin_w, origin_h), Image.BICUBIC) prediction = np.array(pred_pil) if self.metrics: #可以以这个参数控制不进行MIoU的计算 self.hist += fast_hist(prediction.flatten(), gt.cpu().numpy().flatten(), self.dataset_cls.num_classes) iou = round(np.nanmean(per_class_iu(self.hist)) * 100, 2) pbar.set_description("Mean IOU: %s" % (str(iou))) ###################################################################### # Dump Images ###################################################################### if self.write_image: if self.inference_mode == 'pooling': img = pool_base_img colorized = self.dataset_cls.colorize_mask(prediction) colorized.save(col_img_name) if args.dataset == 'kitti' and args.split == 'test': blend = Image.blend(small_img.convert("RGBA"), colorized.convert("RGBA"), 0.5) else: blend = Image.blend(img.convert("RGBA"), colorized.convert("RGBA"), 0.5) blend.save(compose_img_name) if gt is not None and args.split != 'test': gt = gt[0].cpu().numpy() # only write diff image if gt is valid diff = (prediction != gt) diff[gt == 255] = 0 diffimg = Image.fromarray(diff.astype('uint8') * 255) PIL.ImageChops.lighter( blend, PIL.ImageOps.invert(diffimg).convert("RGBA")).save( diff_img_name) label_out = np.zeros_like(prediction) for label_id, train_id in self.dataset_cls.id_to_trainid.items(): label_out[np.where(prediction == train_id)] = label_id cv2.imwrite(pred_img_name, label_out) #这里将会转换成label_id
def inf(self, imgs, img_names, gt, inference, net, scales, pbar, base_img, pos): ###################################################################### # Run inference ###################################################################### self.img_name = img_names[0] col_img_name = '{}/{}_color.png'.format(self.rgb_path, self.img_name) pred_img_name = '{}/{}.png'.format(self.pred_path, self.img_name) diff_img_name = '{}/{}_diff.png'.format(self.diff_path, self.img_name) compose_img_name = '{}/{}_compose.png'.format(self.compose_path, self.img_name) to_pil = transforms.ToPILImage() if self.inference_mode == 'pooling': img = imgs pool_base_img = to_pil(base_img[0]) else: img = to_pil(imgs[0]) prediction_pre_argmax_collection = inference(net, img, scales, pos) # print(len(prediction_pre_argmax_collection)) # print(prediction_pre_argmax_collection[0].shape) if self.inference_mode == 'pooling': prediction = prediction_pre_argmax_collection prediction = np.concatenate(prediction, axis=0) else: prediction_pre_argmax = np.mean(prediction_pre_argmax_collection, axis=0) prediction = np.argmax(prediction_pre_argmax, axis=0) if self.metrics: self.hist += fast_hist(prediction.flatten(), gt.cpu().numpy().flatten(), self.dataset_cls.num_classes) iou_w = round(np.nanmean(per_class_iu(self.hist)) * 100, 2) # acc_w = np.diag(self.hist).sum() / self.hist.sum() H, W = prediction.shape pred_split = np.split(prediction, [H // 4, (H // 4) * 2, (H // 4) * 3], axis=0) gt_split = np.split(gt.cpu().numpy(), [H // 4, (H // 4) * 2, (H // 4) * 3], axis=1) self.hist_up += fast_hist(pred_split[0].flatten(), gt_split[0].flatten(), self.dataset_cls.num_classes) iou_u = round(np.nanmean(per_class_iu(self.hist_up)) * 100, 2) # acc_u = np.diag(self.hist_up).sum() / self.hist_up.sum() self.hist_mid1 += fast_hist(pred_split[1].flatten(), gt_split[1].flatten(), self.dataset_cls.num_classes) iou_m1 = round(np.nanmean(per_class_iu(self.hist_mid1)) * 100, 2) # acc_m1 = np.diag(self.hist_mid1).sum() / self.hist_mid1.sum() self.hist_mid2 += fast_hist(pred_split[2].flatten(), gt_split[2].flatten(), self.dataset_cls.num_classes) iou_m2 = round(np.nanmean(per_class_iu(self.hist_mid2)) * 100, 2) # acc_m2 = np.diag(self.hist_mid2).sum() / self.hist_mid2.sum() self.hist_down += fast_hist(pred_split[3].flatten(), gt_split[3].flatten(), self.dataset_cls.num_classes) iou_d = round(np.nanmean(per_class_iu(self.hist_down)) * 100, 2) # acc_d = np.diag(self.hist_down).sum() / self.hist_down.sum() pbar.set_description( "Mean IOU (Whole,Up,Mid1,Mid2,DOWN): %s %s %s %s %s" % (str(iou_w), str(iou_u), str(iou_m1), str(iou_m2), str(iou_d))) # pbar.set_description("ACC (Whole,Up,Mid,DOWN): %s %s %s %s" % (str(acc_w), str(acc_u), str(acc_m), str(acc_d))) ###################################################################### # Dump Images ###################################################################### if self.write_image: if self.inference_mode == 'pooling': img = pool_base_img colorized = self.dataset_cls.colorize_mask(prediction) colorized.save(col_img_name) blend = Image.blend(img.convert("RGBA"), colorized.convert("RGBA"), 0.5) blend.save(compose_img_name) if gt is not None: gt = gt[0].cpu().numpy() # only write diff image if gt is valid diff = (prediction != gt) diff[gt == 255] = 0 diffimg = Image.fromarray(diff.astype('uint8') * 255) PIL.ImageChops.lighter( blend, PIL.ImageOps.invert(diffimg).convert("RGBA")).save( diff_img_name) label_out = np.zeros_like(prediction) for label_id, train_id in self.dataset_cls.id_to_trainid.items(): label_out[np.where(prediction == train_id)] = label_id cv2.imwrite(pred_img_name, label_out)