| 421 | return log |
| 422 | |
| 423 | def test_step(self, batch, batch_idx): |
| 424 | ts_w = batch['ts_w'] |
| 425 | grid = batch['grid'] |
| 426 | mk = batch['masks'] |
| 427 | grid_c = batch['grid_c'] |
| 428 | W, H = self.hparams.img_wh |
| 429 | self.seq_len = batch['seq_len'] |
| 430 | if self.hparams.canonical_dir is not None: |
| 431 | self.canonical_img = batch['canonical_img'] |
| 432 | self.img_wh = batch['img_wh'] |
| 433 | |
| 434 | save_dir = os.path.join('results', |
| 435 | self.hparams.root_dir.split('/')[0], |
| 436 | self.hparams.root_dir.split('/')[1], |
| 437 | self.hparams.exp_name) |
| 438 | sample_name = self.hparams.root_dir.split('/')[1] |
| 439 | if self.hparams.canonical_dir is not None: |
| 440 | test_dir = f'{save_dir}_transformed' |
| 441 | video_name = f'{sample_name}_{self.hparams.exp_name}_transformed' |
| 442 | else: |
| 443 | test_dir = f'{save_dir}' |
| 444 | video_name = f'{sample_name}_{self.hparams.exp_name}' |
| 445 | Path(test_dir).mkdir(parents=True, exist_ok=True) |
| 446 | |
| 447 | if batch_idx > 0 and self.hparams.save_video: |
| 448 | self.video_visualizer.set_path(os.path.join( |
| 449 | test_dir, f'{video_name}.mp4')) |
| 450 | self.raw_video_visualizer.set_path(os.path.join( |
| 451 | test_dir, f'{video_name}_raw.mp4')) |
| 452 | self.dual_video_visualizer.set_path(os.path.join( |
| 453 | test_dir, f'{video_name}_dual.mp4')) |
| 454 | |
| 455 | if batch_idx == 0 and self.hparams.canonical_dir is None: |
| 456 | # Save the canonical image. |
| 457 | ret = self(ts_w, grid_c, False, self.global_step) |
| 458 | |
| 459 | ret_n = self(ts_w, grid, self.hparams.encode_w, self.global_step) |
| 460 | |
| 461 | img = np.zeros((H * W, 3), dtype=np.float32) |
| 462 | for i in range(self.num_models): |
| 463 | if batch_idx == 0 and self.hparams.canonical_dir is None: |
| 464 | results_c = ret.rgbs[i] |
| 465 | if self.hparams.canonical_wh: |
| 466 | img_c = results_c.view(self.hparams.canonical_wh[1], |
| 467 | self.hparams.canonical_wh[0], |
| 468 | 3).float().cpu().numpy() |
| 469 | else: |
| 470 | img_c = results_c.view(H, W, 3).float().cpu().numpy() |
| 471 | |
| 472 | img_c = cv2.cvtColor(img_c, cv2.COLOR_BGR2RGB) |
| 473 | cv2.imwrite(f'{test_dir}/canonical_{i}.png', img_c * 255) |
| 474 | |
| 475 | mk_n = rearrange(mk[i], 'b h w c -> (b h w) c') |
| 476 | mk_n = mk_n.sum(dim=-1) > 0.05 |
| 477 | mk_n = mk_n.cpu().numpy() |
| 478 | results = ret_n.rgbs[i] |
| 479 | results = results.cpu().numpy() # (3, H, W) |
| 480 | img[mk_n] = results[mk_n] |