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)
