Human-Falling-Detect/export_fastpose_onnx.py at main · zzy7569/Human-Falling-Detect · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
import torch
import sys
from SPPE.src.models.FastPose import FastPose
from SPPE.src.opt import opt

def export_fastpose_onnx(backbone, weight_path, onnx_path, input_height, input_width, device='cpu'):
    model = FastPose(backbone, num_join=17).to(device)
    model.load_state_dict(torch.load(weight_path, map_location=device))
    model.eval()
    dummy_input = torch.randn(1, 3, input_height, input_width, device=device)
    torch.onnx.export(
        model, dummy_input, onnx_path,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
        opset_version=12
    )
    print(f"FastPose exported to {onnx_path}")

if __name__ == '__main__':
    # 导出resnet50
    export_fastpose_onnx(
        backbone='resnet50',
        weight_path='Models/sppe/fast_res50_256x192.pth',
        onnx_path='Models/sppe/fast_res50_256x192.onnx',
        input_height=256,
        input_width=192
    )

    # 导出resnet101
    # export_fastpose_onnx(
    #     backbone='resnet101',
    #     weight_path='Models/sppe/fast_res101_320x256.pth',
    #     onnx_path='Models/sppe/fast_res101_320x256.onnx',
    #     input_height=320,
    #     input_width=256
    # )