def eval(self):
self.metric.reset()
self.model.eval()
if self.args.distributed:
model = self.model.module
else:
model = self.model
logger.info("Start validation, Total sample: {:d}".format(
len(self.val_loader)))
for i, (image, target) in enumerate(self.val_loader):
image = image.to(self.device)
target = target.to(self.device)
with torch.no_grad():
outputs = model(image)
self.metric.update(outputs[0], target)
pixAcc, mIoU = self.metric.get()
logger.info(
"Sample: {:d}, validation pixAcc: {:.3f}, mIoU: {:.3f}".format(
i + 1, pixAcc * 100, mIoU * 100))
if self.args.save_pred:
pred = torch.argmax(outputs[0], 1)
pred = pred.cpu().data.numpy()
predict = pred.squeeze(0)
mask = get_color_pallete(predict, self.args.dataset)
# mask.save(os.path.join(outdir, os.path.splitext(filename[0])[0] + '.png'))
synchronize()
def eval(self):
self.metric.reset()
self.model.eval()
if self.args.distributed:
model = self.model.module
else:
model = self.model
logger.info("Start validation, Total sample: {:d}".format(
len(self.val_loader)))
for i, (image, target) in enumerate(self.val_loader):
img = data_process(image)
img = img.to(self.device)
target = target.to(self.device)
with torch.no_grad():
outputs = model(img)
self.metric.update(outputs, target)
pixAcc, mIoU = self.metric.get()
logger.info(
"Sample: {:d}, validation pixAcc: {:.3f}, mIoU: {:.3f}".format(
i + 1, pixAcc * 100, mIoU * 100))
if self.args.save_pred:
pred = torch.argmax(outputs[0], 1)
pred = pred.cpu().data.numpy()
predict = pred.squeeze(0)
img_show = image[0].numpy()
img_show = img_show.astype('uint8')
plt.subplot(1, 3, 1)
plt.title('image')
plt.imshow(img_show)
mask = target.cpu().data.numpy()
mask = mask.reshape(mask.shape[1], mask.shape[2])
mask = mask_to_image(mask)
plt.subplot(1, 3, 2)
plt.title('mask')
plt.imshow(mask)
predict = mask_to_image(predict)
plt.subplot(1, 3, 3)
plt.title('pred')
plt.imshow(predict)
save_file = "save_fig_val"
os.makedirs(save_file, exist_ok=True)
plt.savefig(os.path.join(save_file, str(i) + '.png'))
synchronize()
def eval(self):
self.metric.reset()
self.model.eval()
if self.args.distributed:
model = self.model.module
else:
model = self.model
logger.info("Start validation, Total sample: {:d}".format(
len(self.val_loader)))
for i, (image, target) in enumerate(self.val_loader):
image = image.to(self.device)
target = target.to(self.device)
with torch.no_grad():
outputs = model(image)
self.metric.update(outputs[0], target)
pixAcc, mIoU = self.metric.get()
logger.info(
"Sample: {:d}, validation pixAcc: {:.3f}, mIoU: {:.3f}".format(
i + 1, pixAcc * 100, mIoU * 100))
if True:
pred = torch.argmax(outputs[0], 1)
pred = pred.cpu().data.numpy()
predict = pred.squeeze(0)
# mask = get_color_pallete(predict, self.args.dataset)
image = image.cpu().data.numpy().squeeze(0).transpose(
(1, 2, 0))
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
colors = np.array([[255, 0, 0], [0, 255, 0], [0, 0, 255]])
res = np.zeros((image.shape[0] * 3, image.shape[1], 3))
inp = ((image + 1) * 127.5).astype('int')
msk = colors[predict]
res[0:image.shape[0], :, :] = inp
res[image.shape[0]:image.shape[0] * 2, :, :] = msk
res[image.shape[0] * 2:, :, :] = cv2.addWeighted(
inp, 0.5, msk, 0.5, 0)
cv2.imwrite(
f'/root/mitya/Lightweight-Segmentation/results/{i}.png',
res)
synchronize()
if __name__ == '__main__':
args = parse_args()
num_gpus = int(
os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
args.distributed = num_gpus > 1
if not args.no_cuda and torch.cuda.is_available():
cudnn.benchmark = True
args.device = "cuda"
else:
args.distributed = False
args.device = "cpu"
if args.distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl",
init_method="env://")
synchronize()
# TODO: optim code
args.save_pred = False
if args.save_pred:
outdir = '../runs/pred_pic/{}_{}'.format(args.model, args.dataset)
if not os.path.exists(outdir):
os.makedirs(outdir)
logger = setup_logger(args.model,
args.log_dir,
get_rank(),
filename='{}_{}_log.txt'.format(
args.model, args.dataset),
mode='a+')