add input image in tensorboard

ZERONE182 · May 21, 2023 · f9ab5b4 · f9ab5b4
1 parent 387794b
commit f9ab5b4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/VAE/train.py b/VAE/train.py
@@ -107,11 +107,13 @@ def validation(model, loader_val, writer, step, batch_size=8, device='cuda'):
         ori_img, R, T, K = next(iter(loader_val))
 
         batch = {'x':ori_img[:,0].to(device), 'z':ori_img[:,1].to(device), 'R': R.to(device), 't': T.to(device), 'K': K.to(device),}
+        input_img = ori_img[:, 0].detach().cpu().numpy()
         gt_img = ori_img[:, 1].detach().cpu().numpy()
         gt_img = ((gt_img.clip(-1, 1)+1)*127.5).astype(np.uint8)
         pred_img = model.module.eval_img(batch, None).detach().cpu().numpy()
         pred_img = ((pred_img.clip(-1, 1)+1)*127.5).astype(np.uint8)
 
+        writer.add_images(f"train/input", input_img, step)
         writer.add_images(f"train/gt", gt_img, step)
         writer.add_images(f"train/pred",pred_img, step)