Rename train.py to train_dtu.py

FangjinhuaWang · Sep 15, 2021 · 19a081d · 19a081d
1 parent c167438
commit 19a081d
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 19 deletions.
diff --git a/datasets/__init__.py b/datasets/__init__.py
@@ -1,8 +0,0 @@
-import importlib
-
-
-# find the dataset definition by name, for example dtu_yao (dtu_yao.py)
-def find_dataset_def(dataset_name):
-    module_name = 'datasets.{}'.format(dataset_name)
-    module = importlib.import_module(module_name)
-    return getattr(module, "MVSDataset")

diff --git a/train.sh b/train.sh
@@ -2,10 +2,5 @@
 
 # train on DTU's training set
 MVS_TRAINING="/home/mvs_training/dtu/"
-
-python train.py --dataset dtu_yao --batch_size 4 --epochs 8 \
---patchmatch_iteration 1 2 2 --patchmatch_range 6 4 2 \
---patchmatch_num_sample 8 8 16 --propagate_neighbors 0 8 16 --evaluate_neighbors 9 9 9 \
---patchmatch_interval_scale 0.005 0.0125 0.025 \
---trainpath=$MVS_TRAINING --trainlist lists/dtu/train.txt --vallist lists/dtu/val.txt \
---logdir ./checkpoints $@
+python train_dtu.py --batch_size 4 --epochs 8 --trainpath=$MVS_TRAINING --trainlist lists/dtu/train.txt \
+--vallist lists/dtu/val.txt --logdir ./checkpoints "$@"
diff --git a/train.py → train_dtu.py b/train.py → train_dtu.py
@@ -9,7 +9,7 @@
 import time
 from typing import List
 from torch.utils.tensorboard import SummaryWriter
-from datasets import find_dataset_def
+from datasets.dtu_yao import MVSDataset
 from models import *
 from utils import *
 import sys
@@ -21,7 +21,6 @@
 parser.add_argument('--mode', default='train', help='train or val', choices=['train', 'val'])
 parser.add_argument('--model', default='PatchmatchNet', help='select model')
 
-parser.add_argument('--dataset', default='dtu_yao', help='select dataset')
 parser.add_argument('--trainpath', help='train datapath')
 parser.add_argument('--valpath', help='validation datapath')
 parser.add_argument('--trainlist', help='train list')
@@ -84,7 +83,6 @@
 print_args(args)
 
 # dataset, dataloader
-MVSDataset = find_dataset_def(args.dataset)
 train_dataset = MVSDataset(args.trainpath, args.trainlist, "train", 5, robust_train=True)
 test_dataset = MVSDataset(args.valpath, args.vallist, "val", 5,  robust_train=False)
 
@@ -326,4 +324,3 @@ def create_stage_images(image: torch.Tensor) -> List[torch.Tensor]:
         train()
     elif args.mode == "val":
         test()
-