update

config/dataset_config.py

@@ -1,5 +1,4 @@
 # Dataset config
-import os
 
 dataset_cfg = {
     'voc': {

dataset/build.py

@@ -8,14 +8,13 @@ try:
     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

train.py

@@ -142,13 +142,13 @@ def train():
     trans_cfg = build_trans_config(model_cfg['trans_type'])
 
     # Transform
-    train_transform, trans_config = build_transform(
+    train_transform, trans_cfg = build_transform(
         args=args, trans_config=trans_cfg, max_stride=model_cfg['max_stride'], is_train=True)
     val_transform, _ = 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_config, train_transform, is_train=True)
+    dataset, dataset_info = build_dataset(args, data_cfg, trans_cfg, train_transform, is_train=True)
 
     # Dataloader
     dataloader = build_dataloader(args, dataset, per_gpu_batch, CollateFunc())