opensource/PaddleOCR
3
0
Fork 0
mirror of https://github.com/PaddlePaddle/PaddleOCR.git synced 2026-09-24 23:33:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'dygraph' of https://github.com/PaddlePaddle/PaddleOCR into table_pr

Browse Source
This commit is contained in:
WenmuZhou
2022-08-14 09:07:19 +00:00
parent ce321153e7 7054013004
commit ba284a2b74
23 changed files with 1717 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/program.py
+41 -5
View File
@@ -25,6 +25,8 @@ import datetime
import paddle
import paddle.distributed as dist
from tqdm import tqdm
import cv2
import numpy as np
from argparse import ArgumentParser, RawDescriptionHelpFormatter
from ppocr.utils.stats import TrainingStats
@@ -262,6 +264,7 @@ def train(config,
config, 'Train', device, logger, seed=epoch)
max_iter = len(train_dataloader) - 1 if platform.system(
) == "Windows" else len(train_dataloader)
for idx, batch in enumerate(train_dataloader):
profiler.add_profiler_step(profiler_options)
train_reader_cost += time.time() - reader_start
@@ -289,7 +292,7 @@ def train(config,
else:
if model_type == 'table' or extra_input:
preds = model(images, data=batch[1:])
elif model_type in ["kie", 'vqa']:
elif model_type in ["kie", 'vqa', 'sr']:
preds = model(batch)
else:
preds = model(images)
@@ -297,11 +300,12 @@ def train(config,
avg_loss = loss['loss']
avg_loss.backward()
optimizer.step()
optimizer.clear_grad()
if cal_metric_during_train and epoch % calc_epoch_interval == 0: # only rec and cls need
batch = [item.numpy() for item in batch]
if model_type in ['kie']:
if model_type in ['kie', 'sr']:
eval_class(preds, batch)
elif model_type in ['table']:
post_result = post_process_class(preds, batch)
@@ -347,8 +351,8 @@ def train(config,
len(train_dataloader) - idx - 1) * eta_meter.avg
eta_sec_format = str(datetime.timedelta(seconds=int(eta_sec)))
strs = 'epoch: [{}/{}], global_step: {}, {}, avg_reader_cost: ' \
'{:.5f} s, avg_batch_cost: {:.5f} s, avg_samples: {}, ' \
'ips: {:.5f} samples/s, eta: {}'.format(
'{:.5f} s, avg_batch_cost: {:.5f} s, avg_samples: {}, ' \
'ips: {:.5f} samples/s, eta: {}'.format(
epoch, epoch_num, global_step, logs,
train_reader_cost / print_batch_step,
train_batch_cost / print_batch_step,
@@ -480,6 +484,7 @@ def eval(model,
leave=True)
max_iter = len(valid_dataloader) - 1 if platform.system(
) == "Windows" else len(valid_dataloader)
sum_images = 0
for idx, batch in enumerate(valid_dataloader):
if idx >= max_iter:
break
@@ -493,6 +498,20 @@ def eval(model,
preds = model(images, data=batch[1:])
elif model_type in ["kie", 'vqa']:
preds = model(batch)
elif model_type in ['sr']:
preds = model(batch)
sr_img = preds["sr_img"]
lr_img = preds["lr_img"]
for i in (range(sr_img.shape[0])):
fm_sr = (sr_img[i].numpy() * 255).transpose(
1, 2, 0).astype(np.uint8)
fm_lr = (lr_img[i].numpy() * 255).transpose(
1, 2, 0).astype(np.uint8)
cv2.imwrite("output/images/{}_{}_sr.jpg".format(
sum_images, i), fm_sr)
cv2.imwrite("output/images/{}_{}_lr.jpg".format(
sum_images, i), fm_lr)
else:
preds = model(images)
else:
@@ -500,6 +519,20 @@ def eval(model,
preds = model(images, data=batch[1:])
elif model_type in ["kie", 'vqa']:
preds = model(batch)
elif model_type in ['sr']:
preds = model(batch)
sr_img = preds["sr_img"]
lr_img = preds["lr_img"]
for i in (range(sr_img.shape[0])):
fm_sr = (sr_img[i].numpy() * 255).transpose(
1, 2, 0).astype(np.uint8)
fm_lr = (lr_img[i].numpy() * 255).transpose(
1, 2, 0).astype(np.uint8)
cv2.imwrite("output/images/{}_{}_sr.jpg".format(
sum_images, i), fm_sr)
cv2.imwrite("output/images/{}_{}_lr.jpg".format(
sum_images, i), fm_lr)
else:
preds = model(images)
@@ -517,12 +550,15 @@ def eval(model,
elif model_type in ['table', 'vqa']:
post_result = post_process_class(preds, batch_numpy)
eval_class(post_result, batch_numpy)
elif model_type in ['sr']:
eval_class(preds, batch_numpy)
else:
post_result = post_process_class(preds, batch_numpy[1])
eval_class(post_result, batch_numpy)
pbar.update(1)
total_frame += len(images)
sum_images += 1
# Get final metric，eg. acc or hmean
metric = eval_class.get_metric()
@@ -617,7 +653,7 @@ def preprocess(is_train=False):
'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE',
'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'LayoutLMv2', 'PREN', 'FCE',
'SVTR', 'ViTSTR', 'ABINet', 'DB++', 'TableMaster', 'SPIN', 'VisionLAN',
'SLANet'
'Gestalt', 'SLANet'
]
if use_xpu: