This commit is contained in:
Glenn Jocher 2019-02-11 22:44:25 +01:00
parent 742908257a
commit 9f145d2aa7
2 changed files with 8 additions and 9 deletions

12
test.py
View File

@ -20,9 +20,9 @@ def test(
device = torch_utils.select_device() device = torch_utils.select_device()
# Configure run # Configure run
data_cfg = parse_data_cfg(data_cfg) data_cfg_dict = parse_data_cfg(data_cfg)
nC = int(data_cfg['classes']) # number of classes (80 for COCO) nC = int(data_cfg_dict['classes']) # number of classes (80 for COCO)
test_path = data_cfg['valid'] test_path = data_cfg_dict['valid']
# Initialize model # Initialize model
model = Darknet(cfg, img_size) model = Darknet(cfg, img_size)
@ -111,7 +111,7 @@ def test(
# Print mAP per class # Print mAP per class
print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP') + '\n\nmAP Per Class:') print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP') + '\n\nmAP Per Class:')
classes = load_classes(data_cfg['names']) # Extracts class labels from file classes = load_classes(data_cfg_dict['names']) # Extracts class labels from file
for i, c in enumerate(classes): for i, c in enumerate(classes):
print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i])) print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i]))
@ -122,8 +122,8 @@ def test(
if __name__ == '__main__': if __name__ == '__main__':
parser = argparse.ArgumentParser(prog='test.py') parser = argparse.ArgumentParser(prog='test.py')
parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch') 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('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
parser.add_argument('--data-cfg', type=str, default='cfg/coco.data', help='path to data config file') parser.add_argument('--data-cfg', type=str, default='cfg/coco.data', help='coco.data file path')
parser.add_argument('--weights', type=str, default='weights/yolov3.pt', help='path to weights file') parser.add_argument('--weights', type=str, default='weights/yolov3.pt', help='path to weights file')
parser.add_argument('--iou-thres', type=float, default=0.5, help='iou threshold required to qualify as detected') 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('--conf-thres', type=float, default=0.3, help='object confidence threshold')

View File

@ -35,8 +35,7 @@ def train(
best = os.path.join(weights, 'best.pt') best = os.path.join(weights, 'best.pt')
# Configure run # Configure run
data_cfg = parse_data_cfg(data_cfg) train_path = parse_data_cfg(data_cfg)['train']
train_path = data_cfg['train']
# Initialize model # Initialize model
model = Darknet(cfg, img_size) model = Darknet(cfg, img_size)
@ -187,8 +186,8 @@ if __name__ == '__main__':
parser.add_argument('--epochs', type=int, default=100, help='number of epochs') 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('--batch-size', type=int, default=16, help='size of each image batch')
parser.add_argument('--accumulated-batches', type=int, default=1, help='number of batches before optimizer step') parser.add_argument('--accumulated-batches', type=int, default=1, help='number of batches before optimizer step')
parser.add_argument('--data-cfg', 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('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
parser.add_argument('--data-cfg', type=str, default='cfg/coco.data', help='coco.data file path')
parser.add_argument('--multi-scale', action='store_true', help='random image sizes per batch 320 - 608') parser.add_argument('--multi-scale', action='store_true', help='random image sizes per batch 320 - 608')
parser.add_argument('--img-size', type=int, default=32 * 13, help='pixels') parser.add_argument('--img-size', type=int, default=32 * 13, help='pixels')
parser.add_argument('--weights', type=str, default='weights', help='path to store weights') parser.add_argument('--weights', type=str, default='weights', help='path to store weights')