Spaces:
Sleeping
Sleeping
File size: 2,318 Bytes
ac8579b |
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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
from cv2 import imshow
from matplotlib import lines
import numpy as np
import onnxruntime
import cv2
import torch
import onnx
from basemodel import TextDetBase
import onnxsim
from models.yolov5.common import Conv
from models.yolov5.yolo import Detect
import torch.nn as nn
import time
from seg_dataset import letterbox
from utils.yolov5_utils import fuse_conv_and_bn
class SiLU(nn.Module): # export-friendly version of nn.SiLU()
@staticmethod
def forward(x):
return x * torch.sigmoid(x)
def concate_models(blk_weights, seg_weights, det_weights, save_path):
textdetector_dict = dict()
textdetector_dict['blk_det'] = torch.load(blk_weights, map_location='cpu')
textdetector_dict['text_seg'] = torch.load(seg_weights, map_location='cpu')['weights']
textdetector_dict['text_det'] = torch.load(det_weights, map_location='cpu')['weights']
torch.save(textdetector_dict, save_path)
def export_onnx(model, im, file, opset, train=False, simplify=True, dynamic=False, inplace=False):
# YOLOv5 ONNX export
f = file + '.onnx'
for k, m in model.named_modules():
if isinstance(m, Conv): # assign export-friendly activations
if isinstance(m.act, nn.SiLU):
m.act = SiLU()
elif isinstance(m, Detect):
m.inplace = inplace
m.onnx_dynamic = False
torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
do_constant_folding=not train,
input_names=['images'],
output_names=['blk', 'seg', 'det'],
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640)
'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
} if dynamic else None)
# Checks
model_onnx = onnx.load(f) # load onnx model
onnx.checker.check_model(model_onnx) # check onnx model
model_onnx, check = onnxsim.simplify(
model_onnx,
dynamic_input_shape=dynamic,
input_shapes={'images': list(im.shape)} if dynamic else None)
assert check, 'assert check failed'
onnx.save(model_onnx, f) |