import os
import sys
from yolo import YOLO

def main(pt_path, onnx_path):
    if os.environ.get('YOLOV3_ROOT') is None:
        raise EnvironmentError(
            "YOLOV3_ROOT environment variable is not set. \
            Please export the path to your YOLOv3 project, e.g.: export YOLOV3_ROOT=/path/to/yolov3")
    classes_path = os.path.join(
        os.environ.get('YOLOV3_ROOT'), 'model_data/coco_classes.txt'
    )
    anchors_path = os.path.join(
        os.environ.get('YOLOV3_ROOT'), 'model_data/yolo_anchors.txt'
    )
    yolo = YOLO(model_path=pt_path, classes_path=classes_path, anchors_path=anchors_path)
    yolo.convert_to_onnx(True, onnx_path)

if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])