class OcrDetCTPN():
def __init__(self, model_path='./checkpoints/CTPN.pth'):
self.model = CTPN_Model()
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
self.model.cuda()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.model.load_state_dict(
torch.load(model_path, map_location=device)['model_state_dict'])
for p in self.model.parameters():
p.requires_grad = False
self.model.eval()
self.prob_thresh = 0.5
def inference(self, image):
image_sz = resize(image, height=ctpn_params.IMAGE_HEIGHT)
# 宽高缩放比例(等比例缩放)
rescale_fac = image.shape[0] / image_sz.shape[0]
h, w = image_sz.shape[:2]
# 减均值
image_sz = image_sz.astype(np.float32) - ctpn_params.IMAGE_MEAN
image_sz = torch.from_numpy(image_sz.transpose(
2, 0, 1)).unsqueeze(0).float()
if self.use_gpu:
image_sz = image_sz.cuda()
cls, regr = self.model(image_sz)
cls_prob = F.softmax(cls, dim=-1).cpu().numpy()
regr = regr.cpu().numpy()
anchor = gen_anchor((int(h / 16), int(w / 16)), 16)
bbox = bbox_transfor_inv(anchor, regr)
bbox = clip_box(bbox, [h, w])
fg = np.where(cls_prob[0, :, 1] > self.prob_thresh)[0]
select_anchor = bbox[fg, :]
select_score = cls_prob[0, fg, 1]
select_anchor = select_anchor.astype(np.int32)
keep_index = filter_bbox(select_anchor, 16)
# nms
select_anchor = select_anchor[keep_index]
select_score = select_score[keep_index]
select_score = np.reshape(select_score, (select_score.shape[0], 1))
nmsbox = np.hstack((select_anchor, select_score))
keep = nms(nmsbox, 0.3)
select_anchor = select_anchor[keep]
select_score = select_score[keep]
# text line-
textConn = TextProposalConnectorOriented()
text = textConn.get_text_lines(select_anchor, select_score, [h, w])
text = [np.hstack((res[:8] * rescale_fac, res[8])) for res in text]
return text
batch_size=1,
shuffle=True,
num_workers=args['num_workers'])
model = CTPN_Model()
model.to(device)
if os.path.exists(checkpoints_weight):
print('using pretrained weight: {}'.format(checkpoints_weight))
cc = torch.load(checkpoints_weight, map_location=device)
model.load_state_dict(cc['model_state_dict'])
try:
resume_epoch = cc['epoch']
except KeyError:
resume_epoch = 0
params_to_uodate = model.parameters()
optimizer = optim.SGD(params_to_uodate, lr=lr, momentum=0.9)
critetion_cls = RPN_CLS_Loss(device)
critetion_regr = RPN_REGR_Loss(device)
best_loss_cls = 100
best_loss_regr = 100
best_loss = 100
best_model = None
epochs += resume_epoch
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
for epoch in range(resume_epoch + 1, epochs):
print(f'Epoch {epoch}/{epochs}')
print('#' * 50)
if __name__ == "__main__":
random_seed = 2020
torch.random.manual_seed(random_seed)
np.random.seed(random_seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(random_seed)
opt = parser.parse_args()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = CTPN_Model()
criterion_cls = RPN_CLS_Loss(device)
criterion_regr = RPN_REGR_Loss(device)
optimizer = optim.SGD([{
'params': model.parameters(),
'initial_lr': 0.001
}],
lr=ctpn_params.lr,
momentum=0.99,
weight_decay=0.0005)
if torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
model = model.cuda()
resume_epoch = 0
# 加载已训练模型，用于断点继续训练
if opt.resume_net != '' and os.path.exists(opt.resume_net):
print('loading pretrained model from %s' % opt.resume_net)
cc = torch.load(opt.resume_net, map_location=device)
model.load_state_dict(cc['model_state_dict'])
