decode_recognition for vintext model
Closed this issue · 2 comments
superkido511 commented
Hello, could you provide the ctc decode logic for the pretrained vintext model? It seems to have only 105 classes.
I tried the decode function from SwinTextSpotter repo but the results are wrong
https://github.com/mxin262/SwinTextSpotter/blob/main/detectron2/utils/visualizer_vintext.py#L671C26-L671C26
import os, sys
import torch
import numpy as np
import cv2
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
import pickle
# with open('chn_cls_list.txt', 'rb') as fp:
# CTLABELS = pickle.load(fp)
# def _decode_recognition(rec):
# s = ''
# for c in rec:
# c = int(c)
# if c < 5461:
# s += str(chr(CTLABELS[c]))
# elif c == 5462:
# s += u''
# return s
# 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
dictionary = "aàáạảãâầấậẩẫăằắặẳẵAÀÁẠẢÃĂẰẮẶẲẴÂẦẤẬẨẪeèéẹẻẽêềếệểễEÈÉẸẺẼÊỀẾỆỂỄoòóọỏõôồốộổỗơờớợởỡOÒÓỌỎÕÔỒỐỘỔỖƠỜỚỢỞỠiìíịỉĩIÌÍỊỈĨuùúụủũưừứựửữƯỪỨỰỬỮUÙÚỤỦŨyỳýỵỷỹYỲÝỴỶỸ"
def make_groups():
groups = []
i = 0
while i < len(dictionary) - 5:
group = [c for c in dictionary[i : i + 6]]
i += 6
groups.append(group)
return groups
groups = make_groups()
TONES = ["", "ˋ", "ˊ", "﹒", "ˀ", "˜"]
SOURCES = ["ă", "â", "Ă", "Â", "ê", "Ê", "ô", "ơ", "Ô", "Ơ", "ư", "Ư", "Đ", "đ"]
TARGETS = ["aˇ", "aˆ", "Aˇ", "Aˆ", "eˆ", "Eˆ", "oˆ", "o˒", "Oˆ", "O˒", "u˒", "U˒", "D-", "d‑"]
def ctc_decode_recognition(rec):
# CTLABELS = "_0123456789abcdefghijklmnopqrstuvwxyz"
# 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','{','|','}','~']
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",
"{",
"|",
"}",
"~",
"ˋ",
"ˊ",
"﹒",
"ˀ",
"˜",
"ˇ",
"ˆ",
"˒",
"‑",
]
# ctc decoding
last_char = False
s = ''
for c in rec:
c = int(c)
if 0<c < 104:# 107:
s += CTLABELS[c-1]
last_char = c
elif c == 0:
s += u''
else:
last_char = False
if len(s) == 0:
s = ' '
s = decoder(s)
return s
def correct_tone_position(word):
word = word[:-1]
if len(word) < 2:
pass
first_ord_char = ""
second_order_char = ""
for char in word:
for group in groups:
if char in group:
second_order_char = first_ord_char
first_ord_char = group[0]
if word[-1] == first_ord_char and second_order_char != "":
pair_chars = ["qu", "Qu", "qU", "QU", "gi", "Gi", "gI", "GI"]
for pair in pair_chars:
if pair in word and second_order_char in ["u", "U", "i", "I"]:
return first_ord_char
return second_order_char
return first_ord_char
def decoder(recognition):
for char in TARGETS:
recognition = recognition.replace(char, SOURCES[TARGETS.index(char)])
if len(recognition) < 1:
return recognition
if recognition[-1] in TONES:
if len(recognition) < 2:
return recognition
replace_char = correct_tone_position(recognition)
tone = recognition[-1]
recognition = recognition[:-1]
for group in groups:
if replace_char in group:
recognition = recognition.replace(replace_char, group[TONES.index(tone)])
return recognition
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 = [ctc_decode_recognition(rrec) for rrec in output['rec']]
When I use the decode function here
https://github.com/mxin262/ESTextSpotter/blob/main/vis.py#L13
The result seems correct but missing diacritical marks
mxin262 commented
superkido511 commented
Thank you