opencv使用最小外接矩形比较方便,但是不能紧包图像。使用近似多边形的方法可以实现紧包图像边界。
因为使用分割模型去训练证件检测,检测出来后是mask,需要从mask中获取到四个顶点的坐标信息。
以下是整个流程,包含推理和后处理,如果只需要后处理部分就从中摘取需要的部分。
1、基于最小外接矩形的处理
# -*- coding : UTF-8 -*-
# @file : gen_bankcard_label.py
# @Time : 2021/6/11 17:45
# @Author : wmz
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
import sys
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(BASE_DIR, '..'))
import argparse
import torch
import cv2
import time
import albumentations as A
import numpy as np
from models.build_BiSeNet import BiSeNet
import xml.dom.minidom as minidom
# 计算任意多边形的面积,顶点按照顺时针或者逆时针方向排列
def compute_polygon_area(points):
point_num = len(points)
if(point_num < 3): return 0.0
s = points[0][1] * (points[point_num-1][0] - points[1][0])
#for i in range(point_num): # (int i = 1 i < point_num ++i):
for i in range(1, point_num): # 有小伙伴发现一个bug,这里做了修改,但是没有测试,需要使用的亲请测试下,以免结果不正确。
s += points[i][1] * (points[i-1][0] - points[(i+1)%point_num][0])
return abs(s/2.0)
def save_points2xml(points, xmlfile):
# write file
# 1.创建DOM树对象
dom = minidom.Document()
# 2.创建根节点。每次都要用DOM对象来创建任何节点。
root_node = dom.createElement('ImageInfo')
# 3.用DOM对象添加根节点
dom.appendChild(root_node)
# 设置该节点的属性
info = "正面[0]"
root_node.setAttribute('bModify', '3')
points = [int(y) for x in points for y in x]
# 用DOM对象创建元素子节点
info_node = dom.createElement('LineInfo')
# 用父节点对象添加元素子节点
root_node.appendChild(info_node)
# 设置该节点的属性
info_node.setAttribute('ptLTX', str(points[0]))
info_node.setAttribute('ptLTY', str(points[1]))
info_node.setAttribute('ptRTX', str(points[2]))
info_node.setAttribute('ptRTY', str(points[3]))
info_node.setAttribute('ptRBX', str(points[4]))
info_node.setAttribute('ptRBY', str(points[5]))
info_node.setAttribute('ptLBX', str(points[6]))
info_node.setAttribute('ptLBY', str(points[7]))
info_node.setAttribute('Chars', info)
info_node.setAttribute('bModify', '3')
with open(xmlfile, 'w', encoding='UTF-8') as fh:
dom.writexml(fh, indent='', addindent='\t', newl='\n', encoding='UTF-8')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
parser = argparse.ArgumentParser(description='Training')
# parser.add_argument('--path_checkpoint', default=r"G:\project_class_bak\results\seg_old_3\04-09_11-23-portrait-512-sup-8500\checkpoint_best.pkl",
# help="path to your dataset")
parser.add_argument('--path_checkpoint', default=r"./results/06-12_23-05/checkpoint_best.pkl",
help="path to your dataset")
# parser.add_argument('--path_img', default=r"G:\deep_learning_data\EG_dataset\dataset\traqining\00004.png",
# help="path to your dataset")
parser.add_argument('--data_root_dir', default=r"D:\GitHub\img_seg\test\yhk",
help="path to your dataset")
parser.add_argument('--result_root_dir', default=r"D:\GitHub\img_seg\test\maskresultyhk",
help="path to your dataset")
args = parser.parse_args()
if __name__ == '__main__':
# step2: 模型加载
# model = BiSeNet(num_classes=1, context_path="resnet101")
model = BiSeNet(num_classes=1, context_path="resnet18")
checkpoint = torch.load(args.path_checkpoint)
model.load_state_dict(checkpoint["model_state_dict"])
model.to(device)
model.eval()
# path_img = args.path_img
data_root_dir = args.data_root_dir
result_root_dir = args.result_root_dir
if not os.path.exists(result_root_dir):
os.makedirs(result_root_dir)
for root, dirs, files in os.walk(data_root_dir):
imgs =