def trainBatch(): # 获取一个batch的数据 [images,label] data = train_iter.next() if opt.BidirDecoder: cpu_images, cpu_texts, cpu_texts_rev = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) t_rev, _ = converter.encode(cpu_texts_rev, scanned=True) utils.loadData(text, t) utils.loadData(text_rev, t_rev) utils.loadData(length, l) # 双向lstm有两个结果 preds0, preds1 = MORAN(image, length, text, text_rev) cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) else: cpu_images, cpu_texts = data utils.loadData(image, cpu_images) # 标签和每个标签的长度 t, l = converter.encode(cpu_texts, scanned=True) utils.loadData(text, t) utils.loadData(length, l) # 单向lstm一个结果 preds = MORAN(image, length, text, text_rev) cost = criterion(preds, text) MORAN.zero_grad() cost.backward() optimizer.step() return cost
def __init__(self, model_path): alphabet = '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:$' self.cuda_flag = torch.cuda.is_available() if self.cuda_flag: self.MORAN = MORAN(1, len(alphabet.split(':')), 256, 32, 100, BidirDecoder=True, CUDA=self.cuda_flag) self.MORAN = self.MORAN.cuda() else: self.MORAN = MORAN(1, len(alphabet.split(':')), 256, 32, 100, BidirDecoder=True, inputDataType='torch.FloatTensor', CUDA=self.cuda_flag) print('loading pretrained model from %s' % model_path) if self.cuda_flag: state_dict = torch.load(model_path) else: state_dict = torch.load(model_path, map_location='cpu') MORAN_state_dict_rename = OrderedDict() for k, v in state_dict.items(): name = k.replace("module.", "") # remove `module.` MORAN_state_dict_rename[name] = v self.MORAN.load_state_dict(MORAN_state_dict_rename) for p in self.MORAN.parameters(): p.requires_grad = False self.MORAN.eval() self.converter = utils.strLabelConverterForAttention(alphabet, ':') self.transformer = dataset.resizeNormalize((100, 32))
def trainBatch(steps): data = train_iter.next() if opt.BidirDecoder: cpu_images, cpu_texts, cpu_texts_rev = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) t_rev, _ = converter.encode(cpu_texts_rev, scanned=True) utils.loadData(text, t) utils.loadData(text_rev, t_rev) utils.loadData(length, l) preds0, preds1 = MORAN(image, length, text, text_rev) cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) else: cpu_images, cpu_texts = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) utils.loadData(text, t) utils.loadData(length, l) preds = MORAN(image, length, text, text_rev) cost = criterion(preds, text) MORAN.zero_grad() cost.backward() # 反向传播 optimizer.step() # 优化器 return cost
def trainBatch(): data = train_iter.next() if opt.BidirDecoder: cpu_images, cpu_texts, cpu_texts_rev = data #读取标签数据 utils.loadData(image, cpu_images) #将图像数据赋值给image t, l = converter.encode(cpu_texts, scanned=True) #将文本编码成类别标签 t_rev, _ = converter.encode(cpu_texts_rev, scanned=True) #将反向文本编码成类别标签 utils.loadData(text, t) #将正向文本标签赋予t utils.loadData(text_rev, t_rev) #将反向文本标签赋予t_rev utils.loadData(length, l) preds0, preds1 = MORAN(image, length, text, text_rev) #输出正向和反向识别概率 cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) #计算交叉熵损失 else: cpu_images, cpu_texts = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) utils.loadData(text, t) utils.loadData(length, l) preds = MORAN(image, length, text, text_rev) cost = criterion(preds, text) MORAN.zero_grad() cost.backward() optimizer.step() return cost
def val(dataset, criterion, max_iter=1000): print('Start val') data_loader = torch.utils.data.DataLoader( dataset, shuffle=False, batch_size=opt.batchSize, num_workers=int(opt.workers)) # opt.batchSize val_iter = iter(data_loader) max_iter = min(max_iter, len(data_loader)) n_correct = 0 n_total = 0 loss_avg = utils.averager() for i in range(max_iter): data = val_iter.next() if opt.BidirDecoder: cpu_images, cpu_texts, cpu_texts_rev = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) t_rev, _ = converter.encode(cpu_texts_rev, scanned=True) utils.loadData(text, t) utils.loadData(text_rev, t_rev) utils.loadData(length, l) preds0, preds1 = MORAN(image, length, text, text_rev, test=True) cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) preds0_prob, preds0 = preds0.max(1) preds0 = preds0.view(-1) preds0_prob = preds0_prob.view(-1) sim_preds0 = converter.decode(preds0.data, length.data) preds1_prob, preds1 = preds1.max(1) preds1 = preds1.view(-1) preds1_prob = preds1_prob.view(-1) sim_preds1 = converter.decode(preds1.data, length.data) sim_preds = [] for j in range(cpu_images.size(0)): text_begin = 0 if j == 0 else length.data[:j].sum() if torch.mean(preds0_prob[text_begin:text_begin+len(sim_preds0[j].split('$')[0]+'$')]).data[0] >\ torch.mean(preds1_prob[text_begin:text_begin+len(sim_preds1[j].split('$')[0]+'$')]).data[0]: sim_preds.append(sim_preds0[j].split('$')[0]+'$') else: sim_preds.append(sim_preds1[j].split('$')[0][-1::-1]+'$') else: cpu_images, cpu_texts = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) utils.loadData(text, t) utils.loadData(length, l) preds = MORAN(image, length, text, text_rev, test=True) cost = criterion(preds, text) _, preds = preds.max(1) preds = preds.view(-1) sim_preds = converter.decode(preds.data, length.data) loss_avg.add(cost) for pred, target in zip(sim_preds, cpu_texts): if pred == target.lower(): n_correct += 1 n_total += 1 print("correct / total: %d / %d, " % (n_correct, n_total)) accuracy = n_correct / float(n_total) print('Test loss: %f, : %f' % (loss_avg.val(), accuracy)) return accuracy
def __init__(self): r = rospkg.RosPack() self.path = r.get_path('moran_text_recog') self.prob_threshold = 0.90 self.cv_bridge = CvBridge() self.commodity_list = [] self.read_commodity(r.get_path('text_msgs') + "/config/commodity_list.txt") self.alphabet = '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:$' self.means = (0.485, 0.456, 0.406) self.stds = (0.229, 0.224, 0.225) self.bbox_thres = 1500 self.color_map = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(255,255,255)] # 0 90 180 270 noise self.objects = [] self.is_compressed = False self.cuda_use = torch.cuda.is_available() if self.cuda_use: cuda_flag = True self.network = MORAN(1, len(self.alphabet.split(':')), 256, 32, 100, BidirDecoder=True, CUDA=cuda_flag) self.network = self.network.cuda() else: self.network = MORAN(1, len(self.alphabet.split(':')), 256, 32, 100, BidirDecoder=True, inputDataType='torch.FloatTensor', CUDA=cuda_flag) model_name = "moran.pth" print "Moran Model Parameters number: " + str(self.count_parameters(self.network)) if self.cuda_use: state_dict = torch.load(os.path.join(self.path, "weights/", model_name)) else: state_dict = torch.load(os.path.join(self.path, "weights/", model_name), map_location='cpu') MORAN_state_dict_rename = OrderedDict() for k, v in state_dict.items(): name = k.replace("module.", "") # remove `module.` MORAN_state_dict_rename[name] = v self.network.load_state_dict(MORAN_state_dict_rename) self.converter = utils.strLabelConverterForAttention(self.alphabet, ':') self.transformer = dataset.resizeNormalize((100, 32)) for p in self.network.parameters(): p.requires_grad = False self.network.eval() #### Publisher self.image_pub = rospy.Publisher("~predict_img", Image, queue_size = 1) self.mask = rospy.Publisher("~mask", Image, queue_size = 1) self.img_bbox_pub = rospy.Publisher("~predict_bbox", Image, queue_size = 1) #### Service self.predict_ser = rospy.Service("~text_recognize_server", text_recognize_srv, self.srv_callback) image_sub1 = rospy.Subscriber('/text_detection_array', text_detection_array, self.callback, queue_size = 1) ### msg filter # image_sub = message_filters.Subscriber('/camera/color/image_raw', Image) # depth_sub = message_filters.Subscriber('/camera/aligned_depth_to_color/image_raw', Image) # ts = message_filters.TimeSynchronizer([image_sub, depth_sub], 10) # ts.registerCallback(self.callback) print "============ Ready ============"
def demo(image_path): image = Image.open(image_path).convert('L') image = transformer(image) if cuda_flag: image = image.cuda() image = image.view(1, *image.size()) image = Variable(image) text = torch.LongTensor(1 * 5) length = torch.IntTensor(1) text = Variable(text) length = Variable(length) max_iter = 20 t, l = converter.encode('0' * max_iter) utils.loadData(text, t) utils.loadData(length, l) output = MORAN(image, length, text, text, test=True, debug=True) # output image preds, preds_reverse = output[0] # 中间输出:moran demo = output[1] _, preds = preds.max(1) _, preds_reverse = preds_reverse.max(1) sim_preds = converter.decode(preds.data, length.data) sim_preds = sim_preds.strip().split('$')[0] sim_preds_reverse = converter.decode(preds_reverse.data, length.data) sim_preds_reverse = sim_preds_reverse.strip().split('$')[0] print(image_path) print('\nResult:\n' + 'Left to Right: ' + sim_preds + '\nRight to Left: ' + sim_preds_reverse) return sim_preds
def trainBatch(): data = train_iter.next() cpu_images, cpu_texts, cpu_texts_rev = data # utils.loadData(image, encode_coordinates_fn(cpu_images)) utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) t_rev, _ = converter.encode(cpu_texts_rev, scanned=True) utils.loadData(text, t) utils.loadData(text_rev, t_rev) utils.loadData(length, l) preds0, preds1 = MORAN(image, length, text, text_rev) cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) MORAN.zero_grad() cost.backward() optimizer.step() return cost
alphabet = '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:$' cuda_flag = False # print(torch.cuda.is_available()) # if torch.cuda.is_available(): # cuda_flag = True # MORAN = MORAN(1, len(alphabet.split(':')), 256, 32, 100, BidirDecoder=True, # CUDA=cuda_flag) # MORAN = MORAN.cuda() # else: # MORAN = MORAN(1, len(alphabet.split(':')), 256, 32, 100, BidirDecoder=True, inputDataType='torch.FloatTensor', # CUDA=cuda_flag) MORAN = MORAN(1, len(alphabet.split(':')), 256, 32, 100, BidirDecoder=True, inputDataType='torch.FloatTensor', CUDA=cuda_flag) print('loading pretrained model from %s' % model_path) if cuda_flag: state_dict = torch.load(model_path) else: state_dict = torch.load(model_path, map_location='cpu') MORAN_state_dict_rename = OrderedDict() for k, v in state_dict.items(): name = k.replace("module.", "") # remove `module.` MORAN_state_dict_rename[name] = v MORAN.load_state_dict(MORAN_state_dict_rename)
import torch from torch.autograd import Variable from models.moran import MORAN MORAN = MORAN(1, 37, 256, 32, 100, BidirDecoder=True, inputDataType='torch.FloatTensor', CUDA=False) input_data = Variable(torch.rand(1, 1, 64, 256)) text = Variable(torch.LongTensor([0, 0])) text_rev = text length = Variable(torch.LongTensor([2])) output = MORAN(input_data, length, text, text_rev, test=True) for i in output: print(i.size())
def train_moran_v2(config_file): import sys sys.path.append('./recognition_model/MORAN_V2') import argparse import random import torch import torch.backends.cudnn as cudnn import torch.optim as optim import torch.utils.data from torch.autograd import Variable import numpy as np import os import tools.utils as utils import tools.dataset as dataset import time from collections import OrderedDict from models.moran import MORAN from alphabet.wordlist import result from yacs.config import CfgNode as CN # from wordlistart import result from alphabet.wordlistlsvt import result def read_config_file(config_file): # 用yaml重构配置文件 f = open(config_file) opt = CN.load_cfg(f) return opt opt = read_config_file(config_file) # Modify opt.alphabet = result assert opt.ngpu == 1, "Multi-GPU training is not supported yet, due to the variant lengths of the text in a batch." if opt.experiment is None: opt.experiment = 'expr' os.system('mkdir {0}'.format(opt.experiment)) opt.manualSeed = random.randint(1, 10000) # fix seed print("Random Seed: ", opt.manualSeed) random.seed(opt.manualSeed) np.random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) cudnn.benchmark = True print(opt) if not torch.cuda.is_available(): assert not opt.cuda, 'You don\'t have a CUDA device.' if torch.cuda.is_available() and not opt.cuda: print("WARNING: You have a CUDA device, so you should probably run with --cuda") train_nips_dataset = dataset.lmdbDataset(root=opt.train_nips, transform=dataset.resizeNormalize((opt.imgW, opt.imgH)), reverse=opt.BidirDecoder) assert train_nips_dataset ''' train_cvpr_dataset = dataset.lmdbDataset(root=opt.train_cvpr, transform=dataset.resizeNormalize((opt.imgW, opt.imgH)), reverse=opt.BidirDecoder) assert train_cvpr_dataset ''' ''' train_dataset = torch.utils.data.ConcatDataset([train_nips_dataset, train_cvpr_dataset]) ''' train_dataset = train_nips_dataset train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=opt.batchSize, shuffle=False, sampler=dataset.randomSequentialSampler(train_dataset, opt.batchSize), num_workers=int(opt.workers)) test_dataset = dataset.lmdbDataset(root=opt.valroot, transform=dataset.resizeNormalize((opt.imgW, opt.imgH)), reverse=opt.BidirDecoder) nclass = len(opt.alphabet.split(opt.sep)) nc = 1 converter = utils.strLabelConverterForAttention(opt.alphabet, opt.sep) criterion = torch.nn.CrossEntropyLoss() if opt.cuda: MORAN = MORAN(nc, nclass, opt.nh, opt.targetH, opt.targetW, BidirDecoder=opt.BidirDecoder, CUDA=opt.cuda) else: MORAN = MORAN(nc, nclass, opt.nh, opt.targetH, opt.targetW, BidirDecoder=opt.BidirDecoder, inputDataType='torch.FloatTensor', CUDA=opt.cuda) if opt.MORAN != '': print('loading pretrained model from %s' % opt.MORAN) if opt.cuda: state_dict = torch.load(opt.MORAN) else: state_dict = torch.load(opt.MORAN, map_location='cpu') MORAN_state_dict_rename = OrderedDict() for k, v in state_dict.items(): name = k.replace("module.", "") # remove `module.` MORAN_state_dict_rename[name] = v MORAN.load_state_dict(MORAN_state_dict_rename, strict=True) image = torch.FloatTensor(opt.batchSize, nc, opt.imgH, opt.imgW) text = torch.LongTensor(opt.batchSize * 5) text_rev = torch.LongTensor(opt.batchSize * 5) length = torch.IntTensor(opt.batchSize) if opt.cuda: MORAN.cuda() MORAN = torch.nn.DataParallel(MORAN, device_ids=range(opt.ngpu)) image = image.cuda() text = text.cuda() text_rev = text_rev.cuda() criterion = criterion.cuda() image = Variable(image) text = Variable(text) text_rev = Variable(text_rev) length = Variable(length) # loss averager loss_avg = utils.averager() # setup optimizer if opt.adam: optimizer = optim.Adam(MORAN.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) elif opt.adadelta: optimizer = optim.Adadelta(MORAN.parameters(), lr=opt.lr) elif opt.sgd: optimizer = optim.SGD(MORAN.parameters(), lr=opt.lr, momentum=0.9) else: optimizer = optim.RMSprop(MORAN.parameters(), lr=opt.lr) def levenshtein(s1, s2): if len(s1) < len(s2): return levenshtein(s2, s1) # len(s1) >= len(s2) if len(s2) == 0: return len(s1) previous_row = range(len(s2) + 1) for i, c1 in enumerate(s1): current_row = [i + 1] for j, c2 in enumerate(s2): insertions = previous_row[j + 1] + 1 deletions = current_row[j] + 1 substitutions = previous_row[j] + (c1 != c2) current_row.append(min(insertions, deletions, substitutions)) previous_row = current_row return previous_row[-1] def val(dataset, criterion, max_iter=1000): print('Start val') data_loader = torch.utils.data.DataLoader( dataset, shuffle=False, batch_size=opt.batchSize, num_workers=int(opt.workers)) # opt.batchSize val_iter = iter(data_loader) max_iter = min(max_iter, len(data_loader)) n_correct = 0 n_total = 0 distance = 0.0 loss_avg = utils.averager() f = open('./log.txt','a',encoding='utf-8') for i in range(max_iter): data = val_iter.next() if opt.BidirDecoder: cpu_images, cpu_texts, cpu_texts_rev = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) t_rev, _ = converter.encode(cpu_texts_rev, scanned=True) utils.loadData(text, t) utils.loadData(text_rev, t_rev) utils.loadData(length, l) preds0, preds1 = MORAN(image, length, text, text_rev, test=True) cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) preds0_prob, preds0 = preds0.max(1) preds0 = preds0.view(-1) preds0_prob = preds0_prob.view(-1) sim_preds0 = converter.decode(preds0.data, length.data) preds1_prob, preds1 = preds1.max(1) preds1 = preds1.view(-1) preds1_prob = preds1_prob.view(-1) sim_preds1 = converter.decode(preds1.data, length.data) sim_preds = [] for j in range(cpu_images.size(0)): text_begin = 0 if j == 0 else length.data[:j].sum() if torch.mean(preds0_prob[text_begin:text_begin+len(sim_preds0[j].split('$')[0]+'$')]).data[0] >\ torch.mean(preds1_prob[text_begin:text_begin+len(sim_preds1[j].split('$')[0]+'$')]).data[0]: sim_preds.append(sim_preds0[j].split('$')[0]+'$') else: sim_preds.append(sim_preds1[j].split('$')[0][-1::-1]+'$') else: cpu_images, cpu_texts = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) utils.loadData(text, t) utils.loadData(length, l) preds = MORAN(image, length, text, text_rev, test=True) cost = criterion(preds, text) _, preds = preds.max(1) preds = preds.view(-1) sim_preds = converter.decode(preds.data, length.data) loss_avg.add(cost) for pred, target in zip(sim_preds, cpu_texts): if pred == target.lower(): n_correct += 1 f.write("预测 %s 目标 %s\n" % ( pred,target ) ) distance += levenshtein(pred,target) / max(len(pred),len(target)) n_total += 1 f.close() print("correct / total: %d / %d, " % (n_correct, n_total)) print('levenshtein distance: %f' % (distance/n_total)) accuracy = n_correct / float(n_total) print('Test loss: %f, accuray: %f' % (loss_avg.val(), accuracy)) return accuracy def trainBatch(): data = train_iter.next() if opt.BidirDecoder: cpu_images, cpu_texts, cpu_texts_rev = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) t_rev, _ = converter.encode(cpu_texts_rev, scanned=True) utils.loadData(text, t) utils.loadData(text_rev, t_rev) utils.loadData(length, l) preds0, preds1 = MORAN(image, length, text, text_rev) cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) else: cpu_images, cpu_texts = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) utils.loadData(text, t) utils.loadData(length, l) preds = MORAN(image, length, text, text_rev) cost = criterion(preds, text) MORAN.zero_grad() cost.backward() optimizer.step() return cost t0 = time.time() acc = 0 acc_tmp = 0 for epoch in range(opt.niter): train_iter = iter(train_loader) i = 0 while i < len(train_loader): # print("main函数里,可迭代次数为 %d" % len(train_loader)) if i % opt.valInterval == 0: for p in MORAN.parameters(): p.requires_grad = False MORAN.eval() acc_tmp = val(test_dataset, criterion) if acc_tmp > acc: acc = acc_tmp torch.save(MORAN.state_dict(), '{0}/{1}_{2}.pth'.format( opt.experiment, i, str(acc)[:6])) if i % opt.saveInterval == 0: torch.save(MORAN.state_dict(), '{0}/{1}_{2}.pth'.format( opt.experiment, epoch, i)) for p in MORAN.parameters(): p.requires_grad = True MORAN.train() cost = trainBatch() loss_avg.add(cost) if i % opt.displayInterval == 0: t1 = time.time() print ('Epoch: %d/%d; iter: %d/%d; Loss: %f; time: %.2f s;' % (epoch, opt.niter, i, len(train_loader), loss_avg.val(), t1-t0)), loss_avg.reset() t0 = time.time() i += 1
from models.moran import MORAN model_path = './demo.pth' img_dir = '/Users/thanhtu/Projects/Python/CRAFT-pytorch/cropped/' alphabet = '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:8:t:u:v:w:x:y:z:$' target_height = 32 target_width = 100 cuda_flag = False if torch.cuda.is_available(): cuda_flag = True MORAN = MORAN(1, len(alphabet.split(':')), 256, target_height, target_width, BidirDecoder=True, CUDA=cuda_flag) MORAN = MORAN.cuda() else: MORAN = MORAN(1, len(alphabet.split(':')), 256, target_height, target_width, BidirDecoder=True, inputDataType='torch.FloatTensor', CUDA=cuda_flag) print('loading pretrained model from %s' % model_path)
transform=dataset.resizeNormalize( (opt.imgW, opt.imgH)), reverse=opt.BidirDecoder) # 分类类别 nclass = len(opt.alphabet.split(opt.sep)) nc = 1 # label格式转化 converter = utils.strLabelConverterForAttention(opt.alphabet, opt.sep) criterion = torch.nn.CrossEntropyLoss() # 交叉熵loss函数 # 模型初始化 if opt.cuda: MORAN = MORAN(nc, nclass, opt.nh, opt.targetH, opt.targetW, BidirDecoder=opt.BidirDecoder, CUDA=opt.cuda) else: MORAN = MORAN(nc, nclass, opt.nh, opt.targetH, opt.targetW, BidirDecoder=opt.BidirDecoder, inputDataType='torch.FloatTensor', CUDA=opt.cuda) # 载入模型,继续训练 if opt.MORAN != '':
from collections import OrderedDict import cv2 from models.moran import MORAN model_path = './demo.pth' img_path = './demo/0.png' alphabet = '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:$' ################# 1-load model ####################################### cuda_flag = False if torch.cuda.is_available(): cuda_flag = True MORAN = MORAN(1, len(alphabet.split(':')), 256, 32, 100, BidirDecoder=True, CUDA=cuda_flag) MORAN = MORAN.cuda() else: MORAN = MORAN(1, len(alphabet.split(':')), 256, 32, 100, BidirDecoder=True, inputDataType='torch.FloatTensor', CUDA=cuda_flag) print('loading pretrained model from %s' % model_path)
import torch from torch.autograd import Variable import tools.utils as utils import tools.dataset as dataset from PIL import Image from collections import OrderedDict import cv2 from models.moran import MORAN model_path = './demo.pth' img_path = './demo/0.png' alphabet = '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:$' MORAN = MORAN(1, len(alphabet.split(':')), 256, 32, 100, BidirDecoder=True) if torch.cuda.is_available(): MORAN = MORAN.cuda() print('loading pretrained model from %s' % model_path) state_dict = torch.load(model_path) MORAN_state_dict_rename = OrderedDict() for k, v in state_dict.items(): name = k.replace("module.", "") # remove `module.` MORAN_state_dict_rename[name] = v MORAN.load_state_dict(MORAN_state_dict_rename) for p in MORAN.parameters(): p.requires_grad = False MORAN.eval() converter = utils.strLabelConverterForAttention(alphabet, ':')
imgW = 100 num_workers = 4 max_iter = 20 span_range_width = .9 span_range_height = .9 grid_height = 4 grid_width = 4 image_height = 32 image_width = 100 model = 'unbounded_stn' data = 'sp-crop' args = params() MORAN = MORAN(args, args.nc, len(args.alphabet.split(':')), args.nh, args.imgH, args.imgW).cuda() if os.path.isfile(args.model_path): model_dict = MORAN.state_dict() pretrained_model = torch.load(args.model_path) pretrained_model = { k.replace('module.', ''): v for k, v in pretrained_model.items() } pretrained_model = { k: v for k, v in pretrained_model.items() if k in model_dict } model_dict.update(pretrained_model) MORAN.load_state_dict(model_dict)
from models.morn import MORN import torch from torchsummary import summary from models.asrn_res import ResNet from models.moran import MORAN device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") # PyTorch v0.4.0 if __name__ == '__main__': # resnet = ResNet(1) # morn = MORN(1, 32, 100) # resnet.to(device) # morn.to(device) # summary(resnet, (1, 32, 100)) # summary(morn, (1, 32, 100)) # 1, 16, 256, 32, 100, # text = torch.LongTensor(opt.batchSize * 5) # text_rev = torch.LongTensor(opt.batchSize * 5) # length = torch.IntTensor(opt.batchSize) moran = MORAN(1, 16, 256, 32, 100, BidirDecoder=True, CUDA=True) moran.to(device) summary(moran, (1, 32, 100)) pass
shuffle=False, sampler=dataset.randomSequentialSampler(train_dataset, args.batchSize), num_workers=args.workers) test_dataset = dataset.lmdbDataset(root=args.valroot, transform=dataset.resizeNormalize( (args.imgW, args.imgH))) nclass = len(args.alphabet.split(args.sep)) # nc = 3 + args.imgH + args.imgW nc = 3 converter = utils.strLabelConverterForAttention(args.alphabet, args.sep) criterion = torch.nn.CrossEntropyLoss() MORAN = MORAN(nc, nclass, args.nh, args.targetH, args.targetW, CUDA=args.cuda) image = torch.FloatTensor(args.batchSize, nc, args.imgH, args.imgW) text = torch.LongTensor(args.batchSize * 5) text_rev = torch.LongTensor(args.batchSize * 5) length = torch.IntTensor(args.batchSize) if args.cuda: MORAN.cuda() image = image.cuda() text = text.cuda() text_rev = text_rev.cuda() criterion = criterion.cuda() # loss averager loss_avg = utils.averager()
num_workers=int(opt.workers)) test_dataset = dataset.lmdbDataset(root=opt.valroot, transform=dataset.resizeNormalize( (opt.imgW, opt.imgH)), reverse=opt.BidirDecoder) nclass = len(opt.alphabet.split(opt.sep)) nc = 1 converter = utils.strLabelConverterForAttention(opt.alphabet, opt.sep) criterion = torch.nn.CrossEntropyLoss() MORAN = MORAN(nc, nclass, opt.nh, opt.targetH, opt.targetW, BidirDecoder=opt.BidirDecoder) if opt.MORAN != '': print('loading pretrained model from %s' % opt.MORAN) state_dict = torch.load(opt.MORAN) MORAN_state_dict_rename = OrderedDict() for k, v in state_dict.items(): name = k.replace("module.", "") # remove `module.` MORAN_state_dict_rename[name] = v MORAN.load_state_dict(MORAN_state_dict_rename, strict=True) image = torch.FloatTensor(opt.batchSize, nc, opt.imgH, opt.imgW) text = torch.LongTensor(opt.batchSize * 5) text_rev = torch.LongTensor(opt.batchSize * 5)
def val(dataset, criterion, max_iter=10000, steps=None): data_loader = torch.utils.data.DataLoader( dataset, shuffle=False, batch_size=opt.batchSize, num_workers=int(opt.workers)) # opt.batchSize val_iter = iter(data_loader) max_iter = min(max_iter, len(data_loader)) n_correct = 0 n_total = 0 distance = 0.0 loss_avg = utils.averager() # f = open('./log.txt', 'a', encoding='utf-8') for i in range(max_iter): # 设置很大的循环数值(达不到此值就会收敛) data = val_iter.next() if opt.BidirDecoder: cpu_images, cpu_texts, cpu_texts_rev = data # data是dataloader导入的东西 utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=False) # 这个encode是将字符encode成id t_rev, _ = converter.encode(cpu_texts_rev, scanned=False) utils.loadData(text, t) utils.loadData(text_rev, t_rev) utils.loadData(length, l) preds0, preds1 = MORAN(image, length, text, text_rev, debug=False, test=True, steps=steps) # 跑模型HARN cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) preds0_prob, preds0 = preds0.max(1) # 取概率最大top1的结果 preds0 = preds0.view(-1) preds0_prob = preds0_prob.view(-1) # 维度的变形(好像是 sim_preds0 = converter.decode(preds0.data, length.data) # 将 id decode为字 preds1_prob, preds1 = preds1.max(1) preds1 = preds1.view(-1) preds1_prob = preds1_prob.view(-1) sim_preds1 = converter.decode(preds1.data, length.data) sim_preds = [] # 预测出来的字 for j in range(cpu_images.size(0)): # 对字典进行处理,把单个字符连成字符串 text_begin = 0 if j == 0 else length.data[:j].sum() if torch.mean(preds0_prob[text_begin:text_begin + len(sim_preds0[j].split('$')[0] + '$')]).item() > \ torch.mean( preds1_prob[text_begin:text_begin + len(sim_preds1[j].split('$')[0] + '$')]).item(): sim_preds.append(sim_preds0[j].split('$')[0] + '$') else: sim_preds.append(sim_preds1[j].split('$')[0][-1::-1] + '$') else: # 用不到的另一种情况 cpu_images, cpu_texts = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) utils.loadData(text, t) utils.loadData(length, l) preds = MORAN(image, length, text, text_rev, test=True) cost = criterion(preds, text) _, preds = preds.max(1) preds = preds.view(-1) sim_preds = converter.decode(preds.data, length.data) loss_avg.add(cost) # 计算loss的平均值 for pred, target in zip( sim_preds, cpu_texts ): # 与GroundTruth的对比,cpu_texts是GroundTruth,sim_preds是连接起来的字符串 if pred == target.lower(): # 完全匹配量 n_correct += 1 # f.write("pred %s\t target %s\n" % (pred, target)) distance += levenshtein(pred, target) / max( len(pred), len(target)) # 莱温斯坦距离 n_total += 1 # 完成了一个单词 # f.close() # print and save # 跑完之后输出到日志中 for pred, gt in zip(sim_preds, cpu_texts): gt = ''.join(gt.split(opt.sep)) print('%-20s, gt: %-20s' % (pred, gt)) print("correct / total: %d / %d, " % (n_correct, n_total)) print('levenshtein distance: %f' % (distance / n_total)) accuracy = n_correct / float(n_total) log.scalar_summary('Validation/levenshtein distance', distance / n_total, steps) log.scalar_summary('Validation/loss', loss_avg.val(), steps) log.scalar_summary('Validation/accuracy', accuracy, steps) print('Test loss: %f, accuray: %f' % (loss_avg.val(), accuracy)) return accuracy
def train_HARN(config_file): import sys sys.path.append('./recognition_model/HARN') import argparse import os import random import io import sys import time from models.moran import MORAN import tools.utils as utils import torch.optim as optim import numpy as np import torch.backends.cudnn as cudnn import torch.utils.data import tools.dataset as dataset from torch.autograd import Variable from collections import OrderedDict from tools.logger import logger # from wordlist import result # from wordlistlsvt import result import warnings warnings.filterwarnings('ignore') # os.environ['CUDA_VISIBLE_DEVICES'] = '1' # 指定GPU # os.environ['CUDA_VISIBLE_DEVICES'] = '5' # 指定GPU from yacs.config import CfgNode as CN def read_config_file(config_file): # 用yaml重构配置文件 f = open(config_file) opt = CN.load_cfg(f) return opt opt = read_config_file(config_file) # 获取了yaml文件 print("配置文件", opt) sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8') # Modify # opt.alphabet = result assert opt.ngpu == 1, "Multi-GPU training is not supported yet, due to the variant lengths of the text in a batch." if opt.experiment is None: opt.experiment = 'expr' os.system('mkdir {0}'.format(opt.experiment)) opt.manualSeed = random.randint(1, 10000) # fix seed random.seed(opt.manualSeed) np.random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) cudnn.benchmark = True # ---------save logger---------# log = logger('./asrn_se50_OCRdata_50_logger') # log = logger('./logger/asrn_se50_lsvt_50') # # 保存日志的路径 / 需要改 # -----------------------------# if not torch.cuda.is_available(): assert not opt.cuda, 'You don\'t have a CUDA device.' if torch.cuda.is_available() and not opt.cuda: print( "WARNING: You have a CUDA device, so you should probably run with --cuda" ) train_nips_dataset = dataset.lmdbDataset(root=opt.train_nips, transform=dataset.resizeNormalize( (opt.imgW, opt.imgH)), reverse=opt.BidirDecoder, alphabet=opt.alphabet) assert train_nips_dataset ''' train_cvpr_dataset = dataset.lmdbDataset(root=opt.train_cvpr, transform=dataset.resizeNormalize((opt.imgW, opt.imgH)), reverse=opt.BidirDecoder) assert train_cvpr_dataset ''' ''' train_dataset = torch.utils.data.ConcatDataset([train_nips_dataset, train_cvpr_dataset]) ''' train_dataset = train_nips_dataset train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=opt.batchSize, shuffle=False, sampler=dataset.randomSequentialSampler(train_dataset, opt.batchSize), num_workers=int(opt.workers)) test_dataset = dataset.lmdbDataset(root=opt.valroot, transform=dataset.resizeNormalize( (opt.imgW, opt.imgH)), reverse=opt.BidirDecoder) nclass = len( opt.alphabet.split(opt.sep) ) # 一共有多少类,英文是36,中文就是wordlist,系统只认名字为wordlist.py的文件,记得将需要用的文件改为这个名字 nc = 1 converter = utils.strLabelConverterForAttention( opt.alphabet, opt.sep) # 给每个字一个编号,例如:中(2)国(30)人(65);convert是id和字符之间的转换 criterion = torch.nn.CrossEntropyLoss() if opt.cuda: MORAN = MORAN(nc, nclass, opt.nh, opt.targetH, opt.targetW, BidirDecoder=opt.BidirDecoder, CUDA=opt.cuda, log=log) else: MORAN = MORAN(nc, nclass, opt.nh, opt.targetH, opt.targetW, BidirDecoder=opt.BidirDecoder, inputDataType='torch.FloatTensor', CUDA=opt.cuda, log=log) if opt.MORAN != '': print('loading pretrained model from %s' % opt.MORAN) if opt.cuda: state_dict = torch.load(opt.MORAN) else: state_dict = torch.load(opt.MORAN, map_location='cpu') MORAN_state_dict_rename = OrderedDict() for k, v in state_dict.items(): name = k.replace("module.", "") # remove `module.` MORAN_state_dict_rename[name] = v MORAN.load_state_dict(MORAN_state_dict_rename, strict=True) image = torch.FloatTensor(opt.batchSize, nc, opt.imgH, opt.imgW) text = torch.LongTensor(opt.batchSize * 5) text_rev = torch.LongTensor(opt.batchSize * 5) length = torch.IntTensor(opt.batchSize) if opt.cuda: MORAN.cuda() MORAN = torch.nn.DataParallel(MORAN, device_ids=range(opt.ngpu)) image = image.cuda() text = text.cuda() text_rev = text_rev.cuda() criterion = criterion.cuda() image = Variable(image) # 把图片转换成 CUDA 可以识别的 Variable 变量 text = Variable(text) text_rev = Variable(text_rev) length = Variable(length) # loss averager loss_avg = utils.averager() # setup optimizer # 优化器的选择,这里用的Adam if opt.adam: optimizer = optim.Adam(MORAN.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) elif opt.adadelta: optimizer = optim.Adadelta(MORAN.parameters(), lr=opt.lr) elif opt.sgd: optimizer = optim.SGD(MORAN.parameters(), lr=opt.lr, momentum=0.9) else: optimizer = optim.RMSprop(MORAN.parameters(), lr=opt.lr) def levenshtein(s1, s2): # 莱温斯坦距离,编辑距离的一种 if len(s1) < len(s2): return levenshtein(s2, s1) # len(s1) >= len(s2) if len(s2) == 0: return len(s1) previous_row = range(len(s2) + 1) for i, c1 in enumerate(s1): current_row = [i + 1] for j, c2 in enumerate(s2): insertions = previous_row[j + 1] + 1 deletions = current_row[j] + 1 substitutions = previous_row[j] + (c1 != c2) current_row.append(min(insertions, deletions, substitutions)) previous_row = current_row return previous_row[-1] def val(dataset, criterion, max_iter=10000, steps=None): data_loader = torch.utils.data.DataLoader( dataset, shuffle=False, batch_size=opt.batchSize, num_workers=int(opt.workers)) # opt.batchSize val_iter = iter(data_loader) max_iter = min(max_iter, len(data_loader)) n_correct = 0 n_total = 0 distance = 0.0 loss_avg = utils.averager() # f = open('./log.txt', 'a', encoding='utf-8') for i in range(max_iter): # 设置很大的循环数值(达不到此值就会收敛) data = val_iter.next() if opt.BidirDecoder: cpu_images, cpu_texts, cpu_texts_rev = data # data是dataloader导入的东西 utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=False) # 这个encode是将字符encode成id t_rev, _ = converter.encode(cpu_texts_rev, scanned=False) utils.loadData(text, t) utils.loadData(text_rev, t_rev) utils.loadData(length, l) preds0, preds1 = MORAN(image, length, text, text_rev, debug=False, test=True, steps=steps) # 跑模型HARN cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) preds0_prob, preds0 = preds0.max(1) # 取概率最大top1的结果 preds0 = preds0.view(-1) preds0_prob = preds0_prob.view(-1) # 维度的变形(好像是 sim_preds0 = converter.decode(preds0.data, length.data) # 将 id decode为字 preds1_prob, preds1 = preds1.max(1) preds1 = preds1.view(-1) preds1_prob = preds1_prob.view(-1) sim_preds1 = converter.decode(preds1.data, length.data) sim_preds = [] # 预测出来的字 for j in range(cpu_images.size(0)): # 对字典进行处理,把单个字符连成字符串 text_begin = 0 if j == 0 else length.data[:j].sum() if torch.mean(preds0_prob[text_begin:text_begin + len(sim_preds0[j].split('$')[0] + '$')]).item() > \ torch.mean( preds1_prob[text_begin:text_begin + len(sim_preds1[j].split('$')[0] + '$')]).item(): sim_preds.append(sim_preds0[j].split('$')[0] + '$') else: sim_preds.append(sim_preds1[j].split('$')[0][-1::-1] + '$') else: # 用不到的另一种情况 cpu_images, cpu_texts = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) utils.loadData(text, t) utils.loadData(length, l) preds = MORAN(image, length, text, text_rev, test=True) cost = criterion(preds, text) _, preds = preds.max(1) preds = preds.view(-1) sim_preds = converter.decode(preds.data, length.data) loss_avg.add(cost) # 计算loss的平均值 for pred, target in zip( sim_preds, cpu_texts ): # 与GroundTruth的对比,cpu_texts是GroundTruth,sim_preds是连接起来的字符串 if pred == target.lower(): # 完全匹配量 n_correct += 1 # f.write("pred %s\t target %s\n" % (pred, target)) distance += levenshtein(pred, target) / max( len(pred), len(target)) # 莱温斯坦距离 n_total += 1 # 完成了一个单词 # f.close() # print and save # 跑完之后输出到日志中 for pred, gt in zip(sim_preds, cpu_texts): gt = ''.join(gt.split(opt.sep)) print('%-20s, gt: %-20s' % (pred, gt)) print("correct / total: %d / %d, " % (n_correct, n_total)) print('levenshtein distance: %f' % (distance / n_total)) accuracy = n_correct / float(n_total) log.scalar_summary('Validation/levenshtein distance', distance / n_total, steps) log.scalar_summary('Validation/loss', loss_avg.val(), steps) log.scalar_summary('Validation/accuracy', accuracy, steps) print('Test loss: %f, accuray: %f' % (loss_avg.val(), accuracy)) return accuracy def trainBatch(steps): data = train_iter.next() if opt.BidirDecoder: cpu_images, cpu_texts, cpu_texts_rev = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) t_rev, _ = converter.encode(cpu_texts_rev, scanned=True) utils.loadData(text, t) utils.loadData(text_rev, t_rev) utils.loadData(length, l) preds0, preds1 = MORAN(image, length, text, text_rev) cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) else: cpu_images, cpu_texts = data utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) utils.loadData(text, t) utils.loadData(length, l) preds = MORAN(image, length, text, text_rev) cost = criterion(preds, text) MORAN.zero_grad() cost.backward() # 反向传播 optimizer.step() # 优化器 return cost t0 = time.time() acc, acc_tmp = 0, 0 print(' === HARN === ') for epoch in range(opt.niter): print(" === Loading Train Data === ") train_iter = iter(train_loader) i = 0 print(" === start training === ") while i < len(train_loader): # len():数据大小 # print("main函数里,可迭代次数为 %d" % len(train_loader)) steps = i + epoch * len(train_loader) # step用来计算什么时候进行存储/打印 if steps % opt.valInterval == 0: for p in MORAN.parameters(): p.requires_grad = False MORAN.eval() print('---------------Please Waiting----------------' ) # train的一些打印信息 acc_tmp = val(test_dataset, criterion, steps=steps) if acc_tmp > acc: acc = acc_tmp try: time.sleep(0.01) torch.save( MORAN.state_dict(), '{0}/{1}_{2}.pth'.format(opt.experiment, i, str(acc)[:6])) print(".pth") except RuntimeError: print("RuntimeError") pass for p in MORAN.parameters(): p.requires_grad = True MORAN.train() cost = trainBatch(steps) loss_avg.add(cost) t1 = time.time() # niter是参数部分设置的epoch数量 print('Epoch: %d/%d; iter: %d/%d; Loss: %f; time: %.2f s;' % (epoch, opt.niter, i, len(train_loader), loss_avg.val(), t1 - t0)), log.scalar_summary('train loss', loss_avg.val(), steps) # 拟合到90多/拟合到1,完全收敛,训练充分 log.scalar_summary('speed batches/persec', steps / (time.time() - t0), steps) loss_avg.reset() t0 = time.time() ''' if i % 100 == 0: t1 = time.time() # niter是参数部分设置的epoch数量 print('Epoch: %d/%d; iter: %d/%d; Loss: %f; time: %.2f s;' % (epoch, opt.niter, i, len(train_loader), loss_avg.val(), t1 - t0)), log.scalar_summary('train loss', loss_avg.val(), i) # 拟合到90多/拟合到1,完全收敛,训练充分 log.scalar_summary('speed batches/persec', i / (time.time() - t0), i) loss_avg.reset() t0 = time.time() ''' ''' if steps % opt.displayInterval == 0: t1 = time.time() # niter是参数部分设置的epoch数量 print('Epoch: %d/%d; iter: %d/%d; Loss: %f; time: %.2f s;' % (epoch, opt.niter, i, len(train_loader), loss_avg.val(), t1 - t0)), log.scalar_summary('train loss', loss_avg.val(), steps) # 拟合到90多/拟合到1,完全收敛,训练充分 log.scalar_summary('speed batches/persec', steps / (time.time() - t0), steps) loss_avg.reset() t0 = time.time() ''' i += 1
def val(dataset, criterion, max_iter=10000, steps=0): data_loader = torch.utils.data.DataLoader( dataset, shuffle=False, batch_size=args.batchSize, num_workers=args.workers) # args.batchSize val_iter = iter(data_loader) max_iter = min(max_iter, len(data_loader)) n_correct = 0 n_total = 0 distance = 0.0 loss_avg = utils.averager() f = open('logger/log.txt', 'w', encoding='utf-8') for i in range(max_iter): data = val_iter.next() cpu_images, cpu_texts, cpu_texts_rev = data # utils.loadData(image, encode_coordinates_fn(cpu_images)) utils.loadData(image, cpu_images) t, l = converter.encode(cpu_texts, scanned=True) t_rev, _ = converter.encode(cpu_texts_rev, scanned=True) utils.loadData(text, t) utils.loadData(text_rev, t_rev) utils.loadData(length, l) preds0, _, preds1, _ = MORAN(image, length, text, text_rev, debug=False, test=True, steps=steps) cost = criterion(torch.cat([preds0, preds1], 0), torch.cat([text, text_rev], 0)) preds0_prob, preds0 = preds0.max(1) preds0 = preds0.view(-1) preds0_prob = preds0_prob.view(-1) sim_preds0 = converter.decode(preds0.data, length.data) preds1_prob, preds1 = preds1.max(1) preds1 = preds1.view(-1) preds1_prob = preds1_prob.view(-1) sim_preds1 = converter.decode(preds1.data, length.data) sim_preds = [] for j in range(cpu_images.size(0)): text_begin = 0 if j == 0 else length.data[:j].sum() if torch.mean(preds0_prob[text_begin:text_begin + len(sim_preds0[j].split('$')[0] + '$')]).item() > \ torch.mean(preds1_prob[text_begin:text_begin + len(sim_preds1[j].split('$')[0] + '$')]).item(): sim_preds.append(sim_preds0[j].split('$')[0] + '$') else: sim_preds.append(sim_preds1[j].split('$')[0][-1::-1] + '$') # img_shape = cpu_images.shape[3] / 100, cpu_images.shape[2] / 100 # input_seq = cpu_texts[0] # output_seq = sim_preds[0] # attention = alpha[0] # attention_image = showAttention(input_seq, output_seq, attention, img_shape) # log.image_summary('map/attention', [attention_image], steps) loss_avg.add(cost) for pred, target in zip(sim_preds, cpu_texts): if pred == target.lower(): n_correct += 1 f.write("pred %s\t\t\t\t\ttarget %s\n" % (pred, target)) distance += levenshtein(pred, target) / max(len(pred), len(target)) n_total += 1 f.close() accuracy = n_correct / float(n_total) log.scalar_summary('Validation/levenshtein distance', distance / n_total, steps) log.scalar_summary('Validation/loss', loss_avg.val(), steps) log.scalar_summary('Validation/accuracy', accuracy, steps) return accuracy