diff --git a/detect.py b/detect.py index 7daa87436daa..c51c6fa4e0c1 100644 --- a/detect.py +++ b/detect.py @@ -172,7 +172,7 @@ def detect(weights='yolov5s.pt', # model.pt path(s) print(f'Done. ({time.time() - t0:.3f}s)') -if __name__ == '__main__': +def parse_opt(): parser = argparse.ArgumentParser() parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') parser.add_argument('--source', type=str, default='data/images', help='file/dir/URL/glob, 0 for webcam') @@ -198,7 +198,15 @@ def detect(weights='yolov5s.pt', # model.pt path(s) parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences') parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference') opt = parser.parse_args() + return opt + + +def main(opt): print(opt) check_requirements(exclude=('tensorboard', 'thop')) - detect(**vars(opt)) + + +if __name__ == "__main__": + opt = parse_opt() + main(opt) diff --git a/models/export.py b/models/export.py index 3c04b07fdc95..8c491dabddc0 100644 --- a/models/export.py +++ b/models/export.py @@ -144,7 +144,7 @@ def export(weights='./yolov5s.pt', # weights path print(f'\nExport complete ({time.time() - t:.2f}s). Visualize with https://github.com/lutzroeder/netron.') -if __name__ == '__main__': +def parse_opt(): parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image (height, width)') @@ -159,7 +159,15 @@ def export(weights='./yolov5s.pt', # weights path parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') parser.add_argument('--opset-version', type=int, default=12, help='ONNX: opset version') opt = parser.parse_args() + return opt + + +def main(opt): print(opt) set_logging() - export(**vars(opt)) + + +if __name__ == "__main__": + opt = parse_opt() + main(opt) diff --git a/test.py b/test.py index a38298da54da..5ebfb36509ea 100644 --- a/test.py +++ b/test.py @@ -294,7 +294,7 @@ def test(data, return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t -if __name__ == '__main__': +def parse_opt(): parser = argparse.ArgumentParser(prog='test.py') parser.add_argument('--data', type=str, default='data/coco128.yaml', help='dataset.yaml path') parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') @@ -319,6 +319,10 @@ def test(data, opt.save_json |= opt.data.endswith('coco.yaml') opt.save_txt |= opt.save_hybrid opt.data = check_file(opt.data) # check file + return opt + + +def main(opt): print(opt) check_requirements(exclude=('tensorboard', 'thop')) @@ -344,3 +348,8 @@ def test(data, np.savetxt(f, y, fmt='%10.4g') # save os.system('zip -r study.zip study_*.txt') plot_study_txt(x=x) # plot + + +if __name__ == "__main__": + opt = parse_opt() + main(opt) diff --git a/train.py b/train.py index 3eb866345d47..8056183242a6 100644 --- a/train.py +++ b/train.py @@ -463,7 +463,7 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary return results -if __name__ == '__main__': +def parse_opt(): parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path') parser.add_argument('--cfg', type=str, default='', help='model.yaml path') @@ -504,6 +504,11 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary # Set DDP variables opt.world_size = int(getattr(os.environ, 'WORLD_SIZE', 1)) opt.global_rank = int(getattr(os.environ, 'RANK', -1)) + return opt + + +def main(opt): + print(opt) set_logging(opt.global_rank) if opt.global_rank in [-1, 0]: check_git_status() @@ -628,3 +633,8 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary plot_evolution(yaml_file) print(f'Hyperparameter evolution complete. Best results saved as: {yaml_file}\n' f'Command to train a new model with these hyperparameters: $ python train.py --hyp {yaml_file}') + + +if __name__ == "__main__": + opt = parse_opt() + main(opt)