from graphviz import Digraph
import yaml

# 加载 YAML 文件
with open('/home/lab/LJ/wampee/ultralytics/ultralytics/cfg/models/v8/yolov8.yaml', 'r') as f:
    model_config = yaml.safe_load(f)

# 创建一个有向图
dot = Digraph(comment='YOLOv8 Model Structure')

# 添加 Backbone 节点
dot.node('Input', 'Input Image')
for idx, layer in enumerate(model_config['backbone']):
    dot.node(f'Backbone_{idx}', f"{layer[2]} ({layer[3]})")
    dot.edge('Input' if idx == 0 else f'Backbone_{idx-1}', f'Backbone_{idx}')

# 添加 Head 节点
for idx, layer in enumerate(model_config['head']):
    dot.node(f'Head_{idx}', f"{layer[2]} ({layer[3]})")
    dot.edge(f'Backbone_{len(model_config["backbone"])-1}' if idx == 0 else f'Head_{idx-1}', f'Head_{idx}')

# 保存和展示结构图
dot.render('yolov8_model_structure', format='png', cleanup=True)
dot.view()