请问yolov5的主函数如何调用？

Question

我想把detect.py的

if __name__ == '__main__':

parser = argparse.ArgumentParser()
parser.add_argument('--weights', nargs='+', type=str, default='./weights/helmet_head_person_l.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str, default=f'inference/images', help='source' )  # file/folder, 0 for webcam
parser.add_argument('--output', type=str, default='inference/output', help='output folder')  # output folder
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--view-img', action='store_true', help='display results')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
parser.add_argument('--augment', action='store_true', help='augmented inference')
parser.add_argument('--update', action='store_true', help='update all models')
opt = parser.parse_args()
print(opt)

with torch.no_grad():
    if opt.update:  # update all models (to fix SourceChangeWarning)
        for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov3-spp.pt']:
            detect()
            create_pretrained(opt.weights, opt.weights)
    else:
        detect()

封装成一个函数，以便其他py文件调用默认就是python detect.py --source 0，
多谢了！！！

Answer 1

将detect.py这个脚本文件的if __name__ == '__main__'之前，定义一个函数，比如your_func，函数体就是if __name__ == '__main__'之后的内容。其他脚本文件这样调用你定义的函数：

 

import detect
detect.your_func()

或者，这样写也可以：

 

from detect import your_func
your_func()