update

config/__init__.py

@@ -1,3 +1,19 @@
+# ------------------------ Dataset Config ------------------------
+from .dataset_config import dataset_cfg
+
+
+def build_dataset_config(args):
+    if args.dataset in ['coco', 'coco-val', 'coco-test']:
+        cfg = dataset_cfg['coco']
+    else:
+        cfg = dataset_cfg[args.dataset]
+
+    print('==============================')
+    print('Dataset Config: {} \n'.format(cfg))
+
+    return cfg
+
+
 # ------------------ Model Config ----------------------
 from .yolov1_config import yolov1_cfg
 from .yolov2_config import yolov2_cfg

config/dataset_config.py

@@ -0,0 +1,47 @@
+# Dataset config
+import os
+
+dataset_cfg = {
+    'voc': {
+        'data_name': 'VOCdevkit',
+        'num_classes': 20,
+        'class_indexs': None,
+        'class_names': ('aeroplane', 'bicycle', 'bird', 'boat',
+                         'bottle', 'bus', 'car', 'cat', 'chair',
+                         'cow', 'diningtable', 'dog', 'horse',
+                         'motorbike', 'person', 'pottedplant',
+                         'sheep', 'sofa', 'train', 'tvmonitor'),
+    },
+
+    'coco':{
+        'data_name': 'COCO',
+        'num_classes': 80,
+        'class_indexs': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20,
+                         21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
+                         41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58,
+                         59, 60, 61, 62, 63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79,
+                         80, 81, 82, 84, 85, 86, 87, 88, 89, 90],
+        'class_names': ('background',
+                        'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck',
+                        'boat', 'traffic light', 'fire hydrant', 'street sign', 'stop sign',
+                        'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
+                        'elephant', 'bear', 'zebra', 'giraffe', 'hat', 'backpack', 'umbrella',
+                        'shoe', 'eye glasses', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
+                        'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
+                        'skateboard', 'surfboard', 'tennis racket', 'bottle', 'plate', 'wine glass',
+                        'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
+                        'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
+                        'couch', 'potted plant', 'bed', 'mirror', 'dining table', 'window', 'desk',
+                        'toilet', 'door', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
+                        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'blender', 'book',
+                        'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'),
+    },
+
+    'ourdataset':{
+        'data_name': 'AnimalDataset',
+        'num_classes': 4,
+        'class_indexs': None,
+        'class_names': ('butterfly', 'cat', 'dog', 'person'),
+    },
+
+}

config/yolov1_config.py

@@ -8,6 +8,7 @@ yolov1_cfg = {
     'backbone': 'resnet18',
     'pretrained': True,
     'stride': 32,  # P5
+    'max_stride': 32,
     # neck
     'neck': 'sppf',
     'expand_ratio': 0.5,

config/yolov2_config.py

@@ -8,6 +8,7 @@ yolov2_cfg = {
     'backbone': 'darknet19',
     'pretrained': True,
     'stride': 32,  # P5
+    'max_stride': 32,
     # neck
     'neck': 'sppf',
     'expand_ratio': 0.5,

config/yolov3_config.py

@@ -9,6 +9,7 @@ yolov3_cfg = {
         'stride': [8, 16, 32],  # P3, P4, P5
         'width': 1.0,
         'depth': 1.0,
+        'max_stride': 32,
         ## Neck
         'neck': 'sppf',
         'expand_ratio': 0.5,
@@ -70,6 +71,7 @@ yolov3_cfg = {
         'stride': [8, 16, 32],  # P3, P4, P5
         'width': 0.25,
         'depth': 0.34,
+        'max_stride': 32,
         ## Neck
         'neck': 'sppf',
         'expand_ratio': 0.5,

config/yolov4_config.py

@@ -9,6 +9,7 @@ yolov4_cfg = {
         'stride': [8, 16, 32],  # P3, P4, P5
         'width': 1.0,
         'depth': 1.0,
+        'max_stride': 32,
         ## Neck
         'neck': 'csp_sppf',
         'expand_ratio': 0.5,
@@ -70,6 +71,7 @@ yolov4_cfg = {
         'stride': [8, 16, 32],  # P3, P4, P5
         'width': 0.25,
         'depth': 0.34,
+        'max_stride': 32,
         ## Neck
         'neck': 'csp_sppf',
         'expand_ratio': 0.5,

config/yolov5_config.py

@@ -12,6 +12,7 @@ yolov5_cfg = {
         'width': 0.25,
         'depth': 0.34,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         ## FPN
         'fpn': 'yolov5_pafpn',
         'fpn_reduce_layer': 'Conv',
@@ -72,6 +73,7 @@ yolov5_cfg = {
         'width': 0.50,
         'depth': 0.34,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         ## FPN
         'fpn': 'yolov5_pafpn',
         'fpn_reduce_layer': 'Conv',
@@ -132,6 +134,7 @@ yolov5_cfg = {
         'width': 0.75,
         'depth': 0.67,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         ## FPN
         'fpn': 'yolov5_pafpn',
         'fpn_reduce_layer': 'Conv',
@@ -192,6 +195,7 @@ yolov5_cfg = {
         'width': 1.0,
         'depth': 1.0,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         ## FPN
         'fpn': 'yolov5_pafpn',
         'fpn_reduce_layer': 'Conv',
@@ -252,6 +256,7 @@ yolov5_cfg = {
         'width': 1.25,
         'depth': 1.34,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         ## FPN
         'fpn': 'yolov5_pafpn',
         'fpn_reduce_layer': 'Conv',

config/yolov7_config.py

@@ -12,6 +12,7 @@ yolov7_cfg = {
         'bk_norm': 'BN',
         'bk_dpw': False,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         # neck
         'neck': 'csp_sppf',
         'expand_ratio': 0.5,
@@ -70,6 +71,7 @@ yolov7_cfg = {
         'bk_norm': 'BN',
         'bk_dpw': False,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         # neck
         'neck': 'csp_sppf',
         'expand_ratio': 0.5,
@@ -128,6 +130,7 @@ yolov7_cfg = {
         'bk_norm': 'BN',
         'bk_dpw': False,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         # neck
         'neck': 'csp_sppf',
         'expand_ratio': 0.5,

config/yolox_config.py

@@ -13,6 +13,7 @@ yolox_cfg = {
         'width': 0.25,
         'depth': 0.34,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         ## FPN
         'fpn': 'yolov5_pafpn',
         'fpn_reduce_layer': 'Conv',
@@ -71,6 +72,7 @@ yolox_cfg = {
         'width': 0.50,
         'depth': 0.34,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         ## FPN
         'fpn': 'yolov5_pafpn',
         'fpn_reduce_layer': 'Conv',
@@ -129,6 +131,7 @@ yolox_cfg = {
         'width': 0.75,
         'depth': 0.67,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         ## FPN
         'fpn': 'yolov5_pafpn',
         'fpn_reduce_layer': 'Conv',
@@ -187,6 +190,7 @@ yolox_cfg = {
         'width': 1.0,
         'depth': 1.0,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         ## FPN
         'fpn': 'yolov5_pafpn',
         'fpn_reduce_layer': 'Conv',
@@ -245,6 +249,7 @@ yolox_cfg = {
         'width': 1.25,
         'depth': 1.34,
         'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
         ## FPN
         'fpn': 'yolov5_pafpn',
         'fpn_reduce_layer': 'Conv',

dataset/build.py

@@ -0,0 +1,97 @@
+import os
+
+try:
+    from .voc import VOCDetection
+    from .coco import COCODataset
+    from .ourdataset import OurDataset
+    from .data_augment.ssd_augment import SSDAugmentation, SSDBaseTransform
+    from .data_augment.yolov5_augment import YOLOv5Augmentation, YOLOv5BaseTransform
+
+except:
+    import sys
+    sys.path.append('.')
+    from voc import VOCDetection
+    from coco import COCODataset
+    from ourdataset import OurDataset
+    from data_augment.ssd_augment import SSDAugmentation, SSDBaseTransform
+    from data_augment.yolov5_augment import YOLOv5Augmentation, YOLOv5BaseTransform
+
+# ------------------------------ Dataset ------------------------------
+def build_dataset(args, data_cfg, trans_config, transform, is_train=False):
+    # Basic parameters
+    data_dir = os.path.join(args.root, data_cfg['data_name'])
+    num_classes = data_cfg['num_classes']
+    class_names = data_cfg['class_names']
+    class_indexs = data_cfg['class_indexs']
+    dataset_info = {
+        'num_classes': num_classes,
+        'class_names': class_names,
+        'class_indexs': class_indexs
+    }
+
+    # Build dataset class
+    ## VOC dataset
+    if args.dataset == 'voc':
+        dataset = VOCDetection(
+            img_size=args.img_size,
+            data_dir=data_dir,
+            image_sets=[('2007', 'trainval'), ('2012', 'trainval')] if is_train else [('2007', 'test')],
+            transform=transform,
+            trans_config=trans_config
+            )
+    ## COCO dataset
+    elif args.dataset == 'coco':
+        dataset = COCODataset(
+            img_size=args.img_size,
+            data_dir=data_dir,
+            image_set='train2017' if is_train else 'val2017',
+            transform=transform,
+            trans_config=trans_config
+            )
+    ## Custom dataset
+    elif args.dataset == 'ourdataset':
+        dataset = OurDataset(
+            data_dir=data_dir,
+            img_size=args.img_size,
+            image_set='train' if is_train else 'val',
+            transform=transform,
+            trans_config=trans_config,
+            )
+
+    return dataset, dataset_info
+
+
+# ------------------------------ Transform ------------------------------
+def build_transform(args, trans_config, max_stride=32, is_train=False):
+    # Modify trans_config
+    ## mosaic prob.
+    if args.mosaic is not None:
+        trans_config['mosaic_prob']=args.mosaic if is_train else 0.0
+    else:
+        trans_config['mosaic_prob']=trans_config['mosaic_prob'] if is_train else 0.0
+    ## mixup prob.
+    if args.mixup is not None:
+        trans_config['mixup_prob']=args.mixup if is_train else 0.0
+    else:
+        trans_config['mixup_prob']=trans_config['mixup_prob']  if is_train else 0.0
+
+    # Transform
+    if trans_config['aug_type'] == 'ssd':
+        if is_train:
+            transform = SSDAugmentation(img_size=args.img_size,)
+        else:
+            transform = SSDBaseTransform(img_size=args.img_size,)
+
+    elif trans_config['aug_type'] == 'yolov5':
+        if is_train:
+            transform = YOLOv5Augmentation(
+                img_size=args.img_size,
+                trans_config=trans_config
+                )
+        else:
+            transform = YOLOv5BaseTransform(
+                img_size=args.img_size,
+                max_stride=max_stride
+                )
+
+    return transform, trans_config

dataset/data_augment/__init__.py

@@ -1,23 +0,0 @@
-from .ssd_augment import SSDAugmentation, SSDBaseTransform
-from .yolov5_augment import YOLOv5Augmentation, YOLOv5BaseTransform
-
-
-def build_transform(img_size, trans_config, is_train=False):
-    if trans_config['aug_type'] == 'ssd':
-        if is_train:
-            transform = SSDAugmentation(img_size=img_size)
-        else:
-            transform = SSDBaseTransform(img_size=img_size)
-
-    elif trans_config['aug_type'] == 'yolov5':
-        if is_train:
-            transform = YOLOv5Augmentation(
-                img_size=img_size,
-                trans_config=trans_config
-                )
-        else:
-            transform = YOLOv5BaseTransform(
-                img_size=img_size,
-                )
-
-    return transform

dataset/data_augment/yolov5_augment.py

@@ -302,12 +302,10 @@ def yolox_mixup_augment(origin_img, origin_target, new_img, new_target, img_size
 # YOLOv5-style TrainTransform
 class YOLOv5Augmentation(object):
     def __init__(self, 
-                 trans_config=None,
-                 img_size=640, 
-                 min_box_size=8):
+                 img_size=640,
+                 trans_config=None):
         self.trans_config = trans_config
         self.img_size = img_size
-        self.min_box_size = min_box_size
 
 
     def __call__(self, image, target, mosaic=False):

demo.py

@@ -8,7 +8,7 @@ import imageio
 import torch
 
 # load transform
-from dataset.data_augment import build_transform
+from dataset.build import build_transform
 
 # load some utils
 from utils.misc import load_weight
@@ -26,6 +26,10 @@ def parse_args():
     # basic
     parser.add_argument('-size', '--img_size', default=640, type=int,
                         help='the max size of input image')
+    parser.add_argument('--mosaic', default=None, type=float,
+                        help='mosaic augmentation.')
+    parser.add_argument('--mixup', default=None, type=float,
+                        help='mixup augmentation.')
     parser.add_argument('--mode', default='image',
                         type=str, help='Use the data from image, video or camera')
     parser.add_argument('--cuda', action='store_true', default=False,
@@ -271,14 +275,15 @@ def run():
     model.to(device).eval()
 
     # transform
-    transform = build_transform(args.img_size, trans_cfg, is_train=False)
+    val_transform, trans_cfg = build_transform(
+        args=args, trans_config=trans_cfg, max_stride=model_cfg['max_stride'], is_train=False)
 
     print("================= DETECT =================")
     # run
     detect(args=args,
            model=model, 
             device=device,
-            transform=transform,
+            transform=val_transform,
             mode=args.mode,
             vis_thresh=args.vis_thresh)
 

eval.py

@@ -9,15 +9,14 @@ from evaluator.coco_evaluator import COCOAPIEvaluator
 from evaluator.ourdataset_evaluator import OurDatasetEvaluator
 
 # load transform
-from dataset.ourdataset import our_class_labels
-from dataset.data_augment import build_transform
+from dataset.build import build_transform
 
 # load some utils
 from utils.misc import load_weight
 from utils.misc import compute_flops
 
+from config import build_dataset_config, build_model_config, build_trans_config
 from models.detectors import build_model
-from config import build_model_config, build_trans_config
 
 
 def parse_args():
@@ -49,6 +48,11 @@ def parse_args():
                         help='data root')
     parser.add_argument('-d', '--dataset', default='coco',
                         help='coco, voc.')
+    parser.add_argument('--mosaic', default=None, type=float,
+                        help='mosaic augmentation.')
+    parser.add_argument('--mixup', default=None, type=float,
+                        help='mixup augmentation.')
+
     # TTA
     parser.add_argument('-tta', '--test_aug', action='store_true', default=False,
                         help='use test augmentation.')
@@ -109,30 +113,13 @@ if __name__ == '__main__':
     else:
         device = torch.device("cpu")
 
-    # dataset
-    if args.dataset == 'voc':
-        print('eval on voc ...')
-        num_classes = 20
-        data_dir = os.path.join(args.root, 'VOCdevkit')
-    elif args.dataset == 'coco-val':
-        print('eval on coco-val ...')
-        num_classes = 80
-        data_dir = os.path.join(args.root, 'COCO')
-    elif args.dataset == 'coco-test':
-        print('eval on coco-test-dev ...')
-        num_classes = 80
-        data_dir = os.path.join(args.root, 'COCO')
-    elif args.dataset == 'ourdataset':
-        print('eval on crowdhuman ...')
-        num_classes = len(our_class_labels)
-        data_dir = os.path.join(args.root, 'OurDataset')
-    else:
-        print('unknow dataset !! we only support voc, coco-val, coco-test !!!')
-        exit(0)
-
-    # config
+    # Dataset & Model Config
+    data_cfg = build_dataset_config(args)
     model_cfg = build_model_config(args)
     trans_cfg = build_trans_config(model_cfg['trans_type'])
+    
+    data_dir = os.path.join(args.root, data_cfg['data_name'])
+    num_classes = data_cfg['num_classes']
 
     # build model
     model = build_model(args, model_cfg, device, num_classes, False)
@@ -152,15 +139,16 @@ if __name__ == '__main__':
     del model_copy
 
     # transform
-    transform = build_transform(args.img_size, trans_cfg, is_train=False)
+    val_transform, trans_cfg = build_transform(
+        args=args, trans_config=trans_cfg, max_stride=model_cfg['max_stride'], is_train=False)
 
     # evaluation
     with torch.no_grad():
         if args.dataset == 'voc':
-            voc_test(model, data_dir, device, transform)
+            voc_test(model, data_dir, device, val_transform)
         elif args.dataset == 'coco-val':
-            coco_test(model, data_dir, device, transform, test=False)
+            coco_test(model, data_dir, device, val_transform, test=False)
         elif args.dataset == 'coco-test':
-            coco_test(model, data_dir, device, transform, test=True)
+            coco_test(model, data_dir, device, val_transform, test=True)
         elif args.dataset == 'ourdataset':
-            our_test(model, data_dir, device, transform)
+            our_test(model, data_dir, device, val_transform)

evaluator/build.py

@@ -0,0 +1,36 @@
+import os
+
+from evaluator.coco_evaluator import COCOAPIEvaluator
+from evaluator.voc_evaluator import VOCAPIEvaluator
+from evaluator.ourdataset_evaluator import OurDatasetEvaluator
+
+
+
+def build_evluator(args, data_cfg, transform, device):
+    # Basic parameters
+    data_dir = os.path.join(args.root, data_cfg['data_name'])
+
+    # Evaluator
+    ## VOC Evaluator
+    if args.dataset == 'voc':
+        evaluator = VOCAPIEvaluator(
+            data_dir=data_dir,
+            device=device,
+            transform=transform)
+    ## COCO Evaluator
+    elif args.dataset == 'coco':
+        evaluator = COCOAPIEvaluator(
+            data_dir=data_dir,
+            device=device,
+            transform=transform
+            )
+    ## Custom dataset Evaluator
+    elif args.dataset == 'ourdataset':
+        evaluator = OurDatasetEvaluator(
+            data_dir=data_dir,
+            device=device,
+            image_set='val',
+            transform=transform
+        )
+
+    return evaluator

evaluator/coco_evaluator.py

@@ -28,19 +28,18 @@ class COCOAPIEvaluator():
             nmsthre (float):
                 IoU threshold of non-max supression ranging from 0 to 1.
         """
-        self.testset = testset
-        if self.testset:
-            image_set = 'test2017'
-        else:
-            image_set = 'val2017'
-
-        self.dataset = COCODataset(data_dir=data_dir, image_set=image_set, is_train=False)
+        # ----------------- Basic parameters -----------------
+        self.image_set = 'test2017' if testset else 'val2017'
         self.transform = transform
         self.device = device
-
+        self.testset = testset
+        # ----------------- Metrics -----------------
         self.map = 0.
         self.ap50_95 = 0.
         self.ap50 = 0.
+        # ----------------- Dataset -----------------
+        self.dataset = COCODataset(data_dir=data_dir, image_set=self.image_set)
+
 
     @torch.no_grad()
     def evaluate(self, model):

evaluator/ourdataset_evaluator.py

@@ -28,14 +28,17 @@ class OurDatasetEvaluator():
             nmsthre (float):
                 IoU threshold of non-max supression ranging from 0 to 1.
         """
-        self.dataset = OurDataset(data_dir=data_dir, image_set=image_set, is_train=False)
+        # ----------------- Basic parameters -----------------
         self.image_set = image_set
         self.transform = transform
         self.device = device
-
+        # ----------------- Metrics -----------------
         self.map = 0.
         self.ap50_95 = 0.
         self.ap50 = 0.
+        # ----------------- Dataset -----------------
+        self.dataset = OurDataset(data_dir=data_dir, image_set=image_set)
+
 
     @torch.no_grad()
     def evaluate(self, model):

test.py

@@ -7,15 +7,14 @@ from copy import deepcopy
 import torch
 
 # load transform
-from dataset.data_augment import build_transform
+from dataset.build import build_dataset, build_transform
 
 # load some utils
-from utils.misc import build_dataset, load_weight
-from utils.misc import compute_flops
+from utils.misc import load_weight, compute_flops
 from utils.box_ops import rescale_bboxes
 
+from config import build_dataset_config, build_model_config, build_trans_config
 from models.detectors import build_model
-from config import build_model_config, build_trans_config
 
 
 def parse_args():
@@ -175,13 +174,18 @@ if __name__ == '__main__':
     else:
         device = torch.device("cpu")
 
-    # config
+    # Dataset & Model Config
+    data_cfg = build_dataset_config(args)
     model_cfg = build_model_config(args)
     trans_cfg = build_trans_config(model_cfg['trans_type'])
 
-    # dataset and evaluator
-    dataset, dataset_info, evaluator = build_dataset(args, trans_cfg, device, is_train=False)
-    num_classes, class_names, class_indexs = dataset_info
+    # Transform
+    val_transform, trans_cfg = build_transform(
+        args=args, trans_config=trans_cfg, max_stride=model_cfg['max_stride'], is_train=False)
+
+    # Dataset
+    dataset, dataset_info = build_dataset(args, data_cfg, trans_cfg, val_transform, is_train=False)
+    num_classes = dataset_info['num_classes']
 
     np.random.seed(0)
     class_colors = [(np.random.randint(255),
@@ -205,17 +209,14 @@ if __name__ == '__main__':
         device=device)
     del model_copy
 
-    # transform
-    transform = build_transform(args.img_size, trans_cfg, is_train=False)
-
     print("================= DETECT =================")
     # run
     test(args=args,
          model=model, 
          device=device, 
          dataset=dataset,
-         transform=transform,
+         transform=val_transform,
          class_colors=class_colors,
-         class_names=class_names,
-         class_indexs=class_indexs
+         class_names=dataset_info['class_names'],
+         class_indexs=dataset_info['class_indexs'],
          )

train.py

@@ -4,19 +4,31 @@ import os
 import argparse
 from copy import deepcopy
 
+# ----------------- Torch Components -----------------
 import torch
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 
+# ----------------- Extra Components -----------------
 from utils import distributed_utils
 from utils.misc import compute_flops
 from utils.misc import ModelEMA, CollateFunc, build_dataset, build_dataloader
+
+# ----------------- Evaluator Components -----------------
+from evaluator.build import build_evluator
+
+# ----------------- Optimizer & LrScheduler Components -----------------
 from utils.solver.optimizer import build_optimizer
 from utils.solver.lr_scheduler import build_lr_scheduler
 
 from engine import train_one_epoch, val_one_epoch
+# ----------------- Config Components -----------------
+from config import build_dataset_config, build_model_config, build_trans_config
+
+# ----------------- Dataset Components -----------------
+from dataset.build import build_dataset, build_transform
 
-from config import build_model_config, build_trans_config
+# ----------------- Model Components -----------------
 from models.detectors import build_model
 
 
@@ -124,23 +136,32 @@ def train():
     else:
         device = torch.device("cpu")
 
-    # config
+    # Dataset & Model & Trans Config
+    data_cfg = build_dataset_config(args)
     model_cfg = build_model_config(args)
     trans_cfg = build_trans_config(model_cfg['trans_type'])
 
-    # dataset and evaluator
-    dataset, dataset_info, evaluator = build_dataset(args, trans_cfg, device, is_train=True)
-    num_classes = dataset_info[0]
+    # Transform
+    train_transform, trans_config = build_transform(
+        args=args, trans_config=trans_cfg, max_stride=model_cfg['max_stride'], is_train=True)
+    val_transform, _ = build_transform(
+        args=args, max_stride=model_cfg['max_stride'], is_train=False)
 
-    # dataloader
+    # Dataset
+    dataset, dataset_info = build_dataset(args, data_cfg, trans_config, train_transform, is_train=True)
+
+    # Dataloader
     dataloader = build_dataloader(args, dataset, per_gpu_batch, CollateFunc())
 
-    # build model
+    # Evaluator
+    evaluator = build_evluator(args, data_cfg, val_transform, device)
+
+    # Build model
     model, criterion = build_model(
         args=args, 
         model_cfg=model_cfg,
         device=device,
-        num_classes=num_classes,
+        num_classes=dataset_info['num_classes'],
         trainable=True,
         )
     model = model.to(device).train()