utils: Edit config dict

1. Add dataset_root, model_name.
2. Change interface of function

Author: synml

demo.py

@@ -11,15 +11,15 @@
 
 if __name__ == '__main__':
     # 0. Load config
-    model_name, config = utils.utils.load_config()
-    print('Activated model: {}'.format(model_name))
+    config = utils.utils.load_config()
+    print('Activated model: {}'.format(config['model_name']))
 
     # 1. Dataset
     _, _, testset, testloader = utils.utils.init_cityscapes_dataset(config)
 
     # 2. Model
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = utils.utils.get_model(model_name, config['num_classes'], config['pretrained_weights']).to(device)
+    model = utils.utils.get_model(config['model_name'], config['num_classes'], config['pretrained_weights']).to(device)
 
     # 이미지 이름 불러오기
     image_names = []
@@ -32,7 +32,7 @@
 
     # 예측 결과 저장
     step = 0
-    result_dir = os.path.join('demo', model_name.lower())
+    result_dir = os.path.join('demo', config['model_name'].lower())
     groundtruth_dir = os.path.join('demo', 'groundtruth')
     os.makedirs(result_dir, exist_ok=True)
     os.makedirs(groundtruth_dir, exist_ok=True)

eval.py

@@ -87,22 +87,22 @@ def evaluate(model, testloader, num_classes: int, device):
 
 if __name__ == '__main__':
     # 0. Load config
-    model_name, config = utils.utils.load_config()
-    print('Activated model: {}'.format(model_name))
+    config = utils.utils.load_config()
+    print('Activated model: {}'.format(config['model_name']))
 
     # 1. Dataset
     _, _, testset, testloader = utils.utils.init_cityscapes_dataset(config)
 
     # 2. Model
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = utils.utils.get_model(model_name, config['num_classes'], config['pretrained_weights']).to(device)
+    model = utils.utils.get_model(config['model_name'], config['num_classes'], config['pretrained_weights']).to(device)
 
     # 모델 평가
     val_loss, iou, miou, fps = evaluate(model, testloader, config['num_classes'], device)
 
     # 평가 결과를 csv 파일로 저장
     os.makedirs('result', exist_ok=True)
-    with open(os.path.join('result', '{}.csv'.format(model_name)), mode='w') as f:
+    with open(os.path.join('result', '{}.csv'.format(config['model_name'])), mode='w') as f:
         writer = csv.writer(f, delimiter=',', lineterminator='\n')
 
         writer.writerow(['Class Number', 'Class Name', 'IoU'])

models/models.ini

@@ -1,3 +1,6 @@
+[dataset]
+root = ../../data/cityscapes
+
 [model]
 activate_model = UNet
 

train.py

@@ -11,23 +11,23 @@
 
 if __name__ == '__main__':
     # 0. Load config
-    model_name, config = utils.utils.load_config()
-    print('Activated model: {}'.format(model_name))
+    config = utils.utils.load_config()
+    print('Activated model: {}'.format(config['model_name']))
 
     # 1. Dataset
     trainset, trainloader, testset, testloader = utils.utils.init_cityscapes_dataset(config)
 
     # 2. Model
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = utils.utils.get_model(model_name, config['num_classes']).to(device)
+    model = utils.utils.get_model(config['model_name'], config['num_classes']).to(device)
 
     # 3. Loss function, optimizer, lr scheduler
     criterion = nn.CrossEntropyLoss()
     optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])
     scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, min_lr=0.0001)
 
     # 4. Tensorboard
-    writer = torch.utils.tensorboard.SummaryWriter(os.path.join('runs', model_name))
+    writer = torch.utils.tensorboard.SummaryWriter(os.path.join('runs', config['model_name']))
     writer.add_graph(model, trainloader.__iter__().__next__()[0].to(device))
 
     # 5. Train and evaluate
@@ -69,6 +69,6 @@
         if miou > prev_miou:
             os.makedirs('weights', exist_ok=True)
             torch.save(model.state_dict(),
-                       os.path.join('weights', '{}_best.pth'.format(model_name)))
+                       os.path.join('weights', '{}_best.pth'.format(config['model_name'])))
             prev_miou = miou
     writer.close()

utils/utils.py

@@ -20,6 +20,8 @@ def load_config():
 
     model_name = parser['model']['activate_model']
     config = {
+        'dataset_root': parser['dataset']['root'],
+        'model_name': model_name,
         'batch_size': parser.getint(model_name, 'batch_size'),
         'epoch': parser.getint(model_name, 'epoch'),
         'image_size': (int(parser[model_name]['image_size'].split('x')[1]),
@@ -29,7 +31,7 @@ def load_config():
         'num_workers': parser.getint(model_name, 'num_workers'),
         'pretrained_weights': parser[model_name]['pretrained_weights'],
     }
-    return model_name, config
+    return config
 
 
 # 모델 불러오기
@@ -98,7 +100,7 @@ def init_cityscapes_dataset(config: dict):
         torchvision.transforms.Resize(config['image_size'], interpolation=0),
         torchvision.transforms.ToTensor(),
     ])
-    trainset = utils.datasets.Cityscapes(root='../../data/cityscapes',
+    trainset = utils.datasets.Cityscapes(root=config['dataset_root'],
                                          split='train',
                                          mode='fine',
                                          target_type='semantic',
@@ -109,7 +111,7 @@ def init_cityscapes_dataset(config: dict):
                                               shuffle=True,
                                               num_workers=config['num_workers'],
                                               pin_memory=True)
-    testset = utils.datasets.Cityscapes(root='../../data/cityscapes',
+    testset = utils.datasets.Cityscapes(root=config['dataset_root'],
                                         split='val',
                                         mode='fine',
                                         target_type='semantic',