From f59c4c4828f63b53e8ea96106b9dc38ae2c40acc Mon Sep 17 00:00:00 2001 From: Maguire93 <610716365@qq.com> Date: Thu, 13 Apr 2023 23:09:29 +0800 Subject: [PATCH] commit DM-NeRF train in replica and scannet for sparse exp --- config.py | 19 +- .../replica/train/m40/office_2_sparse_inv.txt | 31 + configs/replica/train/office_0.txt | 5 +- configs/replica/train/office_0_full.txt | 27 + configs/replica/train/office_0_sparse_inv.txt | 30 + configs/replica/train/office_2_full.txt | 28 + configs/replica/train/office_2_sparse_inv.txt | 28 + configs/replica/train/office_3.txt | 12 +- configs/replica/train/office_3_full.txt | 28 + configs/replica/train/office_3_sparse_inv.txt | 30 + configs/replica/train/office_4_full.txt | 27 + configs/replica/train/office_4_sparse_inv.txt | 27 + configs/replica/train/room_0_full.txt | 27 + configs/replica/train/room_0_sparse_inv.txt | 26 + configs/replica/train/room_1_full.txt | 28 + configs/replica/train/room_1_sparse_inv.txt | 29 + configs/replica/train/room_2_full.txt | 28 + configs/replica/train/room_2_sparse_inv.txt | 30 + configs/scannet/train/scene0010_00_full.txt | 30 + configs/scannet/train/scene0010_00_sparse.txt | 31 + configs/scannet/train/scene0012_00_full.txt | 31 + configs/scannet/train/scene0012_00_sparse.txt | 31 + configs/scannet/train/scene0024_00_sparse.txt | 34 + configs/scannet/train/scene0033_00_sparse.txt | 35 + configs/scannet/train/scene0038_00_full.txt | 31 + configs/scannet/train/scene0038_00_sparse.txt | 31 + configs/scannet/train/scene0088_00_full.txt | 31 + configs/scannet/train/scene0088_00_sparse.txt | 31 + configs/scannet/train/scene0113_00_full.txt | 30 + configs/scannet/train/scene0113_00_sparse.txt | 31 + configs/scannet/train/scene0192_00_full.txt | 33 + configs/scannet/train/scene0192_00_sparse.txt | 33 + data/color_dict-old.json | 661 ++++++++++++++++++ data/color_dict.json | 19 +- datasets/loader_scannet.py | 37 +- mani_replica.py | 69 ++ networks/evaluator.py | 112 ++- networks/render.py | 1 + networks/tester.py | 264 ++++++- reload_train_scannet.py | 183 +++++ requirements.txt | 8 +- retrain_scannet_sparse.sh | 11 + test_replica.py | 5 + test_replica_office_sparse.sh | 9 + test_scannet.py | 4 +- test_scannet.sh | 10 + tools/visualizer.py | 2 +- train_replica.py | 47 +- train_replica.sh | 7 + train_replica_office_full.sh | 10 + train_replica_office_sparse.sh | 9 + train_replica_room_full.sh | 8 + train_replica_room_sparse.sh | 10 + train_scannet.py | 57 +- train_scannet_full.sh | 11 + train_scannet_sparse.sh | 11 + 56 files changed, 2412 insertions(+), 56 deletions(-) create mode 100644 configs/replica/train/m40/office_2_sparse_inv.txt create mode 100644 configs/replica/train/office_0_full.txt create mode 100644 configs/replica/train/office_0_sparse_inv.txt create mode 100644 configs/replica/train/office_2_full.txt create mode 100644 configs/replica/train/office_2_sparse_inv.txt create mode 100644 configs/replica/train/office_3_full.txt create mode 100644 configs/replica/train/office_3_sparse_inv.txt create mode 100644 configs/replica/train/office_4_full.txt create mode 100644 configs/replica/train/office_4_sparse_inv.txt create mode 100644 configs/replica/train/room_0_full.txt create mode 100644 configs/replica/train/room_0_sparse_inv.txt create mode 100644 configs/replica/train/room_1_full.txt create mode 100644 configs/replica/train/room_1_sparse_inv.txt create mode 100644 configs/replica/train/room_2_full.txt create mode 100644 configs/replica/train/room_2_sparse_inv.txt create mode 100644 configs/scannet/train/scene0010_00_full.txt create mode 100644 configs/scannet/train/scene0010_00_sparse.txt create mode 100644 configs/scannet/train/scene0012_00_full.txt create mode 100644 configs/scannet/train/scene0012_00_sparse.txt create mode 100644 configs/scannet/train/scene0024_00_sparse.txt create mode 100644 configs/scannet/train/scene0033_00_sparse.txt create mode 100644 configs/scannet/train/scene0038_00_full.txt create mode 100644 configs/scannet/train/scene0038_00_sparse.txt create mode 100644 configs/scannet/train/scene0088_00_full.txt create mode 100644 configs/scannet/train/scene0088_00_sparse.txt create mode 100644 configs/scannet/train/scene0113_00_full.txt create mode 100644 configs/scannet/train/scene0113_00_sparse.txt create mode 100644 configs/scannet/train/scene0192_00_full.txt create mode 100644 configs/scannet/train/scene0192_00_sparse.txt create mode 100644 data/color_dict-old.json create mode 100644 mani_replica.py create mode 100644 reload_train_scannet.py create mode 100644 retrain_scannet_sparse.sh create mode 100644 test_replica_office_sparse.sh create mode 100644 test_scannet.sh create mode 100644 train_replica.sh create mode 100644 train_replica_office_full.sh create mode 100644 train_replica_office_sparse.sh create mode 100644 train_replica_room_full.sh create mode 100644 train_replica_room_sparse.sh create mode 100644 train_scannet_full.sh create mode 100644 train_scannet_sparse.sh diff --git a/config.py b/config.py index f93802f..44c7edf 100644 --- a/config.py +++ b/config.py @@ -24,6 +24,10 @@ def config_parser(): help='layers in network') parser.add_argument("--netwidth", type=int, default=256, help='channels per layer') + + parser.add_argument("--label_sparse_inv", type=int, default=1, + help='label sparse inv') + # 32*32*4 parser.add_argument("--N_train", type=int, default=4096, help='batch size (number of random rays per gradient step)') @@ -77,7 +81,7 @@ def config_parser(): help='frequency of console printout and metric loggin') parser.add_argument("--i_img", type=int, default=500, help='frequency of tensorboard image logging') - parser.add_argument("--i_save", type=int, default=10000, + parser.add_argument("--i_save", type=int, default=50000, help='frequency of weight ckpt saving') parser.add_argument("--i_test", type=int, default=50000, help='frequency of testset saving') @@ -142,6 +146,19 @@ def initial(): parser = config_parser() args = parser.parse_args() + sparse_inv = args.label_sparse_inv + if args.expname[-4:] != 'full': + expname_list = args.expname.split('_') + num = int(args.expname[-2:]) + try: + flag = (int(expname_list[-1]) == sparse_inv) or (num == sparse_inv) + except: + flag = (num == sparse_inv) + + if not flag: + exp_str = '-new_sparse' + str(sparse_inv).zfill(2) + args.expname += exp_str + # get log time if args.log_time is None: args.log_time = time.strftime("%Y%m%d%H%M", time.localtime()) diff --git a/configs/replica/train/m40/office_2_sparse_inv.txt b/configs/replica/train/m40/office_2_sparse_inv.txt new file mode 100644 index 0000000..c67b9ed --- /dev/null +++ b/configs/replica/train/m40/office_2_sparse_inv.txt @@ -0,0 +1,31 @@ +expname = office_2_sparse10 +# sparse inv +label_sparse_inv = 10 + +basedir = ./logs/replica +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_2 +# audodl-3090 +# datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/office_2 +# m40 +datadir = /media/alpha4TB/exchange/nerf/DM-NeRF/replica/replica/office_2 + + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1024 +N_test = 1024 + + +near = 0.0 +far = 5.7 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 diff --git a/configs/replica/train/office_0.txt b/configs/replica/train/office_0.txt index c261436..6f7e415 100644 --- a/configs/replica/train/office_0.txt +++ b/configs/replica/train/office_0.txt @@ -1,12 +1,13 @@ expname = office_0 basedir = ./logs/replica -datadir = ./data/replica/office_0 +# 2080ti +datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_0 lrate_decay = 500 N_samples = 64 N_importance = 128 -N_train = 3072 +N_train = 1024 N_test = 4096 diff --git a/configs/replica/train/office_0_full.txt b/configs/replica/train/office_0_full.txt new file mode 100644 index 0000000..3897194 --- /dev/null +++ b/configs/replica/train/office_0_full.txt @@ -0,0 +1,27 @@ +expname = office_0_full + +label_sparse_inv = 1 + +basedir = ./logs/replica +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_0 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/office_0 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +near = 0.0 +far = 4.7 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 diff --git a/configs/replica/train/office_0_sparse_inv.txt b/configs/replica/train/office_0_sparse_inv.txt new file mode 100644 index 0000000..34a97ba --- /dev/null +++ b/configs/replica/train/office_0_sparse_inv.txt @@ -0,0 +1,30 @@ + +expname = office_0_sparse10 +# sparse inv +label_sparse_inv = 10 + + +basedir = ./logs/replica +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_0 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/office_0 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + + +near = 0.0 +far = 4.7 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 diff --git a/configs/replica/train/office_2_full.txt b/configs/replica/train/office_2_full.txt new file mode 100644 index 0000000..53f69d7 --- /dev/null +++ b/configs/replica/train/office_2_full.txt @@ -0,0 +1,28 @@ +expname = office_2_full + +label_sparse_inv = 1 + +basedir = ./logs/replica +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_2 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/office_2 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1024 +N_test = 1024 + + +near = 0.0 +far = 5.7 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 diff --git a/configs/replica/train/office_2_sparse_inv.txt b/configs/replica/train/office_2_sparse_inv.txt new file mode 100644 index 0000000..1f79c28 --- /dev/null +++ b/configs/replica/train/office_2_sparse_inv.txt @@ -0,0 +1,28 @@ +expname = office_2_sparse10 +# sparse inv +label_sparse_inv = 10 + +basedir = ./logs/replica +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_2 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/office_2 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1024 +N_test = 1024 + + +near = 0.0 +far = 5.7 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 diff --git a/configs/replica/train/office_3.txt b/configs/replica/train/office_3.txt index adb2912..b5b8dee 100644 --- a/configs/replica/train/office_3.txt +++ b/configs/replica/train/office_3.txt @@ -1,19 +1,25 @@ expname = office_3 basedir = ./logs/replica -datadir = ./data/replica/office_3 +# 2080ti +datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_3 lrate_decay = 500 N_samples = 64 N_importance = 128 -N_train = 3072 +N_train = 1024 N_test = 4096 near = 0.0 far = 6.6 -testskip = 4 +testskip = 1 penalize tolerance = 0.05 deta_w = 0.05 + +# render test step +i_test = -1 + +label_sparse_inv = 1 \ No newline at end of file diff --git a/configs/replica/train/office_3_full.txt b/configs/replica/train/office_3_full.txt new file mode 100644 index 0000000..b996dc8 --- /dev/null +++ b/configs/replica/train/office_3_full.txt @@ -0,0 +1,28 @@ +expname = office_3_full +# sparse inv +label_sparse_inv = 1 + +basedir = ./logs/replica +# 2080ti +datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_3 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/office_3 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 2048 +N_test = 2048 + + +near = 0.0 +far = 6.6 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 diff --git a/configs/replica/train/office_3_sparse_inv.txt b/configs/replica/train/office_3_sparse_inv.txt new file mode 100644 index 0000000..31850fe --- /dev/null +++ b/configs/replica/train/office_3_sparse_inv.txt @@ -0,0 +1,30 @@ +expname = office_3_sparse10 +# sparse inv +label_sparse_inv = 10 + +basedir = ./logs/replica +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/office_3 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_3 + + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1024 +N_test = 1024 + + +near = 0.0 +far = 6.6 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 + diff --git a/configs/replica/train/office_4_full.txt b/configs/replica/train/office_4_full.txt new file mode 100644 index 0000000..d174877 --- /dev/null +++ b/configs/replica/train/office_4_full.txt @@ -0,0 +1,27 @@ +expname = office_4_full + +label_sparse_inv = 1 + +basedir = ./logs/replica +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_4 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/office_4 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +near = 0.0 +far = 6.2 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 diff --git a/configs/replica/train/office_4_sparse_inv.txt b/configs/replica/train/office_4_sparse_inv.txt new file mode 100644 index 0000000..e58a5ba --- /dev/null +++ b/configs/replica/train/office_4_sparse_inv.txt @@ -0,0 +1,27 @@ +expname = office_4_sparse10 + +label_sparse_inv = 10 + +basedir = ./logs/replica +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/office_4 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/office_4 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +near = 0.0 +far = 6.2 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 diff --git a/configs/replica/train/room_0_full.txt b/configs/replica/train/room_0_full.txt new file mode 100644 index 0000000..cf8259c --- /dev/null +++ b/configs/replica/train/room_0_full.txt @@ -0,0 +1,27 @@ +expname = room_0_full + +label_sparse_inv = 1 + +basedir = ./logs/replica/room +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/room_0 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/room_0 +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + + +near = 0.0 +far = 6.5 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/replica/train/room_0_sparse_inv.txt b/configs/replica/train/room_0_sparse_inv.txt new file mode 100644 index 0000000..450ea5b --- /dev/null +++ b/configs/replica/train/room_0_sparse_inv.txt @@ -0,0 +1,26 @@ +expname = room_0_sparse10 + +label_sparse_inv = 10 + +basedir = ./logs/replica/room +# 2080ti +datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/room_0 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 2048 +N_test = 2048 + + +near = 0.0 +far = 6.5 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/replica/train/room_1_full.txt b/configs/replica/train/room_1_full.txt new file mode 100644 index 0000000..02d032b --- /dev/null +++ b/configs/replica/train/room_1_full.txt @@ -0,0 +1,28 @@ +expname = room_1_full + +label_sparse_inv = 1 + +basedir = ./logs/replica/room +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/room_1 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/room_1 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + + +near = 0.0 +far = 5.2 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/replica/train/room_1_sparse_inv.txt b/configs/replica/train/room_1_sparse_inv.txt new file mode 100644 index 0000000..2ac1233 --- /dev/null +++ b/configs/replica/train/room_1_sparse_inv.txt @@ -0,0 +1,29 @@ +expname = room_1_sparse10 + +label_sparse_inv = 10 + +basedir = ./logs/replica/room +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/room_1 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/room_1 + + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + + +near = 0.0 +far = 6.5 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/replica/train/room_2_full.txt b/configs/replica/train/room_2_full.txt new file mode 100644 index 0000000..f0ecd5a --- /dev/null +++ b/configs/replica/train/room_2_full.txt @@ -0,0 +1,28 @@ +expname = room_2_full + +label_sparse_inv = 1 + +basedir = ./logs/replica/room +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/room_2 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/room_2 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + + +near = 0.0 +far = 6.0 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/replica/train/room_2_sparse_inv.txt b/configs/replica/train/room_2_sparse_inv.txt new file mode 100644 index 0000000..5495eb8 --- /dev/null +++ b/configs/replica/train/room_2_sparse_inv.txt @@ -0,0 +1,30 @@ +expname = room_2_sparse10 + +label_sparse_inv = 10 + +basedir = ./logs/replica/room +# 2080ti +datadir = /data/dzy_data/nerf/datasets/nerf_replica/replica_ins/room_2 +# audodl-3090 +# datadir = /root/autodl-tmp/datasets/nerf_replica/replica_ins/room_2 +# xp-9 +# datadir = /data/jtc/zty/dataset/replica_DM-NeRF/room_2 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + + +near = 0.0 +far = 6.0 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0010_00_full.txt b/configs/scannet/train/scene0010_00_full.txt new file mode 100644 index 0000000..877788f --- /dev/null +++ b/configs/scannet/train/scene0010_00_full.txt @@ -0,0 +1,30 @@ +expname = scene0010_00_full +basedir = ./logs/scannet + +# datadir = ./data/scannet/scene0010_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0010_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0010_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 2048 +N_test = 2048 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 7.6 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0010_00_sparse.txt b/configs/scannet/train/scene0010_00_sparse.txt new file mode 100644 index 0000000..ef1ef8b --- /dev/null +++ b/configs/scannet/train/scene0010_00_sparse.txt @@ -0,0 +1,31 @@ +expname = scene0010_00_sparse10 +label_sparse_inv = 10 +basedir = ./logs/scannet + +# datadir = ./data/scannet/scene0010_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0010_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0010_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 7.6 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0012_00_full.txt b/configs/scannet/train/scene0012_00_full.txt new file mode 100644 index 0000000..981d064 --- /dev/null +++ b/configs/scannet/train/scene0012_00_full.txt @@ -0,0 +1,31 @@ +expname = scene0012_00_full +label_sparse_inv = 1 + +basedir = ./logs/scannet +# datadir = ./data/scannet/scene0012_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0012_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0012_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 2048 +N_test = 2048 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 9.5 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0012_00_sparse.txt b/configs/scannet/train/scene0012_00_sparse.txt new file mode 100644 index 0000000..00bc9e4 --- /dev/null +++ b/configs/scannet/train/scene0012_00_sparse.txt @@ -0,0 +1,31 @@ +expname = scene0012_00_sparse10 +label_sparse_inv = 10 + +basedir = ./logs/scannet +# datadir = ./data/scannet/scene0012_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0012_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0012_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 9.5 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0024_00_sparse.txt b/configs/scannet/train/scene0024_00_sparse.txt new file mode 100644 index 0000000..3e8dec6 --- /dev/null +++ b/configs/scannet/train/scene0024_00_sparse.txt @@ -0,0 +1,34 @@ + +expname = scene0024_00_sparse10 +label_sparse_inv = 10 + +basedir = ./logs/scannet +# datadir = ./data/scannet/scene0024_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0024_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0024_00 + + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +crop_width = 600 +crop_height = 450 +resize + +near = 0.0 +far = 9.8 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0033_00_sparse.txt b/configs/scannet/train/scene0033_00_sparse.txt new file mode 100644 index 0000000..0dc41af --- /dev/null +++ b/configs/scannet/train/scene0033_00_sparse.txt @@ -0,0 +1,35 @@ +expname = scene0033_00_sparse10 + +label_sparse_inv = 10 + +basedir = ./logs/scannet +# M40 +# datadir = /media/alpha4TB/exchange/nerf/nerf_scannet/scene0033_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0033_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0033_00 + + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 2048 +N_test = 2048 + +crop_width = 600 +crop_height = 450 +resize + +near = 0.0 +far = 9.8 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0038_00_full.txt b/configs/scannet/train/scene0038_00_full.txt new file mode 100644 index 0000000..6da77b8 --- /dev/null +++ b/configs/scannet/train/scene0038_00_full.txt @@ -0,0 +1,31 @@ +expname = scene0038_00_full +label_sparse_inv = 1 +basedir = ./logs/scannet + +# datadir = ./data/scannet/scene0038_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0038_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0038_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 9.2 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0038_00_sparse.txt b/configs/scannet/train/scene0038_00_sparse.txt new file mode 100644 index 0000000..61564a4 --- /dev/null +++ b/configs/scannet/train/scene0038_00_sparse.txt @@ -0,0 +1,31 @@ +expname = scene0038_00_sparse10 +label_sparse_inv = 10 +basedir = ./logs/scannet + +# datadir = ./data/scannet/scene0038_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0038_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0038_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 9.2 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0088_00_full.txt b/configs/scannet/train/scene0088_00_full.txt new file mode 100644 index 0000000..9b89fa0 --- /dev/null +++ b/configs/scannet/train/scene0088_00_full.txt @@ -0,0 +1,31 @@ +expname = scene0088_00_sparse10 +label_sparse_inv = 1 +basedir = ./logs/scannet + +# datadir = ./data/scannet/scene0088_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0088_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0088_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 2048 +N_test = 2048 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 9.5 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0088_00_sparse.txt b/configs/scannet/train/scene0088_00_sparse.txt new file mode 100644 index 0000000..d8ab129 --- /dev/null +++ b/configs/scannet/train/scene0088_00_sparse.txt @@ -0,0 +1,31 @@ +expname = scene0088_00_sparse10 +label_sparse_inv = 10 +basedir = ./logs/scannet + +# datadir = ./data/scannet/scene0088_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0088_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0088_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 2048 +N_test = 2048 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 9.5 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0113_00_full.txt b/configs/scannet/train/scene0113_00_full.txt new file mode 100644 index 0000000..480e4be --- /dev/null +++ b/configs/scannet/train/scene0113_00_full.txt @@ -0,0 +1,30 @@ +expname = scene0113_00_full +basedir = ./logs/scannet + +# datadir = ./data/scannet/scene0113_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0113_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0113_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 2048 +N_test = 2048 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 7.6 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0113_00_sparse.txt b/configs/scannet/train/scene0113_00_sparse.txt new file mode 100644 index 0000000..4e74c95 --- /dev/null +++ b/configs/scannet/train/scene0113_00_sparse.txt @@ -0,0 +1,31 @@ +expname = scene0113_00_sparse10 +label_sparse_inv = 10 +basedir = ./logs/scannet + +# datadir = ./data/scannet/scene0113_00 +# audodl-3090 +# datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0113_00 +# 2080ti +datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0113_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 7.6 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0192_00_full.txt b/configs/scannet/train/scene0192_00_full.txt new file mode 100644 index 0000000..dc062d4 --- /dev/null +++ b/configs/scannet/train/scene0192_00_full.txt @@ -0,0 +1,33 @@ +expname = scene0192_00_full +label_sparse_inv = 1 + +basedir = ./logs/scannet + +# datadir = ./data/scannet/scene0192_00 +# audodl-3090 +# datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0192_00 +# 2080ti +datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0192_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 8.2 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + + +# render test step +i_test = -1 \ No newline at end of file diff --git a/configs/scannet/train/scene0192_00_sparse.txt b/configs/scannet/train/scene0192_00_sparse.txt new file mode 100644 index 0000000..8e15886 --- /dev/null +++ b/configs/scannet/train/scene0192_00_sparse.txt @@ -0,0 +1,33 @@ +expname = scene0192_00_sparse10 +label_sparse_inv = 10 + +basedir = ./logs/scannet + +# datadir = ./data/scannet/scene0192_00 +# audodl-3090 +datadir = /root/autodl-tmp/datasets/nerf_scannet/scene0192_00 +# 2080ti +# datadir = /data/dzy_data/nerf/datasets/nerf_scannet/scene0192_00 + +lrate_decay = 500 + +N_samples = 64 +N_importance = 128 +N_train = 1500 +N_test = 1500 + +crop_width = 640 +crop_height = 480 +resize + +near = 0.0 +far = 8.2 +testskip = 1 + +penalize +tolerance = 0.05 +deta_w = 0.05 + + +# render test step +i_test = -1 \ No newline at end of file diff --git a/data/color_dict-old.json b/data/color_dict-old.json new file mode 100644 index 0000000..4b5cb95 --- /dev/null +++ b/data/color_dict-old.json @@ -0,0 +1,661 @@ +{ + "dmsr": { + "restroom": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8 + }, + "dinning": { + "1": 0, + "3": 1, + "4": 2, + "6": 3, + "8": 4, + "9": 5, + "10": 6, + "11": 7, + "12": 8, + "13": 9 + }, + "bedroom": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5 + }, + "reception": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "5": 4, + "6": 5, + "7": 6, + "8": 7, + "9": 8, + "10": 9 + }, + "kitchen": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "5": 4 + }, + "bathroom": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7 + }, + "study": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10, + "11": 11, + "12": 12 + }, + "office": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5 + } + }, + "scannet": { + "scene0010_00": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6 + }, + "scene0012_00": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4 + }, + "scene0024_00": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10 + }, + "scene0033_00": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8 + }, + "scene0038_00": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7 + }, + "scene0088_00": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10, + "11": 11, + "12": 12, + "13": 13, + "14": 14, + "15": 15 + }, + "scene0113_00": { + "0": 0, + "1": 1, + "2": 2, + "3": 3 + }, + "scene0192_00": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4 + } + }, + "replica": { + "office_0": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10, + "11": 11, + "12": 12, + "13": 13, + "14": 14, + "15": 15, + "16": 16, + "17": 17, + "18": 18, + "19": 19, + "20": 20, + "21": 21, + "22": 22, + "23": 23, + "24": 24, + "25": 25, + "26": 26, + "27": 27, + "28": 28, + "29": 29, + "30": 30, + "31": 31, + "32": 32, + "33": 33, + "34": 34, + "35": 35, + "36": 36, + "37": 37, + "38": 38, + "39": 39, + "40": 40, + "42": 41, + "43": 42, + "44": 43, + "45": 44, + "46": 45, + "48": 46, + "49": 47, + "50": 48, + "51": 49, + "54": 50, + "55": 51, + "56": 52, + "57": 53, + "58": 54, + "61": 55, + "63": 56, + "64": 57, + "66": 58 + }, + "office_2": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10, + "11": 11, + "12": 12, + "13": 13, + "14": 14, + "17": 15, + "19": 16, + "23": 17, + "27": 18, + "28": 19, + "29": 20, + "31": 21, + "32": 22, + "34": 23, + "38": 24, + "39": 25, + "40": 26, + "41": 27, + "43": 28, + "44": 29, + "46": 30, + "47": 31, + "49": 32, + "51": 33, + "52": 34, + "54": 35, + "57": 36, + "58": 37, + "59": 38, + "60": 39, + "63": 40, + "65": 41, + "67": 42, + "68": 43, + "69": 44, + "70": 45, + "71": 46, + "72": 47, + "73": 48, + "74": 49, + "75": 50, + "76": 51, + "77": 52, + "78": 53, + "79": 54, + "80": 55, + "81": 56, + "82": 57, + "83": 58, + "84": 59, + "85": 60, + "86": 61, + "87": 62, + "88": 63, + "89": 64, + "90": 65, + "91": 66, + "92": 67, + "93": 68 + }, + "office_3": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "6": 4, + "7": 5, + "8": 6, + "9": 7, + "10": 8, + "11": 9, + "12": 10, + "13": 11, + "14": 12, + "15": 13, + "16": 14, + "17": 15, + "18": 16, + "19": 17, + "21": 18, + "22": 19, + "23": 20, + "25": 21, + "26": 22, + "27": 23, + "29": 24, + "30": 25, + "31": 26, + "32": 27, + "33": 28, + "34": 29, + "36": 30, + "38": 31, + "39": 32, + "42": 33, + "43": 34, + "46": 35, + "49": 36, + "50": 37, + "51": 38, + "54": 39, + "55": 40, + "56": 41, + "58": 42, + "60": 43, + "61": 44, + "65": 45, + "66": 46, + "70": 47, + "72": 48, + "75": 49, + "76": 50, + "78": 51, + "79": 52, + "81": 53, + "82": 54, + "83": 55, + "84": 56, + "85": 57, + "86": 58, + "87": 59, + "88": 60, + "89": 61, + "90": 62, + "91": 63, + "92": 64, + "93": 65, + "94": 66, + "95": 67, + "96": 68, + "97": 69, + "98": 70, + "99": 71, + "100": 72, + "101": 73, + "102": 74, + "103": 75, + "105": 76, + "107": 77, + "108": 78, + "110": 79, + "111": 80, + "112": 81 + }, + "office_4": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10, + "11": 11, + "12": 12, + "13": 13, + "14": 14, + "15": 15, + "16": 16, + "17": 17, + "18": 18, + "19": 19, + "20": 20, + "21": 21, + "22": 22, + "23": 23, + "24": 24, + "26": 25, + "28": 26, + "29": 27, + "30": 28, + "31": 29, + "33": 30, + "34": 31, + "35": 32, + "36": 33, + "39": 34, + "41": 35, + "43": 36, + "47": 37, + "48": 38, + "49": 39, + "51": 40, + "52": 41, + "53": 42, + "54": 43, + "55": 44, + "56": 45, + "57": 46, + "58": 47, + "60": 48, + "61": 49, + "62": 50, + "65": 51, + "66": 52, + "69": 53, + "70": 54, + "71": 55 + }, + "room_0": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10, + "11": 11, + "12": 12, + "13": 13, + "14": 14, + "15": 15, + "16": 16, + "17": 17, + "18": 18, + "19": 19, + "20": 20, + "21": 21, + "22": 22, + "23": 23, + "24": 24, + "25": 25, + "26": 26, + "27": 27, + "28": 28, + "29": 29, + "30": 30, + "31": 31, + "32": 32, + "33": 33, + "34": 34, + "35": 35, + "36": 36, + "37": 37, + "38": 38, + "39": 39, + "40": 40, + "41": 41, + "42": 42, + "43": 43, + "44": 44, + "45": 45, + "46": 46, + "47": 47, + "48": 48, + "49": 49, + "50": 50, + "51": 51, + "52": 52, + "53": 53, + "54": 54, + "55": 55, + "56": 56, + "57": 57, + "58": 58, + "59": 59, + "60": 60, + "61": 61, + "62": 62, + "63": 63, + "64": 64, + "65": 65, + "66": 66, + "67": 67, + "68": 68, + "69": 69, + "70": 70, + "71": 71, + "72": 72, + "73": 73, + "74": 74, + "75": 75, + "77": 76, + "78": 77, + "79": 78, + "80": 79, + "81": 80, + "82": 81, + "83": 82, + "84": 83, + "85": 84, + "86": 85, + "87": 86, + "88": 87, + "89": 88, + "90": 89, + "91": 90, + "92": 91, + "93": 92 + }, + "room_1": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10, + "11": 11, + "12": 12, + "13": 13, + "14": 14, + "15": 15, + "16": 16, + "17": 17, + "18": 18, + "19": 19, + "20": 20, + "21": 21, + "22": 22, + "23": 23, + "24": 24, + "25": 25, + "26": 26, + "27": 27, + "28": 28, + "29": 29, + "30": 30, + "31": 31, + "32": 32, + "33": 33, + "35": 34, + "36": 35, + "37": 36, + "39": 37, + "40": 38, + "41": 39, + "42": 40, + "43": 41, + "44": 42, + "45": 43, + "46": 44, + "48": 45, + "49": 46, + "50": 47, + "51": 48, + "52": 49, + "53": 50, + "54": 51, + "55": 52, + "56": 53 + }, + "room_2": { + "0": 0, + "1": 1, + "2": 2, + "3": 3, + "4": 4, + "5": 5, + "6": 6, + "7": 7, + "8": 8, + "9": 9, + "10": 10, + "11": 11, + "12": 12, + "13": 13, + "14": 14, + "15": 15, + "16": 16, + "17": 17, + "18": 18, + "19": 19, + "20": 20, + "21": 21, + "22": 22, + "23": 23, + "24": 24, + "25": 25, + "26": 26, + "27": 27, + "28": 28, + "29": 29, + "30": 30, + "31": 31, + "32": 32, + "33": 33, + "34": 34, + "35": 35, + "36": 36, + "37": 37, + "38": 38, + "39": 39, + "40": 40, + "41": 41, + "42": 42, + "43": 43, + "44": 44, + "45": 45, + "46": 46, + "47": 47, + "48": 48, + "49": 49, + "50": 50, + "51": 51, + "52": 52, + "53": 53, + "54": 54, + "55": 55, + "56": 56, + "57": 57, + "58": 58, + "59": 59, + "60": 60, + "61": 61 + } + } +} \ No newline at end of file diff --git a/data/color_dict.json b/data/color_dict.json index 4b5cb95..9946c45 100644 --- a/data/color_dict.json +++ b/data/color_dict.json @@ -219,15 +219,16 @@ "49": 47, "50": 48, "51": 49, - "54": 50, - "55": 51, - "56": 52, - "57": 53, - "58": 54, - "61": 55, - "63": 56, - "64": 57, - "66": 58 + "52": 50, + "54": 51, + "55": 52, + "56": 53, + "57": 54, + "58": 55, + "61": 56, + "63": 57, + "64": 58, + "66": 59 }, "office_2": { "0": 0, diff --git a/datasets/loader_scannet.py b/datasets/loader_scannet.py index cef1253..e157986 100644 --- a/datasets/loader_scannet.py +++ b/datasets/loader_scannet.py @@ -5,7 +5,7 @@ import imageio import numpy as np from tools.pose_generator import pose_spherical - +import matplotlib.pyplot as plt def img_i(f): if f.endswith('png'): @@ -88,10 +88,10 @@ def load_rgb(self): i_split = [np.arange(counts[i], counts[i + 1]) for i in range(2)] if not self.resize: - intrinsic_f = os.path.join(self.data_dir, 'intrinsic', 'intrinsic_color.txt') + intrinsic_f = os.path.join(self.data_dir,splits[0] ,'intrinsic', 'intrinsic_color.txt') intrinsic = np.loadtxt(intrinsic_f, delimiter=' ') else: - intrinsic_f = os.path.join(self.data_dir, 'intrinsic', 'intrinsic_depth.txt') + intrinsic_f = os.path.join(self.data_dir,splits[0] , 'intrinsic', 'intrinsic_depth.txt') intrinsic = np.loadtxt(intrinsic_f, delimiter=' ') return all_rgb, all_pose, i_split, intrinsic @@ -103,7 +103,7 @@ def __init__(self, data_dir, testskip=1, resize=True): self.testskip = testskip self.resize = resize - def load_semantic_instance(self): + def load_semantic_instance(self,load_npz = False): splits = ['train', 'test'] all_ins = [] for s in splits: @@ -113,9 +113,14 @@ def load_semantic_instance(self): skip = self.testskip indices = np.loadtxt(os.path.join(self.data_dir, f'{s}_split.txt')).astype(np.int16) file_train = os.path.join(self.data_dir, s) + if load_npz: + ins_fnames = [os.path.join(file_train, f'{s}_ins', f'{index}.npz') for index in indices] + gt_labels = np.array([ins_npz_i(f) for f in ins_fnames]) + else: + ins_fnames = [os.path.join(file_train, f'{s}_ins_full', f'{index}.png') for index in indices] + gt_labels = [imageio.imread(f) for f in ins_fnames] + gt_labels = np.array(gt_labels).astype(np.float32) - ins_fnames = [os.path.join(file_train, f'{s}_ins', f'{index}.npz') for index in indices] - gt_labels = np.array([ins_npz_i(f) for f in ins_fnames]) index = np.arange(0, len(gt_labels), skip) gt_labels = gt_labels[index] if self.resize: @@ -128,11 +133,25 @@ def load_semantic_instance(self): ins_rgbs = f['datasets'][:] f.close() unique_labels = np.unique(gt_labels) - ins_num = len(unique_labels) - 1 - ins_rgbs = ins_rgbs[:ins_num] - gt_labels[gt_labels == -1] = ins_num + if load_npz: + ins_num = len(unique_labels) - 1 + ins_rgbs = ins_rgbs[:ins_num] + gt_labels[gt_labels == -1] = ins_num + else: + # ins_num = len(unique_labels) - 1 + ins_num = unique_labels.shape[0] return gt_labels, ins_rgbs, ins_num + def show(self,item,ins_rgbs): + h, w = item.shape + ra_se_im_t = np.zeros(shape=(h, w, 3)) + unique_labels = np.unique(item) + for index, label in enumerate(unique_labels): + ra_se_im_t[item == label] = ins_rgbs[int(label)] + ra_se_im_t = ra_se_im_t.astype(np.uint8) + plt.imshow(ra_se_im_t) + plt.show() + def selected_pixels(self, full_ins, ins_num, crop_mask): N, H, W = full_ins.shape full_ins = np.reshape(full_ins, [N, -1]) # (N, H*W) diff --git a/mani_replica.py b/mani_replica.py new file mode 100644 index 0000000..c4e50c5 --- /dev/null +++ b/mani_replica.py @@ -0,0 +1,69 @@ +import torch + +from datasets.loader_replica import * +from config import create_nerf, initial +from networks import manipulator +from networks.tester import render_test +from networks.manipulator import manipulator_demo +from tools import pose_generator + + +def test(): + model_coarse.eval() + model_fine.eval() + args.is_train = False + with torch.no_grad(): + + if args.mani_eval: + print('Manipulating', args.mani_mode, '......') + """this operations list can re-design""" + in_images = torch.Tensor(images) + # in_instances = torch.Tensor(instances).type(torch.int8) + in_instances = instances.type(torch.int8) + # in_poses = torch.Tensor(poses) + in_poses = poses + testsavedir = os.path.join(args.basedir, args.expname, args.log_time, + 'mani_eval_{:06d}'.format(iteration), + 'target_label' + str(args.target_label) + ) + + os.makedirs(testsavedir, exist_ok=True) + + mani_center = pose_generator.get_scene_center(in_poses) + mani_pose_save_path = os.path.join(testsavedir, 'transformation_matrix.json') + pose_generator.generate_poses_eval(args,mani_center = mani_center,save_path=mani_pose_save_path) + trans_dicts = pose_generator.load_mani_poses(args,load_path=mani_pose_save_path) + + manipulator.manipulator_eval(position_embedder, view_embedder, model_coarse, model_fine, in_poses, hwk, + trans_dicts=trans_dicts, save_dir=testsavedir, ins_rgbs=ins_colors, args=args, + gt_rgbs=in_images, gt_labels=in_instances) + print('Manipulating Done', testsavedir) + return + + +if __name__ == '__main__': + + args = initial() + # load data + images, poses, hwk, i_split, instances, ins_colors, args.ins_num = load_data(args) + print('Load data from', args.datadir) + + H, W, K = hwk + i_train, i_test = i_split + position_embedder, view_embedder, model_coarse, model_fine, args = create_nerf(args) + + ckpt_path = os.path.join(args.basedir, args.expname, args.log_time, args.test_model) + print('Reloading from', ckpt_path) + ckpt = torch.load(ckpt_path) + iteration = ckpt['iteration'] + # Load model + model_coarse.load_state_dict(ckpt['network_coarse_state_dict']) + model_fine.load_state_dict(ckpt['network_fine_state_dict']) + + # i_test = i_test[::10] + images = torch.Tensor(images[i_test]) + instances = torch.Tensor(instances[i_test]).type(torch.int16) + poses = torch.Tensor(poses[i_test]) + args.perturb = False + + test() diff --git a/networks/evaluator.py b/networks/evaluator.py index a426ac5..68aa7f6 100644 --- a/networks/evaluator.py +++ b/networks/evaluator.py @@ -56,8 +56,13 @@ def reorder(cost_matrix, valid_ins_num): gt_ins = gt_ins.permute([1, 0]) pred_ins = pred_ins[None, :, :] gt_ins = gt_ins[:, None, :] - - cost_ce = torch.mean(-gt_ins * torch.log(pred_ins + 1e-8) - (1 - gt_ins) * torch.log(1 - pred_ins + 1e-8), dim=-1) + try: + pred_ins = pred_ins.cpu() + gt_ins = gt_ins.cpu() + cost_ce = torch.mean(-gt_ins * torch.log(pred_ins + 1e-8) - (1 - gt_ins) * torch.log(1 - pred_ins + 1e-8), dim=-1) + # cost_ce = np.mean(-gt_ins * np.log(pred_ins + 1e-8) - (1 - gt_ins) * np.log(1 - pred_ins + 1e-8), dim=-1) + except: + cost_ce = torch.mean(-gt_ins * torch.log(pred_ins + 1e-8) - (1 - gt_ins) * torch.log(1 - pred_ins + 1e-8), dim=-1) # get soft iou score between prediction and ground truth, don't need do mean operation TP = torch.sum(pred_ins * gt_ins, dim=-1) @@ -122,6 +127,59 @@ def integral_method(prec, rec): return ap_list +def ins_eval_except(pred_ins, gt_ins, gt_ins_num, ins_num, mask=None): + if mask is None: + pred_label = torch.argmax(pred_ins, dim=-1) + valid_pred_labels = torch.unique(pred_label) + else: + pred_label = torch.argmax(pred_ins, dim=-1) + + pred_label[mask == 0] = ins_num # unlabeled index for prediction set as -1 + valid_pred_labels = torch.unique(pred_label) + # valid_pred_labels = torch.unique(pred_label)[:-1] + + valid_pred_num = len(valid_pred_labels) + # prepare confidence masks and confidence scores + pred_conf_mask = np.max(pred_ins.numpy(), axis=-1) + + pred_conf_list = [] + valid_pred_labels = valid_pred_labels.numpy().tolist() + for label in valid_pred_labels: + index = torch.where(pred_label == label) + ssm = pred_conf_mask[index[0], index[1]] + pred_obj_conf = np.median(ssm) + pred_conf_list.append(pred_obj_conf) + pred_conf_scores = torch.from_numpy(np.array(pred_conf_list)) + + # change predicted labels to each signal object masks not existed padding as zero + pred_ins = torch.zeros_like(gt_ins) + pred_ins[..., :valid_pred_num] = F.one_hot(pred_label)[..., valid_pred_labels] + + cost_ce, cost_iou, order_row, order_col = hungarian(pred_ins.reshape((-1, ins_num)), + gt_ins.reshape((-1, ins_num)), + gt_ins_num, ins_num) + + valid_inds = order_col[:gt_ins_num] + ious_metrics = 1 - cost_iou[order_row, valid_inds] + + # prepare confidence values + confidence = torch.zeros(size=[gt_ins_num]) + for i, valid_ind in enumerate(valid_inds): + if valid_ind < valid_pred_num: + confidence[i] = pred_conf_scores[valid_ind] + else: + confidence[i] = 0 + + ap = calculate_ap(ious_metrics, gt_ins_num, confidence=confidence, function_select='integral') + + invalid_mask = valid_inds >= valid_pred_num + valid_inds[invalid_mask] = 0 + valid_pred_labels = torch.from_numpy(np.array(valid_pred_labels)) + return_labels = valid_pred_labels[valid_inds].cpu().numpy() + return_labels[invalid_mask] = -1 + + return pred_label, ap, return_labels,ious_metrics + def ins_eval(pred_ins, gt_ins, gt_ins_num, ins_num, mask=None): if mask is None: pred_label = torch.argmax(pred_ins, dim=-1) @@ -172,4 +230,52 @@ def ins_eval(pred_ins, gt_ins, gt_ins_num, ins_num, mask=None): return_labels = valid_pred_labels[valid_inds].cpu().numpy() return_labels[invalid_mask] = -1 - return pred_label, ap, return_labels + return pred_label, ap, return_labels,ious_metrics + +def get_non_robust_classes(confusion_matrix, robustness_thres): + axis_0 = np.sum(confusion_matrix, axis=0) + axis_1 = np.sum(confusion_matrix, axis=1) + total_labels = axis_0.sum() + non_robust_0 = axis_0 / total_labels < robustness_thres + non_robust_1 = axis_1 / total_labels < robustness_thres + return np.where(np.logical_and(non_robust_0, non_robust_1))[0].tolist() + + +def calculate_miou(confusion_matrix, ignore_class=None, robust=0.005): + MIoU = np.divide(np.diag(confusion_matrix), (np.sum(confusion_matrix, axis=1) + np.sum(confusion_matrix, axis=0) - np.diag(confusion_matrix))) + if ignore_class is not None: + ignore_class += get_non_robust_classes(confusion_matrix, robust) + for i in ignore_class: + MIoU[i] = float('nan') + MIoU = np.nanmean(MIoU) + return MIoU + + +class ConfusionMatrix: + + def __init__(self, num_classes, ignore_class=None, robust=0.005): + np.seterr(divide='ignore', invalid='ignore') + self.num_class = num_classes + self.ignore_class = ignore_class + self.confusion_matrix = np.zeros((self.num_class,) * 2) + self.robust = robust + + def _generate_matrix(self, gt_image, pre_image): + mask = (gt_image >= 0) & (gt_image < self.num_class) + label = self.num_class * gt_image[mask].astype('int') + pre_image[mask] + count = np.bincount(label, minlength=self.num_class**2) + confusion_matrix = count.reshape(self.num_class, self.num_class) + return confusion_matrix + + def add_batch(self, gt_image, pre_image, return_miou=False): + assert gt_image.shape == pre_image.shape + confusion_matrix = self._generate_matrix(gt_image, pre_image) + self.confusion_matrix += confusion_matrix + if return_miou: + return calculate_miou(confusion_matrix, self.ignore_class, self.robust) + + def get_miou(self): + return calculate_miou(self.confusion_matrix, self.ignore_class, self.robust) + + def reset(self): + self.confusion_matrix = np.zeros((self.num_class,) * 2) diff --git a/networks/render.py b/networks/render.py index 0d34afe..b2d5b19 100644 --- a/networks/render.py +++ b/networks/render.py @@ -24,6 +24,7 @@ def render_train(raw, z_vals, rays_d): ins_map = torch.sum(weights_ins[..., None] * ins_labels, -2) # [N_rays, Class_number] ins_map = torch.sigmoid(ins_map) ins_map = ins_map[..., :-1] + # ins_map = ins_map[..., ] return rgb_map, weights, depth_map, ins_map diff --git a/networks/tester.py b/networks/tester.py index a389e07..ab2ae78 100644 --- a/networks/tester.py +++ b/networks/tester.py @@ -9,7 +9,7 @@ from skimage import metrics from networks.render import dm_nerf -from networks.evaluator import to8b, ins_eval +from networks.evaluator import to8b, ins_eval, ConfusionMatrix, ins_eval_except from networks.helpers import get_rays_k, z_val_sample from tools.visualizer import render_label2img, render_gt_label2img @@ -17,11 +17,11 @@ def render_test(position_embedder, view_embedder, model_coarse, model_fine, render_poses, hwk, args, gt_imgs=None, gt_labels=None, ins_rgbs=None, savedir=None, matched_file=None, crop_mask=None): data_info = args.datadir.split('/') - dataset_name = data_info[2] + dataset_name = data_info[-2] scene_name = data_info[-1] H, W, K = hwk cropped_imgs, cropped_labels, psnrs, ssims, lpipses, aps = [], [], [], [], [], [] - + metric_ious = [] if crop_mask is not None: crop_mask = crop_mask.reshape(-1) for index in range(len(gt_imgs)): @@ -49,7 +49,18 @@ def render_test(position_embedder, view_embedder, model_coarse, model_fine, rend import json gt_color_dict_path = './data/color_dict.json' gt_color_dict = json.load(open(gt_color_dict_path, 'r')) - color_dict = gt_color_dict[dataset_name][scene_name] + scannet_full = True + try: + color_dict = gt_color_dict['replica'][scene_name] + scannet_full = False + except: + + color_dict = gt_color_dict['scannet'][scene_name] + if args.ins_num is not None: + if len(color_dict) != args.ins_num + 1: + for i in range(args.ins_num + 1): + color_dict[str(i)] = i + full_map = {} for i, c2w in enumerate(render_poses): @@ -96,13 +107,22 @@ def render_test(position_embedder, view_embedder, model_coarse, model_fine, rend gt_label = gt_labels[i].cpu() if crop_mask is not None: - valid_gt_labels = torch.unique(gt_label)[:-1] + if scannet_full: + valid_gt_labels = torch.unique(gt_label) + else: + valid_gt_labels = torch.unique(gt_label)[:-1] valid_gt_num = len(valid_gt_labels) gt_ins[..., :valid_gt_num] = F.one_hot(gt_label.long())[..., valid_gt_labels.long()] gt_label_np = valid_gt_labels.cpu().numpy() if valid_gt_num > 0: mask = (gt_label < args.ins_num).type(torch.float32) - pred_label, ap, pred_matched_order = ins_eval(ins.cpu(), gt_ins, valid_gt_num, args.ins_num, mask) + try: + pred_label, ap, pred_matched_order,iou_ins_eval = ins_eval(ins.cpu(), gt_ins, valid_gt_num, args.ins_num, mask) + except: + pred_label, ap, pred_matched_order, iou_ins_eval = ins_eval_except(ins.cpu(), gt_ins, valid_gt_num, + args.ins_num, mask) + # pred_label, ap, pred_matched_order, iou_ins_eval = ins_eval_except(ins.cpu(), gt_ins, valid_gt_num, + # args.ins_num, mask) else: pred_label = -1 * torch.ones([H, W]) ap = torch.tensor([1.0]) @@ -112,7 +132,7 @@ def render_test(position_embedder, view_embedder, model_coarse, model_fine, rend gt_ins[..., :valid_gt_num] = F.one_hot(gt_label.long())[..., valid_gt_labels.long()] gt_label_np = valid_gt_labels.cpu().numpy() # change name if valid_gt_num > 0: - pred_label, ap, pred_matched_order = ins_eval(ins.cpu(), gt_ins, valid_gt_num, args.ins_num) + pred_label, ap, pred_matched_order,iou_ins_eval = ins_eval(ins.cpu(), gt_ins, valid_gt_num, args.ins_num) else: pred_label = -1 * torch.ones([H, W]) ap = torch.tensor([1.0]) @@ -127,17 +147,35 @@ def render_test(position_embedder, view_embedder, model_coarse, model_fine, rend aps.append(ap) print(f"APs: {ap}") + print("iou_ins_eval:") + # try: + # train_cm = ConfusionMatrix(num_classes=valid_gt_num) + # metric_iou = train_cm.add_batch(np.array(gt_ins.cpu().numpy(), dtype=np.uint8), + # np.array(ins.cpu().numpy(), dtype=np.uint8), + # return_miou=True) + # except: + # metric_iou = [] + # metric_ious.append(metric_iou) + # print(f"metric_iou: {metric_iou}") + + + if savedir is not None: rgb8 = to8b(rgb.cpu().numpy()) ins_img = render_label2img(pred_label, ins_rgbs, color_dict, ins_map) filename = os.path.join(savedir, '{:03d}.png'.format(i)) fileins = os.path.join(savedir, f"instance_{str(i).zfill(3)}.png") - cv2.imwrite(fileins, ins_img) + # cv2.imwrite(fileins, ins_img) + pred_ins_file = os.path.join(savedir, f'{i}_ins_pred_mask.png') + imageio.imwrite(pred_ins_file, np.array(pred_label.cpu().numpy(), dtype=np.uint8)) + imageio.imwrite(fileins, ins_img) imageio.imwrite(filename, rgb8) gt_ins_img = render_gt_label2img(gt_label, ins_rgbs, color_dict) gt_img_file = os.path.join(savedir, f'{i}_ins_gt.png') - cv2.imwrite(gt_img_file, gt_ins_img) + # cv2.imwrite(gt_img_file, gt_ins_img) + imageio.imwrite(gt_img_file, gt_ins_img) + gt_ins_file = os.path.join(savedir, f'{i}_ins_gt_mask.png') imageio.imwrite(gt_ins_file, np.array(gt_label.cpu().numpy(), dtype=np.uint8)) @@ -147,9 +185,18 @@ def render_test(position_embedder, view_embedder, model_coarse, model_fine, rend json.dump(full_map, f) aps = np.array(aps) + # try: + # metric_ious = np.array(metric_ious) + + # output = np.stack([psnrs, ssims, lpipses, aps[:, 0], aps[:, 1], aps[:, 2], aps[:, 3], aps[:, 4], aps[:, 5], metric_ious[:]]) output = np.stack([psnrs, ssims, lpipses, aps[:, 0], aps[:, 1], aps[:, 2], aps[:, 3], aps[:, 4], aps[:, 5]]) output = output.transpose([1, 0]) out_ap = np.mean(aps, axis=0) + # out_miou = np.mean(metric_ious, axis=0) + # mean_output = np.array([np.mean(psnrs), np.mean(ssims), np.mean(lpipses), + # out_ap[0], out_ap[1], out_ap[2], out_ap[3], out_ap[4], out_ap[5],out_miou]) + # mean_output = mean_output.reshape([1, 10]) + mean_output = np.array([np.mean(psnrs), np.mean(ssims), np.mean(lpipses), out_ap[0], out_ap[1], out_ap[2], out_ap[3], out_ap[4], out_ap[5]]) mean_output = mean_output.reshape([1, 9]) @@ -160,3 +207,202 @@ def render_test(position_embedder, view_embedder, model_coarse, model_fine, rend print('PSNR: {:.4f}, SSIM: {:.4f}, LPIPS: {:.4f} '.format(np.mean(psnrs), np.mean(ssims), np.mean(lpipses))) print('AP50: {:.4f}, AP75: {:.4f}, AP80: {:.4f}, AP85: {:.4f}, AP90: {:.4f}, AP95: {:.4f}' .format(out_ap[0], out_ap[1], out_ap[2], out_ap[3], out_ap[4], out_ap[5])) + # print("miou: {:.4f}",out_miou) + + + +def render_test_scannet(position_embedder, view_embedder, model_coarse, model_fine, render_poses, hwk, args, gt_imgs=None, + gt_labels=None, ins_rgbs=None, savedir=None, matched_file=None, crop_mask=None): + data_info = args.datadir.split('/') + dataset_name = data_info[2] + scene_name = data_info[-1] + H, W, K = hwk + cropped_imgs, cropped_labels, psnrs, ssims, lpipses, aps = [], [], [], [], [], [] + metric_ious = [] + if crop_mask is not None: + crop_mask = crop_mask.reshape(-1) + for index in range(len(gt_imgs)): + gt_img = gt_imgs[index].reshape(-1, 3) + gt_img = gt_img[crop_mask == 1] + gt_img = gt_img.reshape([args.crop_height, args.crop_width, gt_img.shape[-1]]) + cropped_imgs.append(gt_img) + + gt_label = gt_labels[index].reshape(-1) + gt_label = gt_label[crop_mask == 1] + gt_label = gt_label.reshape([args.crop_height, args.crop_width]) + cropped_labels.append(gt_label) + # vis_crop(gt_img) + + gt_imgs = torch.stack(cropped_imgs, 0) + gt_labels = torch.stack(cropped_labels, 0) + gt_imgs_cpu = gt_imgs.cpu().numpy() + gt_imgs_gpu = gt_imgs.to(args.device) + lpips_vgg = lpips.LPIPS(net="vgg").to(args.device) + gt_ins = torch.zeros(size=(gt_labels.shape[1], gt_labels.shape[2], args.ins_num)).cpu() + if matched_file is not None: + if os.path.exists(matched_file): + os.remove(matched_file) + + import json + gt_color_dict_path = './data/color_dict.json' + gt_color_dict = json.load(open(gt_color_dict_path, 'r')) + try: + color_dict = gt_color_dict['replica'][scene_name] + except: + + color_dict = gt_color_dict['scannet'][scene_name] + if args.ins_num is not None: + if len(color_dict) != args.ins_num + 1: + for i in range(args.ins_num + 1): + color_dict[str(i)] = i + + full_map = {} + + for i, c2w in enumerate(render_poses): + print('=' * 50, i, '=' * 50) + t = time.time() + z_val_coarse = z_val_sample(args.N_test, args.near, args.far, args.N_samples) + rays_o, rays_d = get_rays_k(H, W, K, torch.Tensor(c2w)) + rays_o = torch.reshape(rays_o, [-1, 3]).float() + rays_d = torch.reshape(rays_d, [-1, 3]).float() + full_rgb, full_ins, full_disp = None, None, None + for step in range(0, H * W, args.N_test): + N_test = args.N_test + if step + N_test > H * W: + N_test = H * W - step + z_val_coarse = z_val_sample(N_test, args.near, args.far, args.N_samples) + rays_io = rays_o[step:step + N_test] # (chuck, 3) + rays_id = rays_d[step:step + N_test] # (chuck, 3) + batch_rays = torch.stack([rays_io, rays_id], dim=0) + all_info = dm_nerf(batch_rays, position_embedder, view_embedder, + model_coarse, model_fine, z_val_coarse, args) + if full_rgb is None and full_ins is None: + full_rgb, full_ins = all_info['rgb_fine'], all_info['ins_fine'] + else: + full_rgb = torch.cat((full_rgb, all_info['rgb_fine']), dim=0) + full_ins = torch.cat((full_ins, all_info['ins_fine']), dim=0) + if crop_mask is not None: + rgb = full_rgb[crop_mask == 1] + ins = full_ins[crop_mask == 1] + rgb = rgb.reshape([args.crop_height, args.crop_width, rgb.shape[-1]]) + ins = ins.reshape([args.crop_height, args.crop_width, ins.shape[-1]]) + else: + rgb = full_rgb.reshape([H, W, full_rgb.shape[-1]]) + ins = full_ins.reshape([H, W, full_ins.shape[-1]]) + + if gt_imgs is not None: + # rgb image evaluation part + psnr = metrics.peak_signal_noise_ratio(rgb.cpu().numpy(), gt_imgs_cpu[i], data_range=1) + ssim = metrics.structural_similarity(rgb.cpu().numpy(), gt_imgs_cpu[i], multichannel=True, data_range=1) + lpips_i = lpips_vgg(rgb.permute(2, 0, 1).unsqueeze(0), gt_imgs_gpu[i].permute(2, 0, 1).unsqueeze(0)) + psnrs.append(psnr) + ssims.append(ssim) + lpipses.append(lpips_i.item()) + print(f"PSNR: {psnr} SSIM: {ssim} LPIPS: {lpips_i.item()}") + + gt_label = gt_labels[i].cpu() + if crop_mask is not None: + valid_gt_labels = torch.unique(gt_label)[:-1] + valid_gt_num = len(valid_gt_labels) + gt_ins[..., :valid_gt_num] = F.one_hot(gt_label.long())[..., valid_gt_labels.long()] + gt_label_np = valid_gt_labels.cpu().numpy() + if valid_gt_num > 0: + mask = (gt_label < args.ins_num).type(torch.float32) + pred_label = torch.argmax(ins, dim=-1) + # pred_label, ap, pred_matched_order,iou_ins_eval = ins_eval(ins.cpu(), gt_ins, valid_gt_num, args.ins_num, mask) + else: + pred_label = -1 * torch.ones([H, W]) + ap = torch.tensor([1.0]) + else: + valid_gt_labels = torch.unique(gt_label) + valid_gt_num = len(valid_gt_labels) + gt_ins[..., :valid_gt_num] = F.one_hot(gt_label.long())[..., valid_gt_labels.long()] + gt_label_np = valid_gt_labels.cpu().numpy() # change name + if valid_gt_num > 0: + pred_label = torch.argmax(ins, dim=-1) + + # pred_label, ap, pred_matched_order,iou_ins_eval = ins_eval(ins.cpu(), gt_ins, valid_gt_num, args.ins_num) + else: + pred_label = -1 * torch.ones([H, W]) + ap = torch.tensor([1.0]) + # get ins_map + # ins_map = {} + # for idx, pred_label_replica in enumerate(pred_matched_order): + # if pred_label_replica != -1: + # ins_map[str(pred_label_replica)] = int(gt_label_np[idx]) + + # full_map[i] = ins_map + + # aps.append(ap) + # print(f"APs: {ap}") + + print("iou_ins_eval:") + # try: + # train_cm = ConfusionMatrix(num_classes=valid_gt_num) + # metric_iou = train_cm.add_batch(np.array(gt_ins.cpu().numpy(), dtype=np.uint8), + # np.array(ins.cpu().numpy(), dtype=np.uint8), + # return_miou=True) + # except: + # metric_iou = [] + # metric_ious.append(metric_iou) + # print(f"metric_iou: {metric_iou}") + + + + if savedir is not None: + rgb8 = to8b(rgb.cpu().numpy()) + # ins_img = render_label2img(pred_label, ins_rgbs, color_dict, ins_map) + ins_img = render_gt_label2img(pred_label, ins_rgbs, color_dict) + filename = os.path.join(savedir, '{:03d}.png'.format(i)) + fileins = os.path.join(savedir, f"instance_{str(i).zfill(3)}.png") + # cv2.imwrite(fileins, ins_img) + pred_ins_file = os.path.join(savedir, f'{i}_ins_pred_mask.png') + imageio.imwrite(pred_ins_file, np.array(pred_label.cpu().numpy(), dtype=np.uint8)) + imageio.imwrite(fileins, ins_img) + imageio.imwrite(filename, rgb8) + + gt_ins_img = render_gt_label2img(gt_label, ins_rgbs, color_dict) + gt_img_file = os.path.join(savedir, f'{i}_ins_gt.png') + # cv2.imwrite(gt_img_file, gt_ins_img) + imageio.imwrite(gt_img_file, gt_ins_img) + + gt_ins_file = os.path.join(savedir, f'{i}_ins_gt_mask.png') + imageio.imwrite(gt_ins_file, np.array(gt_label.cpu().numpy(), dtype=np.uint8)) + + if gt_imgs is not None: + map_result_file = os.path.join(savedir, 'matching_log.json') + with open(map_result_file, 'w') as f: + json.dump(full_map, f) + + aps = np.array(aps) + # try: + # metric_ious = np.array(metric_ious) + try: + # output = np.stack([psnrs, ssims, lpipses, aps[:, 0], aps[:, 1], aps[:, 2], aps[:, 3], aps[:, 4], aps[:, 5], metric_ious[:]]) + output = np.stack([psnrs, ssims, lpipses, aps[:, 0], aps[:, 1], aps[:, 2], aps[:, 3], aps[:, 4], aps[:, 5]]) + output = output.transpose([1, 0]) + out_ap = np.mean(aps, axis=0) + # out_miou = np.mean(metric_ious, axis=0) + # mean_output = np.array([np.mean(psnrs), np.mean(ssims), np.mean(lpipses), + # out_ap[0], out_ap[1], out_ap[2], out_ap[3], out_ap[4], out_ap[5],out_miou]) + # mean_output = mean_output.reshape([1, 10]) + + mean_output = np.array([np.mean(psnrs), np.mean(ssims), np.mean(lpipses), + out_ap[0], out_ap[1], out_ap[2], out_ap[3], out_ap[4], out_ap[5]]) + mean_output = mean_output.reshape([1, 9]) + except: + output = np.stack([psnrs, ssims, lpipses]) + output = output.transpose([1, 0]) + mean_output = np.array([np.mean(psnrs), np.mean(ssims), np.mean(lpipses)]) + mean_output = mean_output.reshape([1, 3]) + + output = np.concatenate([output, mean_output], 0) + test_result_file = os.path.join(savedir, 'test_results.txt') + np.savetxt(fname=test_result_file, X=output, fmt='%.6f', delimiter=' ') + print('=' * 49, 'Avg', '=' * 49) + print('PSNR: {:.4f}, SSIM: {:.4f}, LPIPS: {:.4f} '.format(np.mean(psnrs), np.mean(ssims), np.mean(lpipses))) + # print('AP50: {:.4f}, AP75: {:.4f}, AP80: {:.4f}, AP85: {:.4f}, AP90: {:.4f}, AP95: {:.4f}' + # .format(out_ap[0], out_ap[1], out_ap[2], out_ap[3], out_ap[4], out_ap[5])) + # print("miou: {:.4f}",out_miou) + + diff --git a/reload_train_scannet.py b/reload_train_scannet.py new file mode 100644 index 0000000..ef095a3 --- /dev/null +++ b/reload_train_scannet.py @@ -0,0 +1,183 @@ +import os +import torch +import numpy as np + +from networks.render import dm_nerf +from networks.tester import render_test, render_test_scannet +from config import initial, create_nerf +from networks.penalizer import ins_penalizer +from datasets.loader_scannet import load_data +from networks.helpers import get_select_crop, z_val_sample +from networks.evaluator import ins_criterion, img2mse, mse2psnr + +import time +from tqdm.auto import tqdm + +np.random.seed(0) +torch.cuda.manual_seed(4) + + +def train(): + start_time = time.time() + + model_fine.train() + model_coarse.train() + N_iters = 400000 + 1 + # N_iters = 1 + 1 + num_img = i_train.shape[0] + + z_val_coarse = z_val_sample(args.N_train, args.near, args.far, args.N_samples) + for i in tqdm(range(iteration, N_iters)): + img_i = i % num_img + # img_i = np.random.choice(i_train) + gt_rgb = images[img_i].to(args.device) + pose = poses[img_i, :3, :4].to(args.device) + gt_label = gt_labels[img_i].to(args.device) + ins_index = ins_indices[img_i] + + target_c, target_i, batch_rays, args.N_ins = get_select_crop(gt_rgb, pose, K, gt_label, ins_index, crop_mask, + args.N_train) + + all_info = dm_nerf(batch_rays, position_embedder, view_embedder, model_coarse, model_fine, z_val_coarse, args) + + # coarse losses + rgb_loss_coarse = img2mse(all_info['rgb_coarse'], target_c) + psnr_coarse = mse2psnr(rgb_loss_coarse) + + ins_loss_coarse, valid_ce_coarse, invalid_ce_coarse, valid_siou_coarse = \ + ins_criterion(all_info['ins_coarse'], target_i, args.ins_num) + + # fine losses + rgb_loss_fine = img2mse(all_info['rgb_fine'], target_c) + psnr_fine = mse2psnr(rgb_loss_fine) + ins_loss_fine, valid_ce_fine, invalid_ce_fine, valid_siou_fine = \ + ins_criterion(all_info['ins_fine'], target_i, args.ins_num) + + # without penalize loss + ins_loss = ins_loss_fine + ins_loss_coarse + rgb_loss = rgb_loss_fine + rgb_loss_coarse + if img_i in i_train_label_sparse: + total_loss = ins_loss + rgb_loss + else: + total_loss = 0. + rgb_loss + + # use penalize + if args.penalize: + emptiness_coarse = ins_penalizer(all_info['raw_coarse'], all_info['z_vals_coarse'], + all_info['depth_coarse'], batch_rays[1], args) + emptiness_fine = ins_penalizer(all_info['raw_fine'], all_info['z_vals_fine'], + all_info['depth_fine'], batch_rays[1], args) + + emptiness_loss = emptiness_fine + emptiness_coarse + total_loss = total_loss + emptiness_loss + + # optimizing + optimizer.zero_grad() + total_loss.backward() + optimizer.step() + + # losses decay + ### update learning rate ### + decay_rate = 0.1 + decay_steps = args.lrate_decay * 1000 + new_lrate = args.lrate * (decay_rate ** ((i) / decay_steps)) + for param_group in optimizer.param_groups: + param_group['lr'] = new_lrate + ################################### + + if i % args.i_print == 0: + print(f"[TRAIN] Iter: {i} PSNR: {psnr_fine.item()} Total_Loss: {total_loss.item()} RGB_Loss: {rgb_loss.item()} Ins_Loss: {ins_loss.item()}") + + if i % args.i_save == 0: + path = os.path.join(args.basedir, args.expname, args.log_time, '{:06d}.tar'.format(i)) + save_model = { + 'iteration': i, + 'network_coarse_state_dict': model_coarse.state_dict(), + 'network_fine_state_dict': model_fine.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + } + torch.save(save_model, path) + + # if i % args.i_test == 0 and args.i_test != -1: + if i == 0: + # if i % args.i_test == 0: + model_coarse.eval() + model_fine.eval() + args.is_train = False + selected_indices = np.random.choice(len(i_test), size=[10], replace=False) + selected_i_test = i_test[selected_indices] + testsavedir = os.path.join(args.basedir, args.expname, args.log_time, 'testset_{:06d}'.format(i)) + matched_file = os.path.join(testsavedir, 'matching_log.txt') + os.makedirs(testsavedir, exist_ok=True) + with torch.no_grad(): + test_poses = torch.Tensor(poses[selected_i_test].to(args.device)) + test_imgs = images[selected_i_test] + test_gt_labels = gt_labels[selected_i_test].to(args.device) + render_test(position_embedder, view_embedder, model_coarse, model_fine, test_poses, hwk, args, + gt_imgs=test_imgs, gt_labels=test_gt_labels, ins_rgbs=ins_rgbs, savedir=testsavedir, + matched_file=matched_file, crop_mask=crop_mask) + # render_test() + # render_test_scannet + print('Training model saved!') + args.is_train = True + model_coarse.train() + model_fine.train() + end_time = time.time() + print(f"training time : {end_time - start_time}") + if args.i_test == -1: + model_coarse.eval() + model_fine.eval() + args.is_train = False + # selected_indices = np.random.choice(len(i_test), size=[10], replace=False) + # selected_i_test = i_test[selected_indices] + selected_i_test = i_test + testsavedir = os.path.join(args.basedir, args.expname, args.log_time, 'testset_{:06d}'.format(i)) + matched_file = os.path.join(testsavedir, 'matching_log.txt') + os.makedirs(testsavedir, exist_ok=True) + with torch.no_grad(): + test_poses = torch.Tensor(poses[selected_i_test].to(args.device)) + test_imgs = images[selected_i_test] + test_gt_labels = gt_labels[selected_i_test].to(args.device) + render_test_scannet(position_embedder, view_embedder, model_coarse, model_fine, test_poses, hwk, args, + gt_imgs=test_imgs, gt_labels=test_gt_labels, ins_rgbs=ins_rgbs, savedir=testsavedir, + matched_file=matched_file, crop_mask=crop_mask) + # render_test + print('Training model saved!') + args.is_train = True + model_coarse.train() + model_fine.train() + +if __name__ == '__main__': + args = initial() + + # load data + images, poses, hwk, i_split, gt_labels, ins_rgbs, args.ins_num, ins_indices, crop_mask = load_data(args) + print('Load data from', args.datadir) + + i_train, i_test = i_split + H, W, K = hwk + + # Create nerf model + position_embedder, view_embedder, model_coarse, model_fine, args = create_nerf(args) + + # Create optimizer + grad_vars = list(model_coarse.parameters()) + list(model_fine.parameters()) + optimizer = torch.optim.Adam(params=grad_vars, lr=args.lrate, betas=(0.9, 0.999)) + + # move data to gpu + images = torch.Tensor(images).cpu() + gt_labels = torch.Tensor(gt_labels).type(torch.int16).cpu() + poses = torch.Tensor(poses).cpu() + + sparse_inv = args.label_sparse_inv + i_train_label_sparse = i_train[::sparse_inv] + print("sparse label num:",i_train_label_sparse.shape[0]) + + ckpt_path = os.path.join(args.basedir, args.expname, args.log_time, args.test_model) + ckpt = torch.load(ckpt_path) + iteration = ckpt['iteration'] + # Load model + model_coarse.load_state_dict(ckpt['network_coarse_state_dict']) + model_fine.load_state_dict(ckpt['network_fine_state_dict']) + + train() diff --git a/requirements.txt b/requirements.txt index 8a3b215..964dd86 100644 --- a/requirements.txt +++ b/requirements.txt @@ -71,7 +71,7 @@ oauthlib==3.1.1 olefile==0.46 open3d==0.13.0 opencv-contrib-python==4.5.3.56 -OpenEXR==1.3.2 +# OpenEXR==1.3.2 packaging==21.0 pandas==1.3.3 pandocfilters==1.5.0 @@ -113,9 +113,9 @@ testpath==0.5.0 threadpoolctl==3.0.0 tifffile==2021.10.12 tomlkit==0.7.2 -torch==1.8.1+cu111 -torchaudio==0.8.1 -torchvision==0.9.1+cu111 +# torch==1.8.1+cu111 +# torchaudio==0.8.1 +# torchvision==0.9.1+cu111 tornado==6.1 tqdm==4.62.3 traitlets==5.1.0 diff --git a/retrain_scannet_sparse.sh b/retrain_scannet_sparse.sh new file mode 100644 index 0000000..5d735b6 --- /dev/null +++ b/retrain_scannet_sparse.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +GPU=$1 +scene=88 +# 010 113 38 +sparse_inv=1 +configs=./configs/scannet/train/scene0${scene}_00_sparse.txt +NOHUP_FILE=../log/dm-nerf/scannet/scene0${scene}_00_sparse${sparse_inv}-GPU_${GPU}.out +export CUDA_VISIBLE_DEVICES=$1 + +nohup python3 reload_train_scannet.py --config configs/scannet/train/scene0088_00_sparse.txt --log_time 202304090218 --label_sparse_inv 1 --test_model 200000.tar >$NOHUP_FILE 2>&1 & \ No newline at end of file diff --git a/test_replica.py b/test_replica.py index b94759f..ee5faac 100644 --- a/test_replica.py +++ b/test_replica.py @@ -13,7 +13,12 @@ def test(): testsavedir = os.path.join(args.basedir, args.expname, args.log_time, 'render_test_{:06d}'.format(iteration)) os.makedirs(testsavedir, exist_ok=True) + # seleted = [178,179] mathed_file = os.path.join(testsavedir, 'matching_log.txt') + + # render_test(position_embedder, view_embedder, model_coarse, model_fine, poses[seleted], hwk, args, + # gt_imgs=images[seleted], gt_labels=instances[seleted], ins_rgbs=ins_colors, savedir=testsavedir, + # matched_file=mathed_file) render_test(position_embedder, view_embedder, model_coarse, model_fine, poses, hwk, args, gt_imgs=images, gt_labels=instances, ins_rgbs=ins_colors, savedir=testsavedir, matched_file=mathed_file) diff --git a/test_replica_office_sparse.sh b/test_replica_office_sparse.sh new file mode 100644 index 0000000..7b4f845 --- /dev/null +++ b/test_replica_office_sparse.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +GPU=$1 +scene=$2 +sparse_inv=$3 +configs=./configs/replica/train/office_${scene}_sparse_inv.txt +NOHUP_FILE=../log/dm-nerf/office/office_${scene}_sparse_inv-GPU_${GPU}.out +export CUDA_VISIBLE_DEVICES=$1 +nohup python3 test_replica.py --config ${configs} --label_sparse_inv ${sparse_inv} >$NOHUP_FILE 2>&1 & diff --git a/test_scannet.py b/test_scannet.py index 43a0228..7607dcd 100644 --- a/test_scannet.py +++ b/test_scannet.py @@ -1,7 +1,7 @@ import torch from datasets.loader_scannet import * -from networks.tester import render_test +from networks.tester import render_test, render_test_scannet from config import create_nerf, initial from networks.manipulator import manipulator_demo @@ -16,6 +16,8 @@ def test(): 'render_test_{:06d}'.format(iteration)) os.makedirs(testsavedir, exist_ok=True) mathed_file = os.path.join(testsavedir, 'matching_log.txt') + # render_test + # render_test_scannet render_test(position_embedder, view_embedder, model_coarse, model_fine, poses, hwk, args, gt_imgs=images, gt_labels=instances, ins_rgbs=ins_colors, savedir=testsavedir, matched_file=mathed_file, crop_mask=crop_mask) diff --git a/test_scannet.sh b/test_scannet.sh new file mode 100644 index 0000000..158a929 --- /dev/null +++ b/test_scannet.sh @@ -0,0 +1,10 @@ +#!/bin/bash + +GPU=$1 +sparse=$2 +# 010 113 38 + +NOHUP_FILE=../log/dm-nerf/scannet/test_0038_sparse${sparse}-GPU_${GPU}.out +export CUDA_VISIBLE_DEVICES=$1 +nohup python3 test_scannet.py --config configs/scannet/train/scene0038_00_sparse.txt --log_time 202304080154 --label_sparse_inv ${sparse} --test_model 300000.tar >$NOHUP_FILE 2>&1 & +#nohup python3 test_scannet.py --config configs/scannet/train/scene0038_00_sparse.txt --log_time 202304080154 --label_sparse_inv 20 --test_model 300000.tar >$NOHUP_FILE 2>&1 & diff --git a/tools/visualizer.py b/tools/visualizer.py index bb540d2..7946f6b 100644 --- a/tools/visualizer.py +++ b/tools/visualizer.py @@ -1,6 +1,6 @@ import torch import numpy as np -import open3d as o3d +# import open3d as o3d import matplotlib.pyplot as plt diff --git a/train_replica.py b/train_replica.py index 638cd6b..8672419 100644 --- a/train_replica.py +++ b/train_replica.py @@ -9,20 +9,26 @@ from datasets.loader_replica import load_data from networks.helpers import get_select_full, z_val_sample from networks.evaluator import ins_criterion, img2mse, mse2psnr +import time +from tqdm.auto import tqdm np.random.seed(0) torch.cuda.manual_seed(3) def train(): + start_time = time.time() + model_fine.train() model_coarse.train() - N_iters = 500000 + 1 - + N_iters = 200000 + 1 + # N_iters = 500000 + 1 + num_img = i_train.shape[0] z_val_coarse = z_val_sample(args.N_train, args.near, args.far, args.N_samples) args.N_ins = None - for i in range(0, N_iters): - img_i = np.random.choice(i_train) + for i in tqdm(range(0, N_iters)): + img_i = i % num_img + # img_i = np.random.choice(i_train) gt_rgb = images[img_i].to(args.device) pose = poses[img_i, :3, :4].to(args.device) gt_label = gt_labels[img_i].to(args.device) @@ -47,7 +53,10 @@ def train(): # without penalize loss ins_loss = ins_loss_fine + ins_loss_coarse rgb_loss = rgb_loss_fine + rgb_loss_coarse - total_loss = ins_loss + rgb_loss + if img_i in i_train_label_sparse: + total_loss = ins_loss + rgb_loss + else: + total_loss = 0. + rgb_loss # use penalize if args.penalize: @@ -87,7 +96,9 @@ def train(): } torch.save(save_model, path) - if i % args.i_test == 0: + # if i % args.i_test == 0: + # if i == 0 : + if i % args.i_test == 0 and args.i_test != -1: model_coarse.eval() model_fine.eval() args.is_train = False @@ -107,6 +118,26 @@ def train(): args.is_train = True model_coarse.train() model_fine.train() + end_time = time.time() + print(f"training time : {end_time - start_time}") + if args.i_test == -1: + model_coarse.eval() + model_fine.eval() + args.is_train = False + # selected_indices = np.random.choice(len(i_test), size=[10], replace=False) + # selected_i_test = i_test[selected_indices] + selected_i_test = i_test + testsavedir = os.path.join(args.basedir, args.expname, args.log_time, 'testset_{:06d}'.format(i)) + matched_file = os.path.join(testsavedir, 'matching_log.txt') + os.makedirs(testsavedir, exist_ok=True) + with torch.no_grad(): + test_poses = torch.Tensor(poses[selected_i_test].to(args.device)) + test_imgs = images[selected_i_test] + test_gt_labels = gt_labels[selected_i_test].to(args.device) + render_test(position_embedder, view_embedder, model_coarse, model_fine, test_poses, hwk, args, + gt_imgs=test_imgs, gt_labels=test_gt_labels, ins_rgbs=ins_rgbs, savedir=testsavedir, + matched_file=matched_file) + print('Training model saved!') if __name__ == '__main__': @@ -131,4 +162,8 @@ def train(): gt_labels = torch.Tensor(gt_labels).type(torch.int16).cpu() poses = torch.Tensor(poses).cpu() + sparse_inv = args.label_sparse_inv + i_train_label_sparse = i_train[::sparse_inv] + print("sparse label num:",i_train_label_sparse.shape[0]) + train() diff --git a/train_replica.sh b/train_replica.sh new file mode 100644 index 0000000..d8c7e1d --- /dev/null +++ b/train_replica.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +GPU=$1 +scene=$2 +configs=./configs/replica/train/office_${scene}_sparse_inv.txt +NOHUP_FILE=../log/dm-nerf/office_${scene}_sparse_inv-GPU_${GPU}.out +CUDA_VISIBLE_DEVICES=$1 python3 train_replica.py --config ${configs} >$NOHUP_FILE 2>&1 & diff --git a/train_replica_office_full.sh b/train_replica_office_full.sh new file mode 100644 index 0000000..2635429 --- /dev/null +++ b/train_replica_office_full.sh @@ -0,0 +1,10 @@ +#!/bin/bash + +GPU=$1 +scene=$2 +configs=./configs/replica/train/office_${scene}_full.txt +NOHUP_FILE=../log/dm-nerf/office/office_${scene}_full-GPU_${GPU}.out +export CUDA_VISIBLE_DEVICES=$1 +nohup python3 train_replica.py --config ${configs} >$NOHUP_FILE 2>&1 & + + diff --git a/train_replica_office_sparse.sh b/train_replica_office_sparse.sh new file mode 100644 index 0000000..6afbc34 --- /dev/null +++ b/train_replica_office_sparse.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +GPU=$1 +scene=$2 +sparse_inv=$3 +configs=./configs/replica/train/office_${scene}_sparse_inv.txt +NOHUP_FILE=../log/dm-nerf/office/office_${scene}_sparse_inv-GPU_${GPU}.out +export CUDA_VISIBLE_DEVICES=$1 +nohup python3 train_replica.py --config ${configs} --label_sparse_inv ${sparse_inv} >$NOHUP_FILE 2>&1 & diff --git a/train_replica_room_full.sh b/train_replica_room_full.sh new file mode 100644 index 0000000..39d00dd --- /dev/null +++ b/train_replica_room_full.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +GPU=$1 +scene=$2 +configs=./configs/replica/train/room_${scene}_full.txt +NOHUP_FILE=../log/dm-nerf/room/room_${scene}_full-GPU_${GPU}.out +export CUDA_VISIBLE_DEVICES=$1 +nohup python3 train_replica.py --config ${configs} >$NOHUP_FILE 2>&1 & diff --git a/train_replica_room_sparse.sh b/train_replica_room_sparse.sh new file mode 100644 index 0000000..febf3c0 --- /dev/null +++ b/train_replica_room_sparse.sh @@ -0,0 +1,10 @@ +#!/bin/bash + +GPU=$1 +scene=$2 +sparse_inv=$3 +configs=./configs/replica/train/room_${scene}_sparse_inv.txt +NOHUP_FILE=../log/dm-nerf/room/room_${scene}_sparse_inv-GPU_${GPU}.out +export CUDA_VISIBLE_DEVICES=$1 +nohup python3 train_replica.py --config ${configs} --label_sparse_inv ${sparse_inv} >$NOHUP_FILE 2>&1 & + diff --git a/train_scannet.py b/train_scannet.py index 47f40a6..c5d6214 100644 --- a/train_scannet.py +++ b/train_scannet.py @@ -3,26 +3,33 @@ import numpy as np from networks.render import dm_nerf -from networks.tester import render_test +from networks.tester import render_test, render_test_scannet from config import initial, create_nerf from networks.penalizer import ins_penalizer from datasets.loader_scannet import load_data from networks.helpers import get_select_crop, z_val_sample from networks.evaluator import ins_criterion, img2mse, mse2psnr +import time +from tqdm.auto import tqdm np.random.seed(0) torch.cuda.manual_seed(4) def train(): + start_time = time.time() + model_fine.train() model_coarse.train() - N_iters = 500000 + 1 + N_iters = 300000 + 1 + # N_iters = 1 + 1 + num_img = i_train.shape[0] z_val_coarse = z_val_sample(args.N_train, args.near, args.far, args.N_samples) - for i in range(0, N_iters): - img_i = np.random.choice(i_train) + for i in tqdm(range(0, N_iters)): + img_i = i % num_img + # img_i = np.random.choice(i_train) gt_rgb = images[img_i].to(args.device) pose = poses[img_i, :3, :4].to(args.device) gt_label = gt_labels[img_i].to(args.device) @@ -49,7 +56,10 @@ def train(): # without penalize loss ins_loss = ins_loss_fine + ins_loss_coarse rgb_loss = rgb_loss_fine + rgb_loss_coarse - total_loss = ins_loss + rgb_loss + if img_i in i_train_label_sparse: + total_loss = ins_loss + rgb_loss + else: + total_loss = 0. + rgb_loss # use penalize if args.penalize: @@ -88,7 +98,9 @@ def train(): } torch.save(save_model, path) - if i % args.i_test == 0: + # if i % args.i_test == 0 and args.i_test != -1: + if i == 0: + # if i % args.i_test == 0: model_coarse.eval() model_fine.eval() args.is_train = False @@ -104,11 +116,36 @@ def train(): render_test(position_embedder, view_embedder, model_coarse, model_fine, test_poses, hwk, args, gt_imgs=test_imgs, gt_labels=test_gt_labels, ins_rgbs=ins_rgbs, savedir=testsavedir, matched_file=matched_file, crop_mask=crop_mask) - print('Training model saved!') + # render_test() + # render_test_scannet + print('Training model saved!') args.is_train = True model_coarse.train() model_fine.train() - + end_time = time.time() + print(f"training time : {end_time - start_time}") + if args.i_test == -1: + model_coarse.eval() + model_fine.eval() + args.is_train = False + # selected_indices = np.random.choice(len(i_test), size=[10], replace=False) + # selected_i_test = i_test[selected_indices] + selected_i_test = i_test + testsavedir = os.path.join(args.basedir, args.expname, args.log_time, 'testset_{:06d}'.format(i)) + matched_file = os.path.join(testsavedir, 'matching_log.txt') + os.makedirs(testsavedir, exist_ok=True) + with torch.no_grad(): + test_poses = torch.Tensor(poses[selected_i_test].to(args.device)) + test_imgs = images[selected_i_test] + test_gt_labels = gt_labels[selected_i_test].to(args.device) + render_test_scannet(position_embedder, view_embedder, model_coarse, model_fine, test_poses, hwk, args, + gt_imgs=test_imgs, gt_labels=test_gt_labels, ins_rgbs=ins_rgbs, savedir=testsavedir, + matched_file=matched_file, crop_mask=crop_mask) + # render_test + print('Training model saved!') + args.is_train = True + model_coarse.train() + model_fine.train() if __name__ == '__main__': args = initial() @@ -132,4 +169,8 @@ def train(): gt_labels = torch.Tensor(gt_labels).type(torch.int16).cpu() poses = torch.Tensor(poses).cpu() + sparse_inv = args.label_sparse_inv + i_train_label_sparse = i_train[::sparse_inv] + print("sparse label num:",i_train_label_sparse.shape[0]) + train() diff --git a/train_scannet_full.sh b/train_scannet_full.sh new file mode 100644 index 0000000..8e9b58e --- /dev/null +++ b/train_scannet_full.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +GPU=$1 +scene=$2 +# 010 113 38 + +configs=./configs/scannet/train/scene0${scene}_00_full.txt +NOHUP_FILE=../log/dm-nerf/scannet/scene0${scene}_00_full-GPU_${GPU}.out +export CUDA_VISIBLE_DEVICES=$1 +nohup python3 train_scannet.py --config ${configs} >$NOHUP_FILE 2>&1 & + diff --git a/train_scannet_sparse.sh b/train_scannet_sparse.sh new file mode 100644 index 0000000..87a8062 --- /dev/null +++ b/train_scannet_sparse.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +GPU=$1 +scene=$2 +# 010 113 38 +sparse_inv=$3 +configs=./configs/scannet/train/scene0${scene}_00_sparse.txt +NOHUP_FILE=../log/dm-nerf/scannet/scene0${scene}_00_sparse${sparse_inv}-GPU_${GPU}.out +export CUDA_VISIBLE_DEVICES=$1 +nohup python3 train_scannet.py --config ${configs} --label_sparse_inv ${sparse_inv} >$NOHUP_FILE 2>&1 & +