yjh0410 há 1 ano atrás
pai
commit
37e209ea2e
5 ficheiros alterados com 17 adições e 29 exclusões
  1. 6 7
      yolo/dataset/coco.py
  2. 1 3
      yolo/dataset/custom.py
  3. 1 1
      yolo/dataset/voc.py
  4. 3 12
      yolo/eval.py
  5. 6 6
      yolo/train.py

+ 6 - 7
yolo/dataset/coco.py

@@ -57,7 +57,7 @@ class COCODataset(Dataset):
             self.copy_paste  = 0.0
             self.mosaic_augment = None
             self.mixup_augment  = None
-        print('==============================')
+        print(' ============ Strong augmentation info. ============ ')
         print('use Mosaic Augmentation: {}'.format(self.mosaic_prob))
         print('use Mixup Augmentation: {}'.format(self.mixup_prob))
         print('use Copy-paste Augmentation: {}'.format(self.copy_paste))
@@ -162,15 +162,14 @@ class COCODataset(Dataset):
 
     def pull_anno(self, index):
         img_id = self.ids[index]
-        im_ann = self.coco.loadImgs(img_id)[0]
-        anno_ids = self.coco.getAnnIds(imgIds=[int(img_id)], iscrowd=False)
-        annotations = self.coco.loadAnns(anno_ids)
-
         # image infor
+        im_ann = self.coco.loadImgs(img_id)[0]
         width = im_ann['width']
         height = im_ann['height']
-        
-        #load a target
+ 
+        # load a target
+        anno_ids = self.coco.getAnnIds(imgIds=[int(img_id)], iscrowd=False)
+        annotations = self.coco.loadAnns(anno_ids)
         bboxes = []
         labels = []
         for anno in annotations:

+ 1 - 3
yolo/dataset/custom.py

@@ -51,9 +51,7 @@ class CustomDataset(Dataset):
             self.copy_paste  = 0.0
             self.mosaic_augment = None
             self.mixup_augment  = None
-        print('==============================')
-        print('Image Set: {}'.format(self.image_set))
-        print('Json file: {}'.format(self.json_file))
+        print(' ============ Strong augmentation info. ============ ')
         print('use Mosaic Augmentation: {}'.format(self.mosaic_prob))
         print('use Mixup Augmentation: {}'.format(self.mixup_prob))
         print('use Copy-paste Augmentation: {}'.format(self.copy_paste))

+ 1 - 1
yolo/dataset/voc.py

@@ -50,7 +50,7 @@ class VOCDataset(Dataset):
             self.copy_paste  = 0.0
             self.mosaic_augment = None
             self.mixup_augment  = None
-        print('==============================')
+        print(' ============ Strong augmentation info. ============ ')
         print('use Mosaic Augmentation: {}'.format(self.mosaic_prob))
         print('use Mixup Augmentation: {}'.format(self.mixup_prob))
         print('use Copy-paste Augmentation: {}'.format(self.copy_paste))

+ 3 - 12
yolo/eval.py

@@ -1,13 +1,8 @@
 import argparse
 import torch
 
-# evaluators
 from evaluator.map_evaluator import MapEvaluator
-
-# load transform
 from dataset.build import build_dataset, build_transform
-
-# load some utils
 from utils.misc import load_weight
 
 from config import build_config
@@ -17,17 +12,17 @@ from models import build_model
 def parse_args():
     parser = argparse.ArgumentParser(description='Real-time Object Detection LAB')
     # Basic setting
-    parser.add_argument('-size', '--img_size', default=640, type=int,
+    parser.add_argument('--img_size', default=640, type=int,
                         help='the max size of input image')
     parser.add_argument('--cuda', action='store_true', default=False,
                         help='Use cuda')
 
     # Model setting
-    parser.add_argument('-m', '--model', default='yolov1', type=str,
+    parser.add_argument('--model', default='yolov1', type=str,
                         help='build yolo')
     parser.add_argument('--weight', default=None,
                         type=str, help='Trained state_dict file path to open')
-    parser.add_argument('-r', '--resume', default=None, type=str,
+    parser.add_argument('--resume', default=None, type=str,
                         help='keep training')
     parser.add_argument('--fuse_conv_bn', action='store_true', default=False,
                         help='fuse Conv & BN')
@@ -38,10 +33,6 @@ def parse_args():
     parser.add_argument('-d', '--dataset', default='coco',
                         help='coco, voc.')
 
-    # TTA
-    parser.add_argument('-tta', '--test_aug', action='store_true', default=False,
-                        help='use test augmentation.')
-
     return parser.parse_args()
 
 

+ 6 - 6
yolo/train.py

@@ -60,27 +60,27 @@ def parse_args():
                         help="Adopting mix precision training.")
     
     # Batchsize
-    parser.add_argument('-bs', '--batch_size', default=16, type=int, 
+    parser.add_argument('--batch_size', default=16, type=int, 
                         help='batch size on all the GPUs.')
 
     # Model
-    parser.add_argument('-m', '--model', default='yolo_n', type=str,
+    parser.add_argument('--model', default='yolo_n', type=str,
                         help='build yolo')
-    parser.add_argument('-p', '--pretrained', default=None, type=str,
+    parser.add_argument('--pretrained', default=None, type=str,
                         help='load pretrained weight')
-    parser.add_argument('-r', '--resume', default=None, type=str,
+    parser.add_argument('--resume', default=None, type=str,
                         help='keep training')
 
     # Dataset
     parser.add_argument('--root', default='D:/python_work/dataset/VOCdevkit/',
                         help='data root')
-    parser.add_argument('-d', '--dataset', default='coco',
+    parser.add_argument('--dataset', default='voc',
                         help='coco, voc')
     parser.add_argument('--num_workers', default=4, type=int, 
                         help='Number of workers used in dataloading')
     
     # DDP train
-    parser.add_argument('-dist', '--distributed', action='store_true', default=False,
+    parser.add_argument('--distributed', action='store_true', default=False,
                         help='distributed training')
     parser.add_argument('--dist_url', default='env://', 
                         help='url used to set up distributed training')