def test(epoch, model, train_loader, phase='test'):
print '\ntest {:s}_files, epoch: {:d}'.format(phase, epoch)
mkdir(datadir+'/data/result')
model.eval()
f1score_list = []
recall_list = []
precision_list = []
word_index_dict = json.load(file.File(args.word_index_json,'r'))
index_word_dict = { v:k for k,v in word_index_dict.items() }
result_file = file.File(datadir+'/data/result/{:d}_{:s}_result.csv'.format(epoch, phase), 'w')
result_file.write('name,content\n')
name_f1score_dict = dict()
# 保存densenet生成的feature
feat_dir = args.data_dir.replace('dataset', 'feats')
mkdir(feat_dir)
feat_dir = os.path.join(feat_dir, phase)
print feat_dir
mkdir(feat_dir)
names = []
probs_all = []
for i,data in enumerate(train_loader):
if i % 50 == 0:
print('step[{:d}] OK...'.format(i))
name = data[0][0].split('/')[-1].split('.seg')[0]
names.append(name)
images, labels = [Variable(x.cuda(async=True)) for x in data[1:3]]
def load_annoataion(p, h):
'''
load annotation from the text file
:param p:
:return:
'''
text_polys = []
text_tags = []
if not os.path.exists(p):
return np.array(text_polys, dtype=np.float32)
with file.File(p, 'r') as f:
reader = csv.reader(f)
for line in reader:
label = line[-1]
# strip BOM. \ufeff for python3, \xef\xbb\bf for python2
line = [i.strip('\ufeff').strip('\xef\xbb\xbf') for i in line]
x1, y1, x2, y2, x3, y3, x4, y4 = list(map(float, line[:8]))
# x1, y1, x2, y2, x3, y3, x4, y4 = y2, h - x2, y3, h - x3, y4, h - x4, y1, h - x1
text_polys.append([[x1, y1], [x2, y2], [x3, y3], [x4, y4]])
if label == '*' or label == '###':
text_tags.append(True)
else:
text_tags.append(False)
return np.array(text_polys, dtype=np.float32), np.array(text_tags,
dtype=np.bool)
def test_single_model(FLAGS):
if FLAGS.eval_pb_path.startswith('s3//'):
pb_model_dir = '/cache/tmp/model'
if os.path.exists(pb_model_dir):
shutil.rmtree(pb_model_dir)
file.copy_parallel(FLAGS.eval_pb_path, pb_model_dir)
else:
pb_model_dir = FLAGS.eval_pb_path
signature_key = 'predict_images'
input_key_1 = 'input_img'
output_key_1 = 'output_score'
config = tf.ConfigProto(allow_soft_placement=True)
with tf.get_default_graph().as_default():
sess = tf.Session(graph=tf.Graph(), config=config)
meta_graph_def = tf.saved_model.loader.load(sess,
[tag_constants.SERVING],
pb_model_dir)
if FLAGS.eval_pb_path.startswith('s3//'):
shutil.rmtree(pb_model_dir)
signature = meta_graph_def.signature_def
input_images_tensor_name = signature[signature_key].inputs[
input_key_1].name
output_score_tensor_name = signature[signature_key].outputs[
output_key_1].name
input_images = sess.graph.get_tensor_by_name(input_images_tensor_name)
output_score = sess.graph.get_tensor_by_name(output_score_tensor_name)
img_names, test_data, test_labels = load_test_data(FLAGS)
right_count = 0
error_infos = []
for index, img in enumerate(test_data):
img = img[np.newaxis, :, :, :]
pred_score = sess.run([output_score], feed_dict={input_images: img})
if pred_score is not None:
pred_label = np.argmax(pred_score[0], axis=1)[0]
test_label = test_labels[index]
if pred_label == test_label:
right_count += 1
else:
error_infos.append('%s, %s, %s\n' %
(img_names[index], test_label, pred_label))
else:
print('pred_score is None')
accuracy = right_count / len(img_names)
print('accuracy: %s' % accuracy)
result_file_name = os.path.join(FLAGS.eval_pb_path, 'accuracy.txt')
with file.File(result_file_name, 'w') as f:
f.write('# predict error files\n')
f.write('####################################\n')
f.write('file_name, true_label, pred_label\n')
f.writelines(error_infos)
f.write('####################################\n')
f.write('accuracy: %s\n' % accuracy)
print('end')
def get_image_csv(data_path):
file_list = file.list_directory(os.path.join(data_path, 'train'))
image_list_train = []
image_list_eval = []
for i in file_list:
if '.txt' in file.list_directory(os.path.join(data_path, 'train', i)):
image_list_train = os.path.join(data_path, 'train', 'image_to_annotation.csv')
image_list_eval = os.path.join(data_path, 'eval', 'image_to_annotation.csv')
break
elif '.xml' in file.list_directory(os.path.join(data_path, 'train', i)):
with file.File(os.path.join(data_path, 'train', 'image_to_annotation.csv'), 'r') as f:
for line in f.readlines()[1:]:
image_path, image_label = line.strip().split(',')
image_list_train.append([os.path.join(data_path, image_path),
os.path.join(data_path, image_label)])
with file.File(os.path.join(data_path, 'eval', 'image_to_annotation.csv'), 'r') as f:
for line in f.readlines()[1:]:
image_path, image_label = line.strip().split(',')
image_list_eval.append([os.path.join(data_path, image_path),
os.path.join(data_path, image_label)])
break
return image_list_train, image_list_eval
def test_single_h5(FLAGS, h5_weights_path):
if not os.path.isfile(h5_weights_path):
print('%s is not a h5 weights file path' % h5_weights_path)
return
# optimizer = adam(lr=FLAGS.learning_rate, clipnorm=0.001)
optimizer = Nadam(lr=FLAGS.learning_rate,
beta_1=0.9,
beta_2=0.999,
epsilon=1e-08,
schedule_decay=0.004)
objective = 'categorical_crossentropy'
metrics = ['accuracy']
model = model_fn(FLAGS, objective, optimizer, metrics)
load_weights(model, FLAGS.eval_weights_path)
img_names, test_data, test_labels = load_test_data(FLAGS)
predictions = model.predict(test_data, verbose=0)
right_count = 0
error_infos = []
for index, pred in enumerate(predictions):
pred_label = np.argmax(pred, axis=0)
test_label = test_labels[index]
if pred_label == test_label:
right_count += 1
else:
error_infos.append('%s, %s, %s\n' %
(img_names[index], test_label, pred_label))
accuracy = right_count / len(img_names)
print('accuracy: %s' % accuracy)
result_file_name = os.path.join(
os.path.dirname(h5_weights_path),
'%s_accuracy.txt' % os.path.basename(h5_weights_path))
with file.File(result_file_name, 'w') as f:
f.write('# predict error files\n')
f.write('####################################\n')
f.write('file_name, true_label, pred_label\n')
f.writelines(error_infos)
f.write('####################################\n')
f.write('accuracy: %s\n' % accuracy)
print('end')
def main():
word_index_dict = json.load(file.File(args.word_index_json,'r'))
num_classes = len(word_index_dict)
image_label_dict = json.load(file.File(args.image_label_json,'r'))
cudnn.benchmark = True
if args.model == 'densenet':
# 两千多种字符，multi-label分类
model = DenseNet121(num_classes).cuda()
elif args.model == 'resnet':
# resnet主要用于文字区域的segmentation以及object detection操作
model = resnet.ResNet(num_classes=num_classes, args=args).cuda()
else:
return
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
# model = torch.nn.DataParallel(model).cuda()
loss = Loss().cuda()
if args.resume:
state_dict = torch.load(args.resume)
model.load_state_dict(state_dict['state_dict'])
best_f1score = state_dict['f1score']
start_epoch = state_dict['epoch'] + 1
else:
best_f1score = 0
if args.model == 'resnet':
start_epoch = 100
else:
start_epoch = 1
args.epoch = start_epoch
print 'best_f1score', best_f1score
# 划分数据集
test_filelist = sorted(file.glob(os.path.join(args.data_dir,'test','*')))
trainval_filelist = sorted(file.glob(os.path.join(args.data_dir,'train','*')))
# 两种输入size训练
# train_filelist1: 长宽比小于8:1的图片，经过padding后变成 64*512 的输入
# train_filelist2: 长宽比大于8:1的图片，经过padding,crop后变成 64*1024的输入
train_filelist1, train_filelist2 = [],[]
# 黑名单，这些图片的label是有问题的
black_list = set(json.load(file.File(args.black_json,'r'))['black_list'])
image_hw_ratio_dict = json.load(file.File(args.image_hw_ratio_json,'r'))
for f in trainval_filelist:
image = f.split('/')[-1]
if image in black_list:
continue
r = image_hw_ratio_dict[image]
if r == 0:
train_filelist1.append(f)
else:
train_filelist2.append(f)
train_val_filelist = train_filelist1 + train_filelist2
val_filelist = train_filelist1[-2048:]
train_filelist1 = train_filelist1[:-2048]
train_filelist2 = train_filelist2
image_size = [512, 64]
if args.phase in ['test', 'val', 'train_val']:
# 测试输出文字检测结果
test_dataset = dataloader.DataSet(
test_filelist,
image_label_dict,
num_classes,
# transform=train_transform,
args=args,
image_size=image_size,
phase='test')
test_loader = DataLoader(
dataset=test_dataset,
batch_size=1,
shuffle=False,
num_workers=8,
pin_memory=True)
train_filelist = train_filelist1[-2048:]
train_dataset = dataloader.DataSet(
train_filelist,
image_label_dict,
num_classes,
image_size=image_size,
args=args,
phase='test')
train_loader = DataLoader(
dataset=train_dataset,
batch_size=1,
shuffle=False,
num_workers=8,
pin_memory=True)
val_dataset = dataloader.DataSet(
val_filelist,
image_label_dict,
num_classes,
image_size=image_size,
args=args,
phase='test')
val_loader = DataLoader(
dataset=val_dataset,
batch_size=1,
shuffle=False,
num_workers=8,
pin_memory=True)
train_val_dataset = dataloader.DataSet(
train_val_filelist,
image_label_dict,
num_classes,
image_size=image_size,
args=args,
phase='test')
train_val_loader= DataLoader(
dataset=train_val_dataset,
batch_size=1,
shuffle=False,
num_workers=8,
pin_memory=True)
if args.phase == 'test':
# test(start_epoch - 1, model, val_loader, 'val')
test(start_epoch - 1, model, test_loader, 'test')
# test(start_epoch - 1, model, train_val_loader, 'train_val')
elif args.phase == 'val':
test(start_epoch - 1, model, train_loader, 'train')
test(start_epoch - 1, model, val_loader, 'val')
elif args.phase == 'train_val':
test(start_epoch - 1, model, train_val_loader, 'train_val')
return
elif args.phase == 'train':
train_dataset1 = dataloader.DataSet(
train_filelist1,
image_label_dict,
num_classes,
image_size=image_size,
args=args,
phase='train')
train_loader1 = DataLoader(
dataset=train_dataset1,
batch_size=args.batch_size,
shuffle=True,
num_workers=8,
pin_memory=True)
# train_dataset2 = dataloader.DataSet(
# train_filelist2,
# image_label_dict,
# num_classes,
# image_size=(1024,64),
# args=args,
# phase='train')
# train_loader2 = DataLoader(
# dataset=train_dataset2,
# batch_size=args.batch_size / 2,
# shuffle=True,
# num_workers=8,
# pin_memory=True)
val_dataset = dataloader.DataSet(
val_filelist,
image_label_dict,
num_classes,
image_size=image_size,
args=args,
phase='val')
val_loader = DataLoader(
dataset=val_dataset,
batch_size=min(8,args.batch_size),
shuffle=False,
num_workers=8,
pin_memory=True)
best_f1score = 0
# eval_mode = 'pretrain-2'
eval_mode = 'eval'
for epoch in range(start_epoch, args.epochs):
args.epoch = epoch
if eval_mode == 'eval':
if best_f1score > 0.8:
args.lr = 0.0001
if best_f1score > 0.7:
args.hard_mining = 1
for param_group in optimizer.param_groups:
param_group['lr'] = args.lr
train_eval(epoch, model, train_loader1, loss, optimizer, 2., 'train-1')
best_f1score = train_eval(epoch, model, val_loader, loss, optimizer, best_f1score, 'eval-{:d}-{:d}'.format(args.batch_size, args.hard_mining))
continue
def main(argv=None):
import os
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
with tf.get_default_graph().as_default():
input_images = tf.placeholder(tf.float32,
shape=[None, None, None, 3],
name='input_images')
global_step = tf.get_variable('global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
f_score, f_geometry = model.model(input_images, is_training=False)
variable_averages = tf.train.ExponentialMovingAverage(
0.997, global_step)
saver = tf.train.Saver(variable_averages.variables_to_restore())
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
ckpt_state = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
model_path = os.path.join(
FLAGS.checkpoint_path,
os.path.basename(ckpt_state.model_checkpoint_path))
print('Restore from {}'.format(model_path))
saver.restore(sess, model_path)
im_fn_list = get_images()
for im_fn in im_fn_list:
im = cv2.imread(im_fn)[:, :, ::-1]
start_time = time.time()
im_resized, (ratio_h, ratio_w) = resize_image(im)
timer = {'net': 0, 'restore': 0, 'nms': 0}
start = time.time()
score, geometry = sess.run(
[f_score, f_geometry],
feed_dict={input_images: [im_resized]})
timer['net'] = time.time() - start
boxes, timer = detect(score_map=score,
geo_map=geometry,
timer=timer)
print(
'{} : net {:.0f}ms, restore {:.0f}ms, nms {:.0f}ms'.format(
im_fn, timer['net'] * 1000, timer['restore'] * 1000,
timer['nms'] * 1000))
if boxes is not None:
boxes = boxes[:, :8].reshape((-1, 4, 2))
boxes[:, :, 0] /= ratio_w
boxes[:, :, 1] /= ratio_h
duration = time.time() - start_time
print('[timing] {}'.format(duration))
# save to file
output_path = FLAGS.output
with file.File(output_path, 'a') as f:
if boxes is not None:
for box in boxes:
# to avoid submitting errors
box = sort_poly(box.astype(np.int32))
if np.linalg.norm(box[0] -
box[1]) < 5 or np.linalg.norm(
box[3] - box[0]) < 5:
continue
f.write('{},{},{},{},{},{},{},{},{}\r\n'.format(
os.path.basename(im_fn).split('.')[0],
box[0, 0],
box[0, 1],
box[1, 0],
box[1, 1],
box[2, 0],
box[2, 1],
box[3, 0],
box[3, 1],
))
cv2.polylines(
im[:, :, ::-1],
[box.astype(np.int32).reshape((-1, 1, 2))],
True,
color=(255, 255, 0),
thickness=1)
