scripts: use data config defined class names

Shorten name of --data-config-path argument to --data-config

detect.py

@@ -10,6 +10,7 @@ from utils import torch_utils
 
 def detect(
         net_config_path,
+        data_config_path,
         images_path,
         weights_file_path='weights/yolov3.pt',
         output='output',
@@ -19,7 +20,6 @@ def detect(
         nms_thres=0.45,
         save_txt=False,
         save_images=False,
-        class_path='data/coco.names',
 ):
 
     device = torch_utils.select_device()
@@ -28,6 +28,8 @@ def detect(
     os.system('rm -rf ' + output)
     os.makedirs(output, exist_ok=True)
 
+    data_config = parse_data_config(data_config_path)
+
     # Load model
     model = Darknet(net_config_path, img_size)
 
@@ -55,7 +57,7 @@ def detect(
     model.to(device).eval()
 
     # Set Dataloader
-    classes = load_classes(class_path)  # Extracts class labels from file
+    classes = load_classes(data_config['names'])  # Extracts class labels from file
     dataloader = load_images(images_path, batch_size=batch_size, img_size=img_size)
 
     imgs = []  # Stores image paths
@@ -151,7 +153,7 @@ if __name__ == '__main__':
     parser.add_argument('--txt-out', type=bool, default=False)
 
     parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
-    parser.add_argument('--class-path', type=str, default='data/coco.names', help='path to class label file')
+    parser.add_argument('--data-config', type=str, default='cfg/coco.data', help='path to data config file')
     parser.add_argument('--conf-thres', type=float, default=0.50, help='object confidence threshold')
     parser.add_argument('--nms-thres', type=float, default=0.45, help='iou threshold for non-maximum suppression')
     parser.add_argument('--batch-size', type=int, default=1, help='size of the batches')
@@ -165,6 +167,7 @@ if __name__ == '__main__':
 
     detect(
         opt.cfg,
+        opt.data_config,
         opt.image_folder,
         output=opt.output_folder,
         batch_size=opt.batch_size,
@@ -173,5 +176,4 @@ if __name__ == '__main__':
         nms_thres=opt.nms_thres,
         save_txt=opt.txt_out,
         save_images=opt.plot_flag,
-        class_path=opt.class_path,
     )

test.py

@@ -11,7 +11,6 @@ def test(
     net_config_path,
     data_config_path,
     weights_file_path,
-    class_path=None,
     batch_size=16,
     img_size=416,
     iou_thres=0.5,
@@ -118,10 +117,9 @@ def test(
     # Print mAP per class
     print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP') + '\n\nmAP Per Class:')
 
-    if class_path:
-        classes = load_classes(class_path)  # Extracts class labels from file
-        for i, c in enumerate(classes):
-            print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i]))
+    classes = load_classes(data_config['names'])  # Extracts class labels from file
+    for i, c in enumerate(classes):
+        print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i]))
 
     # Return mAP
     return mean_mAP, mean_R, mean_P
@@ -132,9 +130,8 @@ if __name__ == '__main__':
     parser = argparse.ArgumentParser(prog='test.py')
     parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
     parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='path to model config file')
-    parser.add_argument('--data-config-path', type=str, default='cfg/coco.data', help='path to data config file')
+    parser.add_argument('--data-config', type=str, default='cfg/coco.data', help='path to data config file')
     parser.add_argument('--weights', type=str, default='weights/yolov3.pt', help='path to weights file')
-    parser.add_argument('--class-path', type=str, default='data/coco.names', help='path to class label file')
     parser.add_argument('--iou-thres', type=float, default=0.5, help='iou threshold required to qualify as detected')
     parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')
     parser.add_argument('--nms-thres', type=float, default=0.45, help='iou threshold for non-maximum suppression')
@@ -147,9 +144,8 @@ if __name__ == '__main__':
 
     mAP = test(
         opt.cfg,
-        opt.data_config_path,
+        opt.data_config,
         opt.weights,
-        class_path=opt.class_path,
         batch_size=opt.batch_size,
         img_size=opt.img_size,
         iou_thres=opt.iou_thres,

train.py

@@ -220,7 +220,7 @@ if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--epochs', type=int, default=100, help='number of epochs')
     parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch')
-    parser.add_argument('--data-config-path', type=str, default='cfg/coco.data', help='data config file path')
+    parser.add_argument('--data-config', type=str, default='cfg/coco.data', help='path to data config file')
     parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
     parser.add_argument('--multi-scale', default=False, help='random image sizes per batch 320 - 608')
     parser.add_argument('--img-size', type=int, default=32 * 13, help='pixels')
@@ -236,7 +236,7 @@ if __name__ == '__main__':
     torch.cuda.empty_cache()
     train(
         opt.cfg,
-        opt.data_config_path,
+        opt.data_config,
         img_size=opt.img_size,
         resume=opt.resume,
         epochs=opt.epochs,