def main(args: argparse.Namespace):
model_path = args.model
image_dir = args.image_dir
output_img_dir = args.output_img_dir
output_txt_dir = args.output_txt_dir
if output_img_dir is not None and not os.path.exists(output_img_dir):
os.makedirs(output_img_dir)
if output_txt_dir is not None and not os.path.exists(output_txt_dir):
os.makedirs(output_txt_dir)
annotation_dir = args.annotation_dir
with_image = True if output_img_dir else False # 是否输出预测的图片
with_gpu = True if torch.cuda.is_available(
) and not args.no_gpu else False # 是否使用gpu
model = load_model(model_path, with_gpu)
if annotation_dir is not None: # 有标注文件就计算预测的各项指标
true_pos, true_neg, false_pos, false_neg = [0] * 4
for image_fn in tqdm(image_dir.glob('*.jpg')):
gt_path = annotation_dir / image_fn.with_suffix(
'.txt').name # 直接将.jpg的文件改成.txt就是对应的标注了
labels = load_annotation(gt_path)
# try:
with torch.no_grad(): # 计算模型在数据集上每个样本的预测值并保存预测图像、文本
polys, im, res = Toolbox.predict(
image_fn, model, with_image, output_img_dir, with_gpu,
labels, output_txt_dir,
strLabelConverter(getattr(common_str, args.keys)))
true_pos += res[0]
false_pos += res[1]
false_neg += res[2]
if (true_pos + false_pos) > 0:
precision = true_pos / (true_pos + false_pos)
else:
precision = 0
if (true_pos + false_neg) > 0:
recall = true_pos / (true_pos + false_neg)
else:
recall = 0
print("TP: %d, FP: %d, FN: %d, precision: %f, recall: %f" %
(true_pos, false_pos, false_neg, precision, recall))
else: # 没有标注文件就仅仅输出预测图像并保存
with torch.no_grad():
for image_fn in tqdm(image_dir.glob('*.jpg')):
Toolbox.predict(
image_fn, model, with_image, output_img_dir, with_gpu,
None, None,
strLabelConverter(getattr(common_str, args.keys)))
def main(args: argparse.Namespace):
model_path = args.model
image_dir = args.image_dir
output_img_dir = args.output_img_dir
output_txt_dir = args.output_txt_dir
if output_img_dir is not None and not os.path.exists(output_img_dir):
os.makedirs(output_img_dir)
if output_txt_dir is not None and not os.path.exists(output_txt_dir):
os.makedirs(output_txt_dir)
annotation_dir = args.annotation_dir
with_image = True if output_img_dir else False
with_gpu = True if torch.cuda.is_available() and not args.no_gpu else False
model = load_model(model_path, with_gpu)
if annotation_dir is not None:
true_pos, true_neg, false_pos, false_neg = [0] * 4
for image_fn in tqdm(image_dir.glob('*.jpg')):
gt_path = annotation_dir / image_fn.with_name('gt_{}'.format(
image_fn.stem)).with_suffix('.txt').name
labels = load_annotation(gt_path)
# try:
with torch.no_grad():
polys, im, res = Toolbox.predict(
image_fn, model, with_image, output_img_dir, with_gpu,
labels, output_txt_dir,
strLabelConverter(getattr(common_str, args.keys)))
true_pos += res[0]
false_pos += res[1]
false_neg += res[2]
if (true_pos + false_pos) > 0:
precision = true_pos / (true_pos + false_pos)
else:
precision = 0
if (true_pos + false_neg) > 0:
recall = true_pos / (true_pos + false_neg)
else:
recall = 0
print("TP: %d, FP: %d, FN: %d, precision: %f, recall: %f" %
(true_pos, false_pos, false_neg, precision, recall))
else:
with torch.no_grad():
for image_fn in tqdm(image_dir.glob('*.jpg')):
Toolbox.predict(
image_fn, model, with_image, output_img_dir, with_gpu,
None, output_txt_dir,
strLabelConverter(getattr(common_str, args.keys)))
def __init__(self,
model,
loss,
metrics,
finetune,
resume,
config,
data_loader,
toolbox: Toolbox,
valid_data_loader=None,
train_logger=None,
keys=custom_1,
writer=None):
super(Trainer, self).__init__(model, loss, metrics, finetune, resume,
config, train_logger)
self.config = config
self.batch_size = data_loader.batch_size
self.data_loader = data_loader
self.valid_data_loader = valid_data_loader
self.valid = True if self.valid_data_loader is not None else False
self.log_step = int(np.sqrt(self.batch_size))
self.toolbox = toolbox
self.labelConverter = strLabelConverter(keys)
self.writer = writer
def __init__(self, model_path=CONFIG.model_path, alphabet=CONFIG.alphabet):
self.model = crnn.CRNN(n_class=len(alphabet) + 1)
self.model.load_state_dict(torch.load(model_path)['state_dict'])
self.converter = util.strLabelConverter(alphabet)
self.toTensor = transforms.ToTensor()
self.model = self.model.cuda()
self.model.eval()
def __init__(self, config):
"""
负责服务初始化
"""
self.model = FOTSModel(config, False)
self.model.eval()
self.config = config
self.model.load_state_dict(
torch.load(config['model_path'])['state_dict'])
self.label_converter = strLabelConverter(
getattr(common_str, self.config['model']['keys']))
if config['cuda']:
self.model.to(torch.device("cuda:0"))
print('init finish')
def main(config_yaml):
'''
Training/Finetune CNN_RNN_Attention Model.
'''
#### Load config settings. ####
f = open(config_yaml, encoding='utf-8')
opt = yaml.load(f)
if os.path.isdir(opt['LOGGER_PATH']) == False:
os.mkdir(opt['LOGGER_PATH'])
logger = Logger(opt['LOGGER_PATH'])
if os.path.isdir(opt['SAVE_PATH']) == False:
os.system('mkdir -p {0}'.format(opt['SAVE_PATH']))
manualSeed = random.randint(1, 10000)
random.seed(manualSeed)
np.random.seed(manualSeed)
torch.manual_seed(manualSeed)
cudnn.benchmark = True
#### Set up DataLoader. ####
train_cfg = opt['TRAIN']
ds_cfg = train_cfg['DATA_SOURCE']
print('Building up dataset:{}'.format(ds_cfg['TYPE']))
if ds_cfg['TYPE'] == 'SYN_DATA':
text_gen = util.TextGenerator(ds_cfg['GEN_SET'], ds_cfg['GEN_LEN'])
ds_train = dataset.synthDataset(ds_cfg['FONT_ROOT'],
ds_cfg['FONT_SIZE'], text_gen)
elif ds_cfg['TYPE'] == 'IMG_DATA':
ds_train = dataset.trainDataset(
ds_cfg['IMG_ROOT'], ds_cfg['TRAIN_SET'],
transform=None) #dataset.graybackNormalize()
assert ds_train
train_loader = torch.utils.data.DataLoader(
ds_train,
batch_size=train_cfg['BATCH_SIZE'],
shuffle=True,
sampler=None,
num_workers=opt['WORKERS'],
collate_fn=dataset.alignCollate(imgH=train_cfg['IMG_H'],
imgW=train_cfg['MAX_W']))
val_cfg = opt['VALIDATION']
ds_val = dataset.testDataset(val_cfg['IMG_ROOT'],
val_cfg['VAL_SET'],
transform=None) #dataset.graybackNormalize()
assert ds_val
val_loader = torch.utils.data.DataLoader(ds_val,
batch_size=32,
shuffle=False,
num_workers=opt['WORKERS'],
collate_fn=dataset.alignCollate(
imgH=train_cfg['IMG_H'],
imgW=train_cfg['MAX_W']))
#### Model construction and Initialization. ####
alphabet = keys.alphabet
nClass = len(alphabet) + 1
if opt['N_GPU'] > 1:
opt['RNN']['multi_gpu'] = True
else:
opt['RNN']['multi_gpu'] = False
model = crann.CRANN(opt, nClass)
#print(model)
#### Train/Val the model. ####
converter = util.strLabelConverter(alphabet)
criterion = CTCLoss()
if opt['CUDA']:
model.cuda()
criterion.cuda()
if opt['OPTIMIZER'] == 'RMSprop':
optimizer = optim.RMSprop(model.parameters(), lr=opt['TRAIN']['LR'])
elif opt['OPTIMIZER'] == 'Adam':
optimizer = optim.Adam(model.parameters(),
lr=opt['TRAIN']['LR'],
betas=(opt['TRAIN']['BETA1'], 0.999))
elif opt['OPTIMIZER'] == 'SGD':
optimizer = optim.SGD(model.parameters(), lr=opt['TRAIN']['LR'])
else:
optimizer = optim.Adadelta(model.parameters(), lr=opt['TRAIN']['LR'])
start_epoch = 0
if opt['VAL_ONLY']:
print('=>loading pretrained model from %s for val only.' %
opt['CRANN'])
checkpoint = torch.load(opt['CRANN'])
model.load_state_dict(checkpoint['state_dict'])
val(model, val_loader, criterion, converter, 0, 0, logger, True)
elif opt['FINETUNE']:
print('=>loading pretrained model from %s for finetuen.' %
opt['CRANN'])
checkpoint = torch.load(opt['CRANN'])
#model.load_state_dict(checkpoint['state_dict'])
model_dict = model.state_dict()
#print(model_dict.keys())
cnn_dict = {
"cnn." + k: v
for k, v in checkpoint.items() if "cnn." + k in model_dict
}
model_dict.update(cnn_dict)
model.load_state_dict(model_dict)
for epoch in range(start_epoch, opt['EPOCHS']):
adjust_lr(optimizer, opt['TRAIN']['LR'], epoch, opt['STEP'])
train(model, train_loader, val_loader, criterion, optimizer, opt,
converter, epoch, logger)
elif opt['RESUME']:
print('=>loading checkpoint from %s for resume training.' %
opt['CRANN'])
checkpoint = torch.load(opt['CRANN'])
start_epoch = checkpoint['epoch'] + 1
print('resume from epoch:{}'.format(start_epoch))
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
for epoch in range(start_epoch, opt['EPOCHS']):
adjust_lr(optimizer, opt['TRAIN']['LR'], epoch, opt['STEP'])
train(model, train_loader, val_loader, criterion, optimizer, opt,
converter, epoch, logger)
else:
print('train from scratch.')
for epoch in range(start_epoch, opt['EPOCHS']):
adjust_lr(optimizer, opt['TRAIN']['LR'], epoch, opt['STEP'])
train(model, train_loader, val_loader, criterion, optimizer, opt,
converter, epoch, logger)
model_path = 'checkpoints/grcnn_art/crann_11_1.pth'
batch_size = 16
imgH = 32
maxW = 100
num_workers = 4
cnn_model = 'grcnn'
rnn_model = 'compositelstm'
n_In = 512
n_Hidden = 256
test_set = '../art_test.txt'
if __name__ == '__main__':
alphabet = keys.alphabet
nClass = len(alphabet) + 1
converter = util.strLabelConverter(alphabet)
model = crann.CRANN(cnn_model, rnn_model, n_In, n_Hidden, nClass).cuda()
if os.path.isfile(model_path):
print("=> loading checkpoint '{}'".format(model_path))
checkpoint = torch.load(model_path)
start_epoch = checkpoint['epoch']
# best_pred = checkpoint['best_pred']
model.load_state_dict(checkpoint['state_dict'])
# print("=> loaded checkpoint '{}' (epoch {} accuracy {})"
# .format(model_path, checkpoint['epoch'], best_pred))
model.eval()
train_set = dataset.imageDataset(test_set) # dataset.graybackNormalize()
test_loader = torch.utils.data.DataLoader(train_set,
})
self.write(json.dumps(to_return))
if __name__ == "__main__":
ag = ArgumentParser()
ag.add_argument("-c", type=str, help='path to config file')
args = ag.parse_args()
with open(args.c, mode='r', encoding='utf-8') as to_read:
config = json.loads(to_read.read())
model = FOTSModel(config, False)
model.eval()
model.load_state_dict(torch.load(config['model_path'])['state_dict'])
label_converter = strLabelConverter(
getattr(common_str, config['model']['keys']))
with_gpu = config['cuda']
if with_gpu:
model.to(torch.device("cuda:0"))
routes = [
('/detect', DetectHandler),
('/recognize', DetectHandler),
('/detect_and_recognize', DetectHandler),
]
application = tornado.web.Application(routes)
http_server = tornado.httpserver.HTTPServer(application)
http_server.listen(config['port'])
print('init finish')
tornado.ioloop.IOLoop.current().start()
def demo_grcnn(config_yaml):
import sys
sys.path.append('./recognition_model/GRCNN')
import torch
import os
from utils import keys
from models import crann
import dataset
from utils import util
import torch.nn.functional as F
import io
import yaml
import tools.utils as utils
import tools.dataset_lmdb as dataset_lmdb
import torchvision.transforms as transforms
import lmdb
import cv2
# 需要在配置文件里体现
# opt.model_path = 'checkpoints/grcnn_art/crann_11_1.pth'
# batch_size = 16
#imgH = 32
# maxW = 100
# num_workers = 4
# cnn_model = 'grcnn'
# rnn_model = 'compositelstm'
# n_In = 512
# n_Hidden = 256
# test_set = '../art_test.txt'
# 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)
f = open(config_yaml, encoding='utf-8')
opt = yaml.load(f)
alphabet = keys.alphabet
nClass = len(alphabet) + 1
converter = util.strLabelConverter(alphabet)
model = crann.CRANN(opt, nClass).cuda()
if os.path.isfile(opt['DEMO']['model_path']):
print("=> loading checkpoint '{}'".format(opt['DEMO']['model_path']))
checkpoint = torch.load(opt['DEMO']['model_path'])
start_epoch = checkpoint['epoch']
# best_pred = checkpoint['best_pred']
model.load_state_dict(checkpoint['state_dict'])
# print("=> loaded checkpoint '{}' (epoch {} accuracy {})"
# .format(opt.model_path, checkpoint['epoch'], best_pred))
model.eval()
# root, mappinggit
train_set = dataset_lmdb.lmdbDataset(opt['DEMO']['test_set_lmdb'])
# train_set = dataset.testDataset(opt['test_set']) # dataset.graybackNormalize()
test_loader = torch.utils.data.DataLoader(
train_set,
batch_size=opt['TRAIN']['BATCH_SIZE'],
shuffle=False,
num_workers=opt['WORKERS'],
collate_fn=dataset.alignCollate(imgH=opt['TRAIN']['IMG_H'],
imgW=opt['TRAIN']['MAX_W']))
file = open('./pred.txt', 'w', encoding='utf-8')
try:
import shutil
shutil.rmtree('./GRCNN_DEMO')
# os.makedirs('./MORAN_DEMO')
except:
pass
os.makedirs('./GRCNN_DEMO')
record_file = open('./GRCNN_DEMO/result.txt', 'a', encoding='utf-8')
index = 0
for i, (cpu_images, targets) in enumerate(test_loader):
# 还可以再改造一下
bsz = cpu_images.size(0)
images = cpu_images.cuda()
predict = model(images)
predict_len = torch.IntTensor([predict.size(0)] * bsz)
_, acc = predict.max(2)
acc = acc.transpose(1, 0).contiguous().view(-1)
prob, _ = F.softmax(predict, dim=2).max(2)
probilities = torch.mean(prob, dim=1)
sim_preds = converter.decode(acc.data, predict_len.data, raw=False)
cnt = 0
for probility, pred, target in zip(probilities, sim_preds, targets):
index += 1
img_key = 'gt_%d' % index
file.write('%s:\t\t\t\t%.3f%%\t%-20s\n' %
(img_key, probility.item() * 100, pred))
# print("调试开始")
# print(images[0].size)
# print("调试结束")
# cv2.imwrite('./GRCNN_DEMO/' + str(index) + '.jpg', (images[cnt].cpu().numpy() + 1.0) * 128)
record_file.write('./GRCNN_DEMO/' + str(index) + '.jpg' + ' ' +
pred + ' ' + target + ' \n')
cnt += 1
file.close()
