#!/usr/bin/env python
# -*- coding:utf-8 -*-

# Created by zhenqin.
# User: zhenqin
# Date: 2024-04-25
# Time: 下午1:58
# Verdor: yiidata.com
#

# 针对有网络模型，但还没有训练保存 .pth 文件的情况
import onnxruntime
import netron
import onnx
import torch.onnx

def printLabels():
    # Load Label from Model's meta data
    session = onnxruntime.InferenceSession("./yolov8s.onnx", providers=['CPUExecutionProvider'])
    meta = session.get_modelmeta().custom_metadata_map  # metadata
    label_names = eval(meta['names'])
    print(label_names)

def addLabels():
    f = "./yolov7.onnx"
    new_f = "./yolov7-425-labels.onnx"
    model_onnx = onnx.load(f)  # load onnx model
    onnx.checker.check_model(model_onnx)  # check onnx model
    labels = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
                "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
                "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
                "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
                "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
                "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
                "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
                "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
                "hair drier", "toothbrush"]
    index = 1
    meta = model_onnx.metadata_props.add()
    meta.key, meta.value = "labels", ",".join(labels)
    print(meta.key + ":" + meta.value)
    onnx.save(model_onnx, new_f)

if __name__ == '__main__':
    # printLabels()
    # netron.start("./yolov8s.onnx")  # 输出网络结构
    netron.start("./yolov7-425-labels.onnx")
    # addLabels()