def test_batch_h5(FLAGS):
"""
test all the h5 weights files in the model_dir
"""
file_paths = file.glob(os.path.join(FLAGS.eval_weights_path, '*.h5'))
for file_path in file_paths:
test_single_h5(FLAGS, file_path)
def get_images():
files = []
files.extend(
file.glob(os.path.join(FLAGS.training_data_path,
'*.{}'.format('jpg'))))
return files
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