Shape mismatch error when loading Vintext model

Hello, I'm trying to do inference with pretrained Vintext checkpoint

Here is my code that I picked up from vis.py

import os, sys
import torch
import numpy as np

from models.ests import build_ests
from util.slconfig import SLConfig
from util.visualizer import COCOVisualizer
from util import box_ops
from PIL import Image
import datasets.transforms as T

CTLABELS = [' ','!','"','#','$','%','&','\'','(',')','*','+',',','-','.','/','0','1','2','3','4','5','6','7','8','9',':',';','<','=','>','?','@','A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z','[','\\',']','^','_','`','a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z','{','|','}','~']
def _decode_recognition(rec):
    s = ''
    rec = rec.tolist()
    for c in rec:
        if c>94:
            continue
        s += CTLABELS[c]
    return s

def build_model_main(args):
    # we use register to maintain models from catdet6 on.
    from models.registry import MODULE_BUILD_FUNCS
    assert args.modelname in MODULE_BUILD_FUNCS._module_dict
    build_func = MODULE_BUILD_FUNCS.get(args.modelname)
    args.device = 'cuda'
    model, criterion, postprocessors = build_func(args)
    return model, criterion, postprocessors

model_config_path = "config/ESTS/ESTS_5scale_vintext_finetune.py" # change the path of the model config file
model_checkpoint_path = "vintext_checkpoint.pth" # change the path of the model checkpoint

args = SLConfig.fromfile(model_config_path) 
model, criterion, postprocessors = build_model_main(args)
checkpoint = torch.load(model_checkpoint_path, map_location='cpu')
model.load_state_dict(checkpoint['model'])
model.eval()
model.cuda()
transform = T.Compose([
    T.RandomResize([800],max_size=1333),
    T.ToTensor(),
    T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])]
)

img_path = 'test.jpg'
image = Image.open(img_path).convert('RGB')
image, _ = transform(image,None)
output = model(image[None].cuda())
output = postprocessors['bbox'](output, torch.Tensor([[1.0, 1.0]]))[0]
rec = [_decode_recognition(i) for i in output['rec']]
thershold = 0.4 # set a thershold
scores = output['scores']
labels = output['labels']
boxes = box_ops.box_xyxy_to_cxcywh(output['boxes'])
select_mask = scores > thershold
recs = []
for i,r in zip(select_mask,rec):
    if i:
        recs.append(r)
# vslzr = COCOVisualizer()
# box_label = ['text' for item in rec[select_mask]]
pred_dict = {
    'boxes': boxes[select_mask],
    'size': torch.tensor([image.shape[1],image.shape[2]]),
    'box_label': recs,
    'image_id': idx,
    'beziers': output['beziers'][select_mask]
}

However, I got shape mismatch error at this line output = model(image[None].cuda()) shape mismatch: value tensor of shape [75, 256, 8, 25] cannot be broadcast to indexing result of shape [75, 256, 8, 8]

mxin262 / ESTextSpotter

Shape mismatch error when loading Vintext model #1