from craft import CRAFT
import torch
from collections import OrderedDict
import imgproc
import numpy as np
import cv2

def copyStateDict(state_dict):
    if list(state_dict.keys())[0].startswith("module"):
        start_idx = 1
    else:
        start_idx = 0
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = ".".join(k.split(".")[start_idx:])
        new_state_dict[name] = v
    return new_state_dict

# 构建输入
input_data = np.random.rand(1, 3, 736, 1280).astype("float32")
net = CRAFT()
net.load_state_dict(copyStateDict(torch.load('craft_mlt_25k.pth')))
net = net.cuda()
net.eval()

# 进行转换
from x2paddle.convert import pytorch2paddle
pytorch2paddle(net, 
          save_dir="paddlemodel", 
          jit_type="trace", 
          input_examples=[torch.tensor(input_data).cuda()])