目录
前言
在ICDAR2021竞赛中,需要取出图片中需要分割的样本,对样本进行仿射变换,背景替换,生成更丰富的样本。
目标
结合mask取出input中的两张卡片
input mask
cards
思路
1.利用cv2.findcontours函数检测物体的轮廓;
2.利用cv2.contourArea函数筛除不合格的轮廓点集;
3.利用cv2.approxPolyDP函数求检测到的轮廓点集中的顶点坐标;
4.有了card的四个顶点,就可以求出card近似的的边长和宽,利用cv2.getPerspectiveTransform函数求出仿射变换矩阵M;
5.利用cv2.warpPerspective函数将input图片中的card提出;
完整代码
#-*- coding:utf-8 -*-
import cv2
import numpy as np
import os
from tqdm import tqdm
def order_points(pts):
''' sort rectangle points by clockwise '''
sort_x = pts[np.argsort(pts[:, 0]), :]
Left = sort_x[:2, :]
Right = sort_x[2:, :]
# Left sort
Left = Left[np.argsort(Left[:, 1])[::-1], :]
# Right sort
Right = Right[np.argsort(Right[:, 1]), :]
return np.concatenate((Left, Right), axis=0)
def get_doc_area(imgs_root, masks_root, dst_path, dst_sub_dirs):
imgs_nm = os.listdir(imgs_root)
masks_nm = os.listdir(masks_root)
for sub_dir in dst_sub_dirs:
if not os.path.exists(os.path.join(dst_path, sub_dir)):
os.makedirs(os.path.join(dst_path, sub_dir))
print('make dir:{}'.format(os.path.join(dst_path, sub_dir)))
for i in tqdm(range(len(imgs_nm))):
nm_head = imgs_nm[i].split('_')[0]
#print(nm_head)
img_path = os.path.join(imgs_root, imgs_nm[i])
mask_path = os.path.join(masks_root, masks_nm[i])
img = cv2.imread(img_path)
mask = cv2.imread(mask_path)
gray = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
ret, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
mask_, contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
for j, cnt in enumerate(contours):
area = cv2.contourArea(cnt)
if area < 0:
continue
epsilon = 0.005 * cv2.arcLength(cnt, True)
approx = cv2.approxPolyDP(cnt, epsilon, True)
approx = np.squeeze(approx)
#print('approx_len:{}'.format(len(approx)))
if len(approx) != 4:
continue
arbitrary_points = np.float32(order_points(approx))
#print(arbitrary_points)
h, w = arbitrary_points[0][1] - arbitrary_points[1][1], \
arbitrary_points[3][0] - arbitrary_points[0][0]
rectangle = np.float32([[0, h], [0, 0], [w, 0], [w, h]])
M = cv2.getPerspectiveTransform(arbitrary_points, rectangle)
doc_dst = cv2.warpPerspective(img, M, (w, h))
mask_dst = cv2.warpPerspective(mask, M, (w, h))
padding = -1
h, w = doc_dst.shape[:2]
doc_dst = doc_dst[-padding:h+padding, -padding:w+padding]
mask_dst = mask_dst[-padding:h+padding, -padding:w+padding]
cv2.imwrite(os.path.join(dst_path, dst_sub_dirs[0], nm_head + '_doc_{}.jpg'.format(j)), doc_dst)
cv2.imwrite(os.path.join(dst_path, dst_sub_dirs[1], nm_head + '_mask_{}.png'.format(j)), mask_dst)
# cv2.imshow('doc_dst', doc_dst)
# cv2.imshow('mask_dst', mask_dst)
# cv2.waitKey(0)
if __name__=='__main__':
imgs_root = r'G:\ICDAR\TrainVal\SBR-Doc-Dataset-V2\input'
masks_root = r'G:\ICDAR\TrainVal\SBR-Doc-Dataset-V2\edge_segmentation_gt'
dst_path = r'G:\ICDAR\Gen\csdn'
dst_sub_dirs = ['docs', 'edge']
get_doc_area(imgs_root, masks_root, dst_path, dst_sub_dirs)
Inference
1.矩形框四个顶点顺时针排序,https://blog.csdn.net/qq_36560894/article/details/113092614
2.cv2.findContours,https://blog.csdn.net/hjxu2016/article/details/77833336/
3.cv2.approxPolyDP,https://vimsky.com/examples/detail/python-method-cv2.approxPolyDP.html
4.仿射变换,https://www.cnblogs.com/pheasant/p/13226182.html
交流
欢迎大家进群交流学习