Chengfang Lu
e72bc061c5
Major changes: - New frontend (platform/web/): Vite + React 18 + TypeScript + Tailwind - 4-module navigation: 数据送标 / 模型管理 / 车队管理 / 系统管理 - Data catalog with charts (DMS/ADAS/Lane 3-tab view) - Quality review workflow (标注质检): Good/Fine/Bad scoring with auto-advance - Audit enhancements: batch operations, rejection categories, Feishu notifications - Operation audit log (操作日志) - World model simulation studio (仿真工坊) - Dataset version management with snapshots and diff - ADAS 7-class dataset integration (138K images organized + compressed) - User management with Feishu integration and pagination - CRUD/search/filter on all pages, card layout redesign - PIL-optimized image overlay rendering - Auto-snapshot on build, in_review workflow stage - Removed embedded algorithm code (now in workspace)
100 lines
3.0 KiB
Python
Executable File
100 lines
3.0 KiB
Python
Executable File
"""Export UFLD .pth checkpoint to ONNX (matches tusimple_res18_4lane_v1 / best.pth)."""
|
|
|
|
import argparse
|
|
import os
|
|
|
|
import torch
|
|
|
|
from model.model import parsingNet
|
|
from utils.common import checkpoint_state_dict
|
|
from utils.config import Config
|
|
|
|
|
|
def export_onnx(
|
|
model_path,
|
|
output_path=None,
|
|
backbone='18',
|
|
griding_num=100,
|
|
num_lanes=4,
|
|
opset=11,
|
|
):
|
|
if output_path is None:
|
|
base, _ = os.path.splitext(model_path)
|
|
output_path = base + '.onnx'
|
|
|
|
cls_num_per_lane = 56 # TuSimple row anchors
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
net = parsingNet(
|
|
pretrained=False,
|
|
backbone=backbone,
|
|
cls_dim=(griding_num + 1, cls_num_per_lane, num_lanes),
|
|
use_aux=False,
|
|
).to(device)
|
|
net.load_state_dict(checkpoint_state_dict(model_path, map_location=device), strict=False)
|
|
net.eval()
|
|
|
|
dummy = torch.randn(1, 3, 288, 800, device=device)
|
|
os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
|
|
torch.onnx.export(
|
|
net,
|
|
dummy,
|
|
output_path,
|
|
opset_version=opset,
|
|
input_names=['input'],
|
|
output_names=['output'],
|
|
dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
|
|
dynamo=False, # legacy exporter embeds weights (PyTorch 2.9+ default may omit)
|
|
)
|
|
print('saved:', os.path.abspath(output_path))
|
|
print('input: [1, 3, 288, 800]')
|
|
print('output:', tuple(net(dummy).shape))
|
|
return output_path
|
|
|
|
|
|
def main():
|
|
default_pth = 'log/20250702_165153_lr_1e-05_b_32_ufld_2lanes_res18/best.pth'
|
|
parser = argparse.ArgumentParser(description='Export UFLD checkpoint to ONNX')
|
|
parser.add_argument('--config', default=None, help='optional py config (overrides defaults)')
|
|
parser.add_argument('--model_path', default=default_pth, help='path to .pth')
|
|
parser.add_argument('--output', default=None, help='output .onnx path')
|
|
parser.add_argument('--backbone', default=None)
|
|
parser.add_argument('--griding_num', type=int, default=None)
|
|
parser.add_argument('--num_lanes', type=int, default=None)
|
|
parser.add_argument('--opset', type=int, default=11)
|
|
args = parser.parse_args()
|
|
|
|
backbone = '18'
|
|
griding_num = 100
|
|
num_lanes = 4
|
|
model_path = args.model_path
|
|
output = args.output
|
|
|
|
if args.config:
|
|
cfg = Config.fromfile(args.config)
|
|
backbone = cfg.backbone
|
|
griding_num = cfg.griding_num
|
|
num_lanes = cfg.num_lanes
|
|
if getattr(cfg, 'test_model', None):
|
|
model_path = cfg.test_model
|
|
|
|
if args.backbone is not None:
|
|
backbone = args.backbone
|
|
if args.griding_num is not None:
|
|
griding_num = args.griding_num
|
|
if args.num_lanes is not None:
|
|
num_lanes = args.num_lanes
|
|
|
|
export_onnx(
|
|
model_path,
|
|
output_path=output,
|
|
backbone=backbone,
|
|
griding_num=griding_num,
|
|
num_lanes=num_lanes,
|
|
opset=args.opset,
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|