"""
NanoDetHeatV4 权重导出: PyTorch. pth → 二进制
"""
import sys, struct, os
sys.path.insert(0, os.path.dirname(__file__))
import torch
import numpy as np
from model_heat_v4 import NanoDetHeatV4

def export(pth_path, out_path):
    model = NanoDetHeatV4(num_classes=3)
    ckpt = torch.load(pth_path, map_location='cpu', weights_only=False)
    if 'model' in ckpt:
        ckpt = ckpt['model']
    model.load_state_dict(ckpt, strict=False)
    model.eval()

    layers = {}
    for name, param in model.named_parameters():
        layers[name] = param.detach().cpu().numpy().astype(np.float32)
    for name, buf in model.named_buffers():
        layers[name] = buf.detach().cpu().numpy().astype(np.float32)

    total = sum(v.size for v in layers.values())
    print(f"层数: {len(layers)}, 总参数: {total:,} ≈ {total/1000:.1f}K")

    for k, v in sorted(layers.items()):
        print(f"  {k:50s}  {list(v.shape)}  {v.size:,}")

    with open(out_path, 'wb') as f:
        f.write(struct.pack('I', len(layers)))
        for name in sorted(layers.keys()):
            data = layers[name]
            nb = name.encode('utf-8')
            f.write(struct.pack('I', len(nb)))
            f.write(nb)
            f.write(struct.pack('I', data.ndim))
            f.write(struct.pack('i' * data.ndim, *data.shape))
            f.write(data.tobytes())

    print(f"\n导出: {out_path}")

    # 测试推理
    x = torch.randn(1, 3, 120, 160)
    with torch.no_grad():
        y = model(x)
    print(f"输入: {x.shape} → 输出: {y.shape}")

if __name__ == '__main__':
    base = os.path.dirname(os.path.abspath(__file__))
    export(os.path.join(base, 'best.pth.1'), os.path.join(base, 'nanodet.bin'))