Slightly modify CLI execution (#3687)

* Slightly modify CLI execution

This simple change makes it easier to run the primary functions of this
repo (train/detect/test) from within Python. An object which represents
`opt` can be constructed and fed to the `main` function of each of these
modules, rather than having to call the lower level functions directly,
or run the module as a script.

* Update export.py

Add CLI parsing update for more convenient module usage within Python.

Co-authored-by: Lewis Belcher <lb@desupervised.io>
This commit is contained in:
lb-desupervisedandLewis Belcher authored and GitHub committed 2021-06-19 12:06:59 +02:00
1 parent bf209f6fe9
commit bfb2276b1d
4 files changed
+41 -6

No files matched your search

+11 -1
View File
@@ -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 @@ if __name__ == '__main__':
# 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 @@ if __name__ == '__main__':
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)