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,