From 5a566454f5efc4ce424c84826358b5de602d7aad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Guillermo=20Garc=C3=ADa?= <ggarcia@redpoints.com>
Date: Wed, 5 Dec 2018 11:55:27 +0100
Subject: [PATCH 1/7] Extract seed and cuda initialization utils

---
 detect.py            |  6 ++++--
 test.py              | 11 ++++++++---
 train.py             | 20 ++++++++++----------
 utils/torch_utils.py | 23 +++++++++++++++++++++++
 utils/utils.py       |  8 ++++++++
 5 files changed, 53 insertions(+), 15 deletions(-)
 create mode 100644 utils/torch_utils.py

diff --git a/detect.py b/detect.py
index 7953f30e..cc2ef727 100755
--- a/detect.py
+++ b/detect.py
@@ -5,8 +5,6 @@ from models import *
 from utils.datasets import *
 from utils.utils import *
 
-cuda = torch.cuda.is_available()
-device = torch.device('cuda:0' if cuda else 'cpu')
 f_path = os.path.dirname(os.path.realpath(__file__)) + '/'
 
 parser = argparse.ArgumentParser()
@@ -28,6 +26,10 @@ print(opt)
 
 
 def main(opt):
+
+    device = torch_utils.select_device()
+    print("Using device: \"{}\"".format(device))
+
     os.system('rm -rf ' + opt.output_folder)
     os.makedirs(opt.output_folder, exist_ok=True)
 
diff --git a/test.py b/test.py
index f65d373c..c0cf476b 100644
--- a/test.py
+++ b/test.py
@@ -4,6 +4,8 @@ from models import *
 from utils.datasets import *
 from utils.utils import *
 
+from utils import torch_utils
+
 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')
@@ -18,11 +20,11 @@ parser.add_argument('-img_size', type=int, default=416, help='size of each image
 opt = parser.parse_args()
 print(opt, end='\n\n')
 
-cuda = torch.cuda.is_available()
-device = torch.device('cuda:0' if cuda else 'cpu')
-
 
 def main(opt):
+    device = torch_utils.select_device()
+    print("Using device: \"{}\"".format(device))
+
     # Configure run
     data_config = parse_data_config(opt.data_config_path)
     nC = int(data_config['classes'])  # number of classes (80 for COCO)
@@ -128,4 +130,7 @@ def main(opt):
 
 
 if __name__ == '__main__':
+
+    init_seeds()
+
     mAP = main(opt)
diff --git a/train.py b/train.py
index 93ac9c96..103d6f03 100644
--- a/train.py
+++ b/train.py
@@ -6,6 +6,8 @@ from models import *
 from utils.datasets import *
 from utils.utils import *
 
+from utils import torch_utils
+
 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')
@@ -26,20 +28,15 @@ print(opt)
 sys.argv[1:] = []  # delete any train.py command-line arguments before they reach test.py
 import test  # must follow sys.argv[1:] = []
 
-cuda = torch.cuda.is_available()
-device = torch.device('cuda:0' if cuda else 'cpu')
 
-random.seed(0)
-np.random.seed(0)
-torch.manual_seed(0)
-if cuda:
-    torch.cuda.manual_seed(0)
-    torch.cuda.manual_seed_all(0)
+def main(opt):
+
+    device = torch_utils.select_device()
+    print("Using device: \"{}\"".format(device))
+
     if not opt.multi_scale:
         torch.backends.cudnn.benchmark = True
 
-
-def main(opt):
     os.makedirs('weights', exist_ok=True)
 
     # Configure run
@@ -217,5 +214,8 @@ def main(opt):
 
 
 if __name__ == '__main__':
+
+    init_seeds()
+
     torch.cuda.empty_cache()
     main(opt)
diff --git a/utils/torch_utils.py b/utils/torch_utils.py
new file mode 100644
index 00000000..58bf5ff4
--- /dev/null
+++ b/utils/torch_utils.py
@@ -0,0 +1,23 @@
+import torch
+
+
+def check_cuda():
+    return torch.cuda.is_available()
+
+
+CUDA_AVAILABLE = check_cuda()
+
+
+def init_seeds(seed=0):
+    torch.manual_seed(seed)
+    if CUDA_AVAILABLE:
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+
+def select_device(force_cpu=False):
+    if force_cpu:
+        device = torch.device('cpu')
+    else:
+        device = torch.device('cuda:0' if CUDA_AVAILABLE else 'cpu')
+    return device
diff --git a/utils/utils.py b/utils/utils.py
index 6fcf5fac..12d161bd 100755
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -5,11 +5,19 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 
+from utils import torch_utils
+
 # Set printoptions
 torch.set_printoptions(linewidth=1320, precision=5, profile='long')
 np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format})  # format short g, %precision=5
 
 
+def init_seeds(seed=0):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch_utils.init_seeds(seed=seed)
+
+
 def load_classes(path):
     """
     Loads class labels at 'path'

From c807c16b7909ad6885f57f35bab678b3f3bc35f0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Guillermo=20Garc=C3=ADa?= <ggarcia@redpoints.com>
Date: Wed, 5 Dec 2018 14:31:08 +0100
Subject: [PATCH 2/7] Fix argument parser bad practice

Keep parsing inside __main__ block and call methods with arguments

Add double -- for long argument names (- reserved for shortcuts)
---
 README.md    |   2 +-
 detect.py    | 111 +++++++++++++++++++++++++++++++--------------------
 test.py      |  81 +++++++++++++++++++++++--------------
 train.py     |  93 +++++++++++++++++++++++++++---------------
 utils/gcp.sh |  14 +++----
 5 files changed, 188 insertions(+), 113 deletions(-)
 mode change 100644 => 100755 utils/gcp.sh

diff --git a/README.md b/README.md
index 0ddad037..cbb37de9 100755
--- a/README.md
+++ b/README.md
@@ -56,7 +56,7 @@ Checkpoints are saved in `/checkpoints` directory. Run `detect.py` to apply trai
 
 Run `test.py` to validate the official YOLOv3 weights `checkpoints/yolov3.weights` against the 5000 validation images. You should obtain a mAP of .581 using this repo (https://github.com/ultralytics/yolov3), compared to .579 as reported in darknet (https://arxiv.org/abs/1804.02767).
 
-Run `test.py -weights_path checkpoints/latest.pt` to validate against the latest training checkpoint.
+Run `test.py --weights checkpoints/latest.pt` to validate against the latest training checkpoint.
 
 # Contact
 
diff --git a/detect.py b/detect.py
index cc2ef727..a104e8a3 100755
--- a/detect.py
+++ b/detect.py
@@ -5,45 +5,39 @@ from models import *
 from utils.datasets import *
 from utils.utils import *
 
-f_path = os.path.dirname(os.path.realpath(__file__)) + '/'
-
-parser = argparse.ArgumentParser()
-# Get data configuration
-
-parser.add_argument('-image_folder', type=str, default='data/samples', help='path to images')
-parser.add_argument('-output_folder', type=str, default='output', help='path to outputs')
-parser.add_argument('-plot_flag', type=bool, default=True)
-parser.add_argument('-txt_out', type=bool, default=False)
-
-parser.add_argument('-cfg', type=str, default=f_path + 'cfg/yolov3.cfg', help='cfg file path')
-parser.add_argument('-class_path', type=str, default=f_path + 'data/coco.names', help='path to class label 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')
-parser.add_argument('-img_size', type=int, default=32 * 13, help='size of each image dimension')
-opt = parser.parse_args()
-print(opt)
+from utils import torch_utils
 
 
-def main(opt):
+def detect(
+        net_config_path,
+        images_path,
+        weights_file_path='weights/yolov3.pt',
+        output='output',
+        batch_size=16,
+        img_size=416,
+        conf_thres=0.3,
+        nms_thres=0.45,
+        save_txt=False,
+        save_images=False,
+        class_path='data/coco.names',
+):
 
     device = torch_utils.select_device()
     print("Using device: \"{}\"".format(device))
 
-    os.system('rm -rf ' + opt.output_folder)
-    os.makedirs(opt.output_folder, exist_ok=True)
+    os.system('rm -rf ' + output)
+    os.makedirs(output, exist_ok=True)
 
     # Load model
-    model = Darknet(opt.cfg, opt.img_size)
+    model = Darknet(net_config_path, img_size)
 
-    weights_path = f_path + 'weights/yolov3.pt'
-    if weights_path.endswith('.pt'):  # pytorch format
-        if weights_path.endswith('weights/yolov3.pt') and not os.path.isfile(weights_path):
-            os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights_path)
+    if weights_file_path.endswith('.pt'):  # pytorch format
+        if weights_file_path.endswith('weights/yolov3.pt') and not os.path.isfile(weights_file_path):
+            os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights_file_path)
     else:  # darknet format
-        load_weights(model, weights_path)
+        load_weights(model, weights_file_path)
 
-        checkpoint = torch.load(weights_path, map_location='cpu')
+        checkpoint = torch.load(weights_file_path, map_location='cpu')
         model.load_state_dict(checkpoint['model'])
         del checkpoint
 
@@ -61,8 +55,8 @@ def main(opt):
     model.to(device).eval()
 
     # Set Dataloader
-    classes = load_classes(opt.class_path)  # Extracts class labels from file
-    dataloader = load_images(opt.image_folder, batch_size=opt.batch_size, img_size=opt.img_size)
+    classes = load_classes(class_path)  # Extracts class labels from file
+    dataloader = load_images(images_path, batch_size=batch_size, img_size=img_size)
 
     imgs = []  # Stores image paths
     img_detections = []  # Stores detections for each image index
@@ -73,10 +67,10 @@ def main(opt):
         # Get detections
         with torch.no_grad():
             pred = model(torch.from_numpy(img).unsqueeze(0).to(device))
-            pred = pred[pred[:, :, 4] > opt.conf_thres]
+            pred = pred[pred[:, :, 4] > conf_thres]
 
             if len(pred) > 0:
-                detections = non_max_suppression(pred.unsqueeze(0), opt.conf_thres, opt.nms_thres)
+                detections = non_max_suppression(pred.unsqueeze(0), conf_thres, nms_thres)
                 img_detections.extend(detections)
                 imgs.extend(img_paths)
 
@@ -93,15 +87,15 @@ def main(opt):
     for img_i, (path, detections) in enumerate(zip(imgs, img_detections)):
         print("image %g: '%s'" % (img_i, path))
 
-        if opt.plot_flag:
+        if save_images:
             img = cv2.imread(path)
 
         # The amount of padding that was added
-        pad_x = max(img.shape[0] - img.shape[1], 0) * (opt.img_size / max(img.shape))
-        pad_y = max(img.shape[1] - img.shape[0], 0) * (opt.img_size / max(img.shape))
+        pad_x = max(img.shape[0] - img.shape[1], 0) * (img_size / max(img.shape))
+        pad_y = max(img.shape[1] - img.shape[0], 0) * (img_size / max(img.shape))
         # Image height and width after padding is removed
-        unpad_h = opt.img_size - pad_y
-        unpad_w = opt.img_size - pad_x
+        unpad_h = img_size - pad_y
+        unpad_w = img_size - pad_x
 
         # Draw bounding boxes and labels of detections
         if detections is not None:
@@ -109,7 +103,7 @@ def main(opt):
             bbox_colors = random.sample(color_list, len(unique_classes))
 
             # write results to .txt file
-            results_img_path = os.path.join(opt.output_folder, path.split('/')[-1])
+            results_img_path = os.path.join(output, path.split('/')[-1])
             results_txt_path = results_img_path + '.txt'
             if os.path.isfile(results_txt_path):
                 os.remove(results_txt_path)
@@ -129,24 +123,55 @@ def main(opt):
                 x1, y1, x2, y2 = max(x1, 0), max(y1, 0), max(x2, 0), max(y2, 0)
 
                 # write to file
-                if opt.txt_out:
+                if save_txt:
                     with open(results_txt_path, 'a') as file:
                         file.write(('%g %g %g %g %g %g \n') % (x1, y1, x2, y2, cls_pred, cls_conf * conf))
 
-                if opt.plot_flag:
+                if save_images:
                     # Add the bbox to the plot
                     label = '%s %.2f' % (classes[int(cls_pred)], conf)
                     color = bbox_colors[int(np.where(unique_classes == int(cls_pred))[0])]
                     plot_one_box([x1, y1, x2, y2], img, label=label, color=color)
 
-        if opt.plot_flag:
+        if save_images:
             # Save generated image with detections
             cv2.imwrite(results_img_path.replace('.bmp', '.jpg').replace('.tif', '.jpg'), img)
 
     if platform == 'darwin':  # MacOS (local)
-        os.system('open ' + opt.output_folder)
+        os.system('open ' + output)
 
 
 if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    # Get data configuration
+
+    parser.add_argument('--image-folder', type=str, default='data/samples', help='path to images')
+    parser.add_argument('--output-folder', type=str, default='output', help='path to outputs')
+    parser.add_argument('--plot-flag', type=bool, default=True)
+    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('--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')
+    parser.add_argument('--img-size', type=int, default=32 * 13, help='size of each image dimension')
+    opt = parser.parse_args()
+    print(opt)
+
     torch.cuda.empty_cache()
-    main(opt)
+
+    init_seeds()
+
+    detect(
+        opt.cfg,
+        opt.image_folder,
+        output=opt.output_folder,
+        batch_size=opt.batch_size,
+        img_size=opt.img_size,
+        conf_thres=opt.conf_thres,
+        nms_thres=opt.nms_thres,
+        save_txt=opt.txt_out,
+        save_images=opt.plot_flag,
+        class_path=opt.class_path,
+    )
diff --git a/test.py b/test.py
index c0cf476b..4630a0ee 100644
--- a/test.py
+++ b/test.py
@@ -6,47 +6,44 @@ from utils.utils import *
 
 from utils import torch_utils
 
-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('-weights_path', 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')
-parser.add_argument('-n_cpu', type=int, default=0, help='number of cpu threads to use during batch generation')
-parser.add_argument('-img_size', type=int, default=416, help='size of each image dimension')
-opt = parser.parse_args()
-print(opt, end='\n\n')
 
-
-def main(opt):
+def test(
+    net_config_path,
+    data_config_path,
+    weights_file_path,
+    class_path=None,
+    batch_size=16,
+    img_size=416,
+    iou_thres=0.5,
+    conf_thres=0.3,
+    nms_thres=0.45,
+    n_cpus=0,
+):
     device = torch_utils.select_device()
     print("Using device: \"{}\"".format(device))
 
     # Configure run
-    data_config = parse_data_config(opt.data_config_path)
+    data_config = parse_data_config(data_config_path)
     nC = int(data_config['classes'])  # number of classes (80 for COCO)
     test_path = data_config['valid']
 
     # Initiate model
-    model = Darknet(opt.cfg, opt.img_size)
+    model = Darknet(net_config_path, img_size)
 
     # Load weights
-    if opt.weights_path.endswith('.pt'):  # pytorch format
-        checkpoint = torch.load(opt.weights_path, map_location='cpu')
+    if weights_file_path.endswith('.pt'):  # pytorch format
+        checkpoint = torch.load(weights_file_path, map_location='cpu')
         model.load_state_dict(checkpoint['model'])
         del checkpoint
     else:  # darknet format
-        load_weights(model, opt.weights_path)
+        load_weights(model, weights_file_path)
 
     model.to(device).eval()
 
     # Get dataloader
     # dataset = load_images_with_labels(test_path)
-    # dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=False, num_workers=opt.n_cpu)
-    dataloader = load_images_and_labels(test_path, batch_size=opt.batch_size, img_size=opt.img_size)
+    # dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=n_cpus)
+    dataloader = load_images_and_labels(test_path, batch_size=batch_size, img_size=img_size)
 
     print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP'))
     outputs, mAPs, mR, mP, TP, confidence, pred_class, target_class = [], [], [], [], [], [], [], []
@@ -55,7 +52,7 @@ def main(opt):
 
         with torch.no_grad():
             output = model(imgs.to(device))
-            output = non_max_suppression(output, conf_thres=opt.conf_thres, nms_thres=opt.nms_thres)
+            output = non_max_suppression(output, conf_thres=conf_thres, nms_thres=nms_thres)
 
         # Compute average precision for each sample
         for sample_i, (labels, detections) in enumerate(zip(targets, output)):
@@ -80,7 +77,7 @@ def main(opt):
                 target_cls = labels[:, 0]
 
                 # Extract target boxes as (x1, y1, x2, y2)
-                target_boxes = xywh2xyxy(labels[:, 1:5]) * opt.img_size
+                target_boxes = xywh2xyxy(labels[:, 1:5]) * img_size
 
                 detected = []
                 for *pred_bbox, conf, obj_conf, obj_pred in detections:
@@ -91,7 +88,7 @@ def main(opt):
                     # Extract index of largest overlap
                     best_i = np.argmax(iou)
                     # If overlap exceeds threshold and classification is correct mark as correct
-                    if iou[best_i] > opt.iou_thres and obj_pred == labels[best_i, 0] and best_i not in detected:
+                    if iou[best_i] > iou_thres and obj_pred == labels[best_i, 0] and best_i not in detected:
                         correct.append(1)
                         detected.append(best_i)
                     else:
@@ -121,9 +118,10 @@ def main(opt):
     # Print mAP per class
     print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP') + '\n\nmAP Per Class:')
 
-    classes = load_classes(opt.class_path)  # Extracts class labels from file
-    for i, c in enumerate(classes):
-        print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i]))
+    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]))
 
     # Return mAP
     return mean_mAP, mean_R, mean_P
@@ -131,6 +129,31 @@ def main(opt):
 
 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('--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')
+    parser.add_argument('--n-cpus', type=int, default=0, help='number of cpu threads to use during batch generation')
+    parser.add_argument('--img-size', type=int, default=416, help='size of each image dimension')
+    opt = parser.parse_args()
+    print(opt, end='\n\n')
+
     init_seeds()
 
-    mAP = main(opt)
+    mAP = test(
+        opt.cfg,
+        opt.data_config_path,
+        opt.weights,
+        class_path=opt.class_path,
+        batch_size=opt.batch_size,
+        img_size=opt.img_size,
+        iou_thres=opt.iou_thres,
+        conf_thres=opt.conf_thres,
+        nms_thres=opt.nms_thres,
+        n_cpus=opt.n_cpus,
+    )
diff --git a/train.py b/train.py
index 103d6f03..5bc9cf9f 100644
--- a/train.py
+++ b/train.py
@@ -8,51 +8,48 @@ from utils.utils import *
 
 from utils import torch_utils
 
-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('-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')
-parser.add_argument('-resume', default=False, help='resume training flag')
-parser.add_argument('-batch_report', default=False, help='report TP, FP, FN, P and R per batch (slower)')
-parser.add_argument('-freeze_darknet53', default=False, help='freeze darknet53.conv.74 layers for first epoch')
-parser.add_argument('-var', type=float, default=0, help='optional test variable')
-opt = parser.parse_args()
-if opt.multi_scale:  # pass maximum multi_scale size
-    opt.img_size = 608
-print(opt)
-
 # Import test.py to get mAP after each epoch
-sys.argv[1:] = []  # delete any train.py command-line arguments before they reach test.py
-import test  # must follow sys.argv[1:] = []
+import test
 
 
-def main(opt):
+def train(
+    net_config_path,
+    data_config_path,
+    img_size=416,
+    resume=False,
+    epochs=100,
+    batch_size=16,
+    report=False,
+    multi_scale=False,
+    freeze_backbone=True,
+    var=0,
+):
 
     device = torch_utils.select_device()
     print("Using device: \"{}\"".format(device))
 
-    if not opt.multi_scale:
+    if not multi_scale:
         torch.backends.cudnn.benchmark = True
 
     os.makedirs('weights', exist_ok=True)
 
     # Configure run
-    data_config = parse_data_config(opt.data_config_path)
+    data_config = parse_data_config(data_config_path)
     num_classes = int(data_config['classes'])
     train_path = '../coco/trainvalno5k.txt'
 
     # Initialize model
-    model = Darknet(opt.cfg, opt.img_size)
+    model = Darknet(net_config_path, img_size)
 
     # Get dataloader
-    dataloader = load_images_and_labels(train_path, batch_size=opt.batch_size, img_size=opt.img_size,
-                                        multi_scale=opt.multi_scale, augment=True)
+    if multi_scale:  # pass maximum multi_scale size
+        img_size = 608
+
+    dataloader = load_images_and_labels(train_path, batch_size=batch_size, img_size=img_size,
+                                        multi_scale=multi_scale, augment=True)
 
     lr0 = 0.001
-    if opt.resume:
+    if resume:
         checkpoint = torch.load('weights/latest.pt', map_location='cpu')
 
         model.load_state_dict(checkpoint['model'])
@@ -103,7 +100,7 @@ def main(opt):
     mean_recall, mean_precision = 0, 0
     print('%11s' * 16 % (
         'Epoch', 'Batch', 'x', 'y', 'w', 'h', 'conf', 'cls', 'total', 'P', 'R', 'nTargets', 'TP', 'FP', 'FN', 'time'))
-    for epoch in range(opt.epochs):
+    for epoch in range(epochs):
         epoch += start_epoch
 
         # Update scheduler (automatic)
@@ -118,7 +115,7 @@ def main(opt):
             g['lr'] = lr
 
         # Freeze darknet53.conv.74 layers for first epoch
-        if opt.freeze_darknet53:
+        if freeze_backbone is not False:
             if epoch == 0:
                 for i, (name, p) in enumerate(model.named_parameters()):
                     if int(name.split('.')[1]) < 75:  # if layer < 75
@@ -143,7 +140,7 @@ def main(opt):
                     g['lr'] = lr
 
             # Compute loss, compute gradient, update parameters
-            loss = model(imgs.to(device), targets, batch_report=opt.batch_report, var=opt.var)
+            loss = model(imgs.to(device), targets, batch_report=report, var=var)
             loss.backward()
 
             # accumulated_batches = 1  # accumulate gradient for 4 batches before stepping optimizer
@@ -156,7 +153,7 @@ def main(opt):
             for key, val in model.losses.items():
                 rloss[key] = (rloss[key] * ui + val) / (ui + 1)
 
-            if opt.batch_report:
+            if report:
                 TP, FP, FN = metrics
                 metrics += model.losses['metrics']
 
@@ -173,7 +170,7 @@ def main(opt):
                     mean_recall = recall[k].mean()
 
             s = ('%11s%11s' + '%11.3g' * 14) % (
-                '%g/%g' % (epoch, opt.epochs - 1), '%g/%g' % (i, len(dataloader) - 1), rloss['x'],
+                '%g/%g' % (epoch, epochs - 1), '%g/%g' % (i, len(dataloader) - 1), rloss['x'],
                 rloss['y'], rloss['w'], rloss['h'], rloss['conf'], rloss['cls'],
                 rloss['loss'], mean_precision, mean_recall, model.losses['nT'], model.losses['TP'],
                 model.losses['FP'], model.losses['FN'], time.time() - t1)
@@ -201,8 +198,13 @@ def main(opt):
             os.system('cp weights/latest.pt weights/backup' + str(epoch) + '.pt')
 
         # Calculate mAP
-        test.opt.weights_path = 'weights/latest.pt'
-        mAP, R, P = test.main(test.opt)
+        mAP, R, P = test.test(
+            net_config_path,
+            data_config_path,
+            'weights/latest.pt',
+            batch_size=batch_size,
+            img_size=img_size,
+        )
 
         # Write epoch results
         with open('results.txt', 'a') as file:
@@ -215,7 +217,32 @@ def main(opt):
 
 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('--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')
+    parser.add_argument('--resume', default=False, help='resume training flag')
+    parser.add_argument('--report', default=False, help='report TP, FP, FN, P and R per batch (slower)')
+    parser.add_argument('--freeze-darknet53', default=False, help='freeze darknet53.conv.74 layers for first epoch')
+    parser.add_argument('--var', type=float, default=0, help='optional test variable')
+    opt = parser.parse_args()
+    print(opt, end='\n\n')
+
     init_seeds()
 
     torch.cuda.empty_cache()
-    main(opt)
+    train(
+        opt.cfg,
+        opt.data_config_path,
+        img_size=opt.img_size,
+        resume=opt.resume,
+        epochs=opt.epochs,
+        batch_size=opt.batch_size,
+        report=opt.report,
+        multi_scale=opt.multi_scale,
+        freeze_backbone=opt.freeze_darknet53,
+        var=opt.var,
+    )
diff --git a/utils/gcp.sh b/utils/gcp.sh
old mode 100644
new mode 100755
index 3e27c169..130f3249
--- a/utils/gcp.sh
+++ b/utils/gcp.sh
@@ -4,28 +4,28 @@
 sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3 && python3 train.py
 
 # Resume
-python3 train.py -resume 1
+python3 train.py --resume 1
 
 # Detect
 gsutil cp gs://ultralytics/yolov3.pt yolov3/weights
 python3 detect.py
 
 # Test
-python3 test.py -img_size 416 -weights_path weights/latest.pt
+python3 test.py --img_size 416 --weights weights/latest.pt
 
 # Test Darknet
-python3 test.py -img_size 416 -weights_path ../darknet/backup/yolov3.backup
+python3 test.py --img_size 416 --weights ../darknet/backup/yolov3.backup
 
 # Download and Test
 sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3
 wget https://pjreddie.com/media/files/yolov3.weights -P weights
-python3 test.py -img_size 416 -weights_path weights/backup5.pt -nms_thres 0.45
+python3 test.py --img_size 416 --weights weights/backup5.pt --nms_thres 0.45
 
 # Download and Resume
 sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3
 wget https://storage.googleapis.com/ultralytics/yolov3.pt -O weights/latest.pt
-python3 train.py -img_size 416 -batch_size 16 -epochs 1 -resume 1
-python3 test.py -img_size 416 -weights_path weights/latest.pt -conf_thres 0.5
+python3 train.py --img_size 416 --batch_size 16 --epochs 1 --resume 1
+python3 test.py --img_size 416 --weights weights/latest.pt --conf_thres 0.5
 
 # Copy latest.pt to bucket
 gsutil cp yolov3/weights/latest.pt gs://ultralytics
@@ -36,6 +36,6 @@ wget https://storage.googleapis.com/ultralytics/latest.pt
 
 # Testing
 sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3
-python3 train.py -epochs 3 -var 64
+python3 train.py --epochs 3 --var 64
 sudo shutdown
 

From b1fb6fa33d87d888c3ee9f38b0374a71841710cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Guillermo=20Garc=C3=ADa?= <ggarcia@redpoints.com>
Date: Wed, 5 Dec 2018 14:34:53 +0100
Subject: [PATCH 3/7] train.py resume argument as store_true

Default is false.

If want to resume, call train.py --resume
---
 README.md    | 2 +-
 train.py     | 2 +-
 utils/gcp.sh | 4 ++--
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index cbb37de9..fd529607 100755
--- a/README.md
+++ b/README.md
@@ -21,7 +21,7 @@ Python 3.7 or later with the following `pip3 install -U -r requirements.txt` pac
 
 **Start Training:** Run `train.py` to begin training after downloading COCO data with `data/get_coco_dataset.sh` and specifying COCO path on line 37 (local) or line 39 (cloud). Training runs about 1 hour per COCO epoch on a 1080 Ti.
 
-**Resume Training:** Run `train.py -resume 1` to resume training from the most recently saved checkpoint `latest.pt`.
+**Resume Training:** Run `train.py --resume` to resume training from the most recently saved checkpoint `latest.pt`.
 
 Each epoch trains on 120,000 images from the train and validate COCO sets, and tests on 5000 images from the COCO validate set. An Nvidia GTX 1080 Ti will process about 10-15 epochs/day depending on image size and augmentation (13 epochs/day at 416 pixels with default augmentation). Loss plots for the bounding boxes, objectness and class confidence should appear similar to results shown here (results in progress to 160 epochs, will update).
 
diff --git a/train.py b/train.py
index 5bc9cf9f..38f86035 100644
--- a/train.py
+++ b/train.py
@@ -224,7 +224,7 @@ if __name__ == '__main__':
     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')
-    parser.add_argument('--resume', default=False, help='resume training flag')
+    parser.add_argument('--resume', action='store_true', help='resume training flag')
     parser.add_argument('--report', default=False, help='report TP, FP, FN, P and R per batch (slower)')
     parser.add_argument('--freeze-darknet53', default=False, help='freeze darknet53.conv.74 layers for first epoch')
     parser.add_argument('--var', type=float, default=0, help='optional test variable')
diff --git a/utils/gcp.sh b/utils/gcp.sh
index 130f3249..32647259 100755
--- a/utils/gcp.sh
+++ b/utils/gcp.sh
@@ -4,7 +4,7 @@
 sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3 && python3 train.py
 
 # Resume
-python3 train.py --resume 1
+python3 train.py --resume
 
 # Detect
 gsutil cp gs://ultralytics/yolov3.pt yolov3/weights
@@ -24,7 +24,7 @@ python3 test.py --img_size 416 --weights weights/backup5.pt --nms_thres 0.45
 # Download and Resume
 sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3
 wget https://storage.googleapis.com/ultralytics/yolov3.pt -O weights/latest.pt
-python3 train.py --img_size 416 --batch_size 16 --epochs 1 --resume 1
+python3 train.py --img_size 416 --batch_size 16 --epochs 1 --resume
 python3 test.py --img_size 416 --weights weights/latest.pt --conf_thres 0.5
 
 # Copy latest.pt to bucket

From 89daa407e5a231d78f35bde36c2bfbef8546cdc5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Guillermo=20Garc=C3=ADa?= <ggarcia@redpoints.com>
Date: Wed, 5 Dec 2018 14:40:34 +0100
Subject: [PATCH 4/7] train.py report argument as store_true

Default is false: python train.py
If want the report: python train.py --report
---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index 38f86035..35f7d52b 100644
--- a/train.py
+++ b/train.py
@@ -225,7 +225,7 @@ if __name__ == '__main__':
     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')
     parser.add_argument('--resume', action='store_true', help='resume training flag')
-    parser.add_argument('--report', default=False, help='report TP, FP, FN, P and R per batch (slower)')
+    parser.add_argument('--report', action='store_true', help='report TP, FP, FN, P and R per batch (slower)')
     parser.add_argument('--freeze-darknet53', default=False, help='freeze darknet53.conv.74 layers for first epoch')
     parser.add_argument('--var', type=float, default=0, help='optional test variable')
     opt = parser.parse_args()

From 9c0c1f23abc26609d63ce0070929e0ee7e548668 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Guillermo=20Garc=C3=ADa?= <ggarcia@redpoints.com>
Date: Wed, 5 Dec 2018 14:52:40 +0100
Subject: [PATCH 5/7] scripts: use data config defined class names

Shorten name of --data-config-path argument to --data-config
---
 detect.py | 10 ++++++----
 test.py   | 14 +++++---------
 train.py  |  4 ++--
 3 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/detect.py b/detect.py
index a104e8a3..b1afa028 100755
--- a/detect.py
+++ b/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,
     )
diff --git a/test.py b/test.py
index 4630a0ee..edcdc9f1 100644
--- a/test.py
+++ b/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,
diff --git a/train.py b/train.py
index 35f7d52b..c8d0640d 100644
--- a/train.py
+++ b/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,

From 868a11675054334fa1601a96c2cf40e05429c842 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Guillermo=20Garc=C3=ADa?= <ggarcia@redpoints.com>
Date: Wed, 5 Dec 2018 15:27:48 +0100
Subject: [PATCH 6/7] train.py remove hardcoded weights/ path for weights.

If I want to store my weights in 'weights2' path:
python train.py --weights-path weights2

Default is the same: weights
---
 train.py | 42 ++++++++++++++++++++++++++++++++----------
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/train.py b/train.py
index c8d0640d..5b0ff69a 100644
--- a/train.py
+++ b/train.py
@@ -11,6 +11,11 @@ from utils import torch_utils
 # Import test.py to get mAP after each epoch
 import test
 
+DARKNET_WEIGHTS_FILENAME = 'darknet53.conv.74'
+DARKNET_WEIGHTS_URL = 'https://pjreddie.com/media/files/{}'.format(
+    DARKNET_WEIGHTS_FILENAME
+)
+
 
 def train(
     net_config_path,
@@ -19,6 +24,7 @@ def train(
     resume=False,
     epochs=100,
     batch_size=16,
+    weights_path='weights',
     report=False,
     multi_scale=False,
     freeze_backbone=True,
@@ -31,12 +37,14 @@ def train(
     if not multi_scale:
         torch.backends.cudnn.benchmark = True
 
-    os.makedirs('weights', exist_ok=True)
+    os.makedirs(weights_path, exist_ok=True)
+    latest_weights_file = os.path.join(weights_path, 'latest.pt')
+    best_weights_file = os.path.join(weights_path, 'best.pt')
 
     # Configure run
     data_config = parse_data_config(data_config_path)
     num_classes = int(data_config['classes'])
-    train_path = '../coco/trainvalno5k.txt'
+    train_path = data_config['train']
 
     # Initialize model
     model = Darknet(net_config_path, img_size)
@@ -50,7 +58,7 @@ def train(
 
     lr0 = 0.001
     if resume:
-        checkpoint = torch.load('weights/latest.pt', map_location='cpu')
+        checkpoint = torch.load(latest_weights_file, map_location='cpu')
 
         model.load_state_dict(checkpoint['model'])
         if torch.cuda.device_count() > 1:
@@ -79,9 +87,13 @@ def train(
         best_loss = float('inf')
 
         # Initialize model with darknet53 weights (optional)
-        if not os.path.isfile('weights/darknet53.conv.74'):
-            os.system('wget https://pjreddie.com/media/files/darknet53.conv.74 -P weights')
-        load_weights(model, 'weights/darknet53.conv.74')
+        def_weight_file = os.path.join(weights_path, DARKNET_WEIGHTS_FILENAME)
+        if not os.path.isfile(def_weight_file):
+            os.system('wget {} -P {}'.format(
+                DARKNET_WEIGHTS_URL,
+                weights_path))
+        assert os.path.isfile(def_weight_file)
+        load_weights(model, def_weight_file)
 
         if torch.cuda.device_count() > 1:
             raise Exception('Multi-GPU not currently supported: https://github.com/ultralytics/yolov3/issues/21')
@@ -187,21 +199,29 @@ def train(
                       'best_loss': best_loss,
                       'model': model.state_dict(),
                       'optimizer': optimizer.state_dict()}
-        torch.save(checkpoint, 'weights/latest.pt')
+        torch.save(checkpoint, latest_weights_file)
 
         # Save best checkpoint
         if best_loss == loss_per_target:
-            os.system('cp weights/latest.pt weights/best.pt')
+            os.system('cp {} {}'.format(
+                latest_weights_file,
+                best_weights_file,
+            ))
 
         # Save backup weights every 5 epochs
         if (epoch > 0) & (epoch % 5 == 0):
-            os.system('cp weights/latest.pt weights/backup' + str(epoch) + '.pt')
+            backup_file_name = 'backup{}.pt'.format(epoch)
+            backup_file_path = os.path.join(weights_path, backup_file_name)
+            os.system('cp {} {}'.format(
+                latest_weights_file,
+                backup_file_path,
+            ))
 
         # Calculate mAP
         mAP, R, P = test.test(
             net_config_path,
             data_config_path,
-            'weights/latest.pt',
+            latest_weights_file,
             batch_size=batch_size,
             img_size=img_size,
         )
@@ -224,6 +244,7 @@ if __name__ == '__main__':
     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')
+    parser.add_argument('--weights-path', type=str, default='weights', help='path to store weights')
     parser.add_argument('--resume', action='store_true', help='resume training flag')
     parser.add_argument('--report', action='store_true', help='report TP, FP, FN, P and R per batch (slower)')
     parser.add_argument('--freeze-darknet53', default=False, help='freeze darknet53.conv.74 layers for first epoch')
@@ -241,6 +262,7 @@ if __name__ == '__main__':
         resume=opt.resume,
         epochs=opt.epochs,
         batch_size=opt.batch_size,
+        weights_path=opt.weights_path,
         report=opt.report,
         multi_scale=opt.multi_scale,
         freeze_backbone=opt.freeze_darknet53,

From d03ce45da5e83527c6225fecd0b17d96887fdf6c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Guillermo=20Garc=C3=ADa?= <ggarcia@redpoints.com>
Date: Wed, 5 Dec 2018 15:55:09 +0100
Subject: [PATCH 7/7] train.py freeze-darknet53 shortened to freeze and action
 store_true

Traing with freeze: python train.py --freeze
Train without freeze: python train.py

Note: in the actual version freeze is only for first epoche
---
 train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index 5b0ff69a..cfe90dea 100644
--- a/train.py
+++ b/train.py
@@ -247,7 +247,7 @@ if __name__ == '__main__':
     parser.add_argument('--weights-path', type=str, default='weights', help='path to store weights')
     parser.add_argument('--resume', action='store_true', help='resume training flag')
     parser.add_argument('--report', action='store_true', help='report TP, FP, FN, P and R per batch (slower)')
-    parser.add_argument('--freeze-darknet53', default=False, help='freeze darknet53.conv.74 layers for first epoch')
+    parser.add_argument('--freeze', action='store_true', help='freeze darknet53.conv.74 layers for first epoche')
     parser.add_argument('--var', type=float, default=0, help='optional test variable')
     opt = parser.parse_args()
     print(opt, end='\n\n')
@@ -265,6 +265,6 @@ if __name__ == '__main__':
         weights_path=opt.weights_path,
         report=opt.report,
         multi_scale=opt.multi_scale,
-        freeze_backbone=opt.freeze_darknet53,
+        freeze_backbone=opt.freeze,
         var=opt.var,
     )