292 lines
11 KiB
Python
292 lines
11 KiB
Python
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
import copy
|
||
|
|
import torch
|
||
|
|
import torch
|
||
|
|
import itertools
|
||
|
|
import torch.nn as nn
|
||
|
|
from torch.autograd import Function, Variable
|
||
|
|
|
||
|
|
def reorder(myPoints):
|
||
|
|
myPoints = myPoints.reshape((4, 2))
|
||
|
|
myPointsNew = np.zeros((4, 1, 2), dtype=np.int32)
|
||
|
|
add = myPoints.sum(1)
|
||
|
|
myPointsNew[0] = myPoints[np.argmin(add)]
|
||
|
|
myPointsNew[3] =myPoints[np.argmax(add)]
|
||
|
|
diff = np.diff(myPoints, axis=1)
|
||
|
|
myPointsNew[1] =myPoints[np.argmin(diff)]
|
||
|
|
myPointsNew[2] = myPoints[np.argmax(diff)]
|
||
|
|
return myPointsNew
|
||
|
|
|
||
|
|
|
||
|
|
def findMiddle(corners,mask,points=[0.25,0.5,0.75]):
|
||
|
|
num_middle_points = len(points)
|
||
|
|
top = [np.array([])]*num_middle_points
|
||
|
|
bottom = [np.array([])]*num_middle_points
|
||
|
|
left = [np.array([])]*num_middle_points
|
||
|
|
right = [np.array([])]*num_middle_points
|
||
|
|
|
||
|
|
center_top = []
|
||
|
|
center_bottom = []
|
||
|
|
center_left = []
|
||
|
|
center_right = []
|
||
|
|
|
||
|
|
center = (int((corners[0][0][1]+corners[3][0][1])/2),int((corners[0][0][0]+corners[3][0][0])/2))
|
||
|
|
for ratio in points:
|
||
|
|
|
||
|
|
center_top.append( (center[0],int(corners[0][0][0]*(1-ratio)+corners[1][0][0]*ratio)) )
|
||
|
|
|
||
|
|
center_bottom.append( (center[0],int(corners[2][0][0]*(1-ratio)+corners[3][0][0]*ratio)) )
|
||
|
|
|
||
|
|
center_left.append( (int(corners[0][0][1]*(1-ratio)+corners[2][0][1]*ratio),center[1]) )
|
||
|
|
|
||
|
|
center_right.append( (int(corners[1][0][1]*(1-ratio)+corners[3][0][1]*ratio),center[1]) )
|
||
|
|
|
||
|
|
for i in range(0,center[0],1):
|
||
|
|
for j in range(num_middle_points):
|
||
|
|
if top[j].size==0:
|
||
|
|
if mask[i,center_top[j][1]]==255:
|
||
|
|
top[j] = np.asarray([center_top[j][1],i])
|
||
|
|
top[j] = top[j].reshape(1,2)
|
||
|
|
|
||
|
|
for i in range(mask.shape[0]-1,center[0],-1):
|
||
|
|
for j in range(num_middle_points):
|
||
|
|
if bottom[j].size==0:
|
||
|
|
if mask[i,center_bottom[j][1]]==255:
|
||
|
|
bottom[j] = np.asarray([center_bottom[j][1],i])
|
||
|
|
bottom[j] = bottom[j].reshape(1,2)
|
||
|
|
|
||
|
|
for i in range(mask.shape[1]-1,center[1],-1):
|
||
|
|
for j in range(num_middle_points):
|
||
|
|
if right[j].size==0:
|
||
|
|
if mask[center_right[j][0],i]==255:
|
||
|
|
right[j] = np.asarray([i,center_right[j][0]])
|
||
|
|
right[j] = right[j].reshape(1,2)
|
||
|
|
|
||
|
|
for i in range(0,center[1]):
|
||
|
|
for j in range(num_middle_points):
|
||
|
|
if left[j].size==0:
|
||
|
|
if mask[center_left[j][0],i]==255:
|
||
|
|
left[j] = np.asarray([i,center_left[j][0]])
|
||
|
|
left[j] = left[j].reshape(1,2)
|
||
|
|
|
||
|
|
return np.asarray(top+bottom+left+right)
|
||
|
|
|
||
|
|
def DP_algorithmv1(contours):
|
||
|
|
biggest = np.array([])
|
||
|
|
max_area = 0
|
||
|
|
step = 0.001
|
||
|
|
count = 0
|
||
|
|
# while biggest.size==0:
|
||
|
|
while True:
|
||
|
|
for i in contours:
|
||
|
|
# print(i.shape)
|
||
|
|
area = cv2.contourArea(i)
|
||
|
|
# print(area,cv2.arcLength(i, True))
|
||
|
|
if area > cv2.arcLength(i, True)*10:
|
||
|
|
peri = cv2.arcLength(i, True)
|
||
|
|
approx = cv2.approxPolyDP(i, (0.01+step*count) * peri, True)
|
||
|
|
if area > max_area and len(approx) == 4:
|
||
|
|
max_area = area
|
||
|
|
biggest_contours = i
|
||
|
|
biggest = approx
|
||
|
|
break
|
||
|
|
if abs(max_area - cv2.contourArea(biggest))/max_area > 0.3:
|
||
|
|
biggest = np.array([])
|
||
|
|
count += 1
|
||
|
|
if count > 200:
|
||
|
|
break
|
||
|
|
temp = biggest[0]
|
||
|
|
return biggest,max_area, biggest_contours
|
||
|
|
|
||
|
|
def DP_algorithm(contours):
|
||
|
|
biggest = np.array([])
|
||
|
|
max_area = 0
|
||
|
|
step = 0.001
|
||
|
|
count = 0
|
||
|
|
|
||
|
|
### largest contours
|
||
|
|
for i in contours:
|
||
|
|
area = cv2.contourArea(i)
|
||
|
|
if area > max_area:
|
||
|
|
max_area = area
|
||
|
|
biggest_contours = i
|
||
|
|
peri = cv2.arcLength(biggest_contours, True)
|
||
|
|
|
||
|
|
### find four corners
|
||
|
|
while True:
|
||
|
|
approx = cv2.approxPolyDP(biggest_contours, (0.01+step*count) * peri, True)
|
||
|
|
if len(approx) == 4:
|
||
|
|
biggest = approx
|
||
|
|
break
|
||
|
|
# if abs(max_area - cv2.contourArea(biggest))/max_area > 0.2:
|
||
|
|
# if abs(max_area - cv2.contourArea(biggest))/max_area > 0.4:
|
||
|
|
# biggest = np.array([])
|
||
|
|
count += 1
|
||
|
|
if count > 200:
|
||
|
|
break
|
||
|
|
return biggest,max_area, biggest_contours
|
||
|
|
|
||
|
|
def drawRectangle(img,biggest,color,thickness):
|
||
|
|
cv2.line(img, (biggest[0][0][0], biggest[0][0][1]), (biggest[1][0][0], biggest[1][0][1]), color, thickness)
|
||
|
|
cv2.line(img, (biggest[0][0][0], biggest[0][0][1]), (biggest[2][0][0], biggest[2][0][1]), color, thickness)
|
||
|
|
cv2.line(img, (biggest[3][0][0], biggest[3][0][1]), (biggest[2][0][0], biggest[2][0][1]), color, thickness)
|
||
|
|
cv2.line(img, (biggest[3][0][0], biggest[3][0][1]), (biggest[1][0][0], biggest[1][0][1]), color, thickness)
|
||
|
|
return img
|
||
|
|
|
||
|
|
def minAreaRect(contours,img):
|
||
|
|
# biggest = np.array([])
|
||
|
|
max_area = 0
|
||
|
|
for i in contours:
|
||
|
|
area = cv2.contourArea(i)
|
||
|
|
if area > max_area:
|
||
|
|
peri = cv2.arcLength(i, True)
|
||
|
|
rect = cv2.minAreaRect(i)
|
||
|
|
points = cv2.boxPoints(rect)
|
||
|
|
max_area = area
|
||
|
|
return points
|
||
|
|
|
||
|
|
def cropRectangle(img,biggest):
|
||
|
|
# print(biggest)
|
||
|
|
w = np.abs(biggest[0][0][0] - biggest[1][0][0])
|
||
|
|
h = np.abs(biggest[0][0][1] - biggest[2][0][1])
|
||
|
|
new_img = np.zeros((w,h,img.shape[-1]),dtype=np.uint8)
|
||
|
|
new_img = img[biggest[0][0][1]:biggest[0][0][1]+h,biggest[0][0][0]:biggest[0][0][0]+w]
|
||
|
|
return new_img
|
||
|
|
|
||
|
|
def cvimg2torch(img,min=0,max=1):
|
||
|
|
'''
|
||
|
|
input:
|
||
|
|
im -> ndarray uint8 HxWxC
|
||
|
|
return
|
||
|
|
tensor -> torch.tensor BxCxHxW
|
||
|
|
'''
|
||
|
|
if len(img.shape)==2:
|
||
|
|
img = np.expand_dims(img,axis=-1)
|
||
|
|
img = img.astype(float) / 255.0
|
||
|
|
img = img.transpose(2, 0, 1) # NHWC -> NCHW
|
||
|
|
img = np.expand_dims(img, 0)
|
||
|
|
img = torch.from_numpy(img).float()
|
||
|
|
return img
|
||
|
|
|
||
|
|
def torch2cvimg(tensor,min=0,max=1):
|
||
|
|
'''
|
||
|
|
input:
|
||
|
|
tensor -> torch.tensor BxCxHxW C can be 1,3
|
||
|
|
return
|
||
|
|
im -> ndarray uint8 HxWxC
|
||
|
|
'''
|
||
|
|
im_list = []
|
||
|
|
for i in range(tensor.shape[0]):
|
||
|
|
im = tensor.detach().cpu().data.numpy()[i]
|
||
|
|
im = im.transpose(1,2,0)
|
||
|
|
im = np.clip(im,min,max)
|
||
|
|
im = ((im-min)/(max-min)*255).astype(np.uint8)
|
||
|
|
im_list.append(im)
|
||
|
|
return im_list
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
class TPSGridGen(nn.Module):
|
||
|
|
def __init__(self, target_height, target_width, target_control_points):
|
||
|
|
'''
|
||
|
|
target_control_points -> torch.tensor num_pointx2 -1~1
|
||
|
|
source_control_points -> torch.tensor batch_size x num_point x 2 -1~1
|
||
|
|
return:
|
||
|
|
grid -> batch_size x hw x 2 -1~1
|
||
|
|
'''
|
||
|
|
super(TPSGridGen, self).__init__()
|
||
|
|
assert target_control_points.ndimension() == 2
|
||
|
|
assert target_control_points.size(1) == 2
|
||
|
|
N = target_control_points.size(0)
|
||
|
|
self.num_points = N
|
||
|
|
target_control_points = target_control_points.float()
|
||
|
|
|
||
|
|
# create padded kernel matrix
|
||
|
|
forward_kernel = torch.zeros(N + 3, N + 3)
|
||
|
|
target_control_partial_repr = self.compute_partial_repr(target_control_points, target_control_points)
|
||
|
|
forward_kernel[:N, :N].copy_(target_control_partial_repr)
|
||
|
|
forward_kernel[:N, -3].fill_(1)
|
||
|
|
forward_kernel[-3, :N].fill_(1)
|
||
|
|
forward_kernel[:N, -2:].copy_(target_control_points)
|
||
|
|
forward_kernel[-2:, :N].copy_(target_control_points.transpose(0, 1))
|
||
|
|
# compute inverse matrix
|
||
|
|
inverse_kernel = torch.inverse(forward_kernel)
|
||
|
|
|
||
|
|
# create target cordinate matrix
|
||
|
|
HW = target_height * target_width
|
||
|
|
target_coordinate = list(itertools.product(range(target_height), range(target_width)))
|
||
|
|
target_coordinate = torch.Tensor(target_coordinate) # HW x 2
|
||
|
|
Y, X = target_coordinate.split(1, dim = 1)
|
||
|
|
Y = Y * 2 / (target_height - 1) - 1
|
||
|
|
X = X * 2 / (target_width - 1) - 1
|
||
|
|
target_coordinate = torch.cat([X, Y], dim = 1) # convert from (y, x) to (x, y)
|
||
|
|
target_coordinate_partial_repr = self.compute_partial_repr(target_coordinate.to(target_control_points.device), target_control_points)
|
||
|
|
target_coordinate_repr = torch.cat([
|
||
|
|
target_coordinate_partial_repr, torch.ones(HW, 1), target_coordinate
|
||
|
|
], dim = 1)
|
||
|
|
|
||
|
|
# register precomputed matrices
|
||
|
|
self.register_buffer('inverse_kernel', inverse_kernel)
|
||
|
|
self.register_buffer('padding_matrix', torch.zeros(3, 2))
|
||
|
|
self.register_buffer('target_coordinate_repr', target_coordinate_repr)
|
||
|
|
|
||
|
|
def forward(self, source_control_points):
|
||
|
|
assert source_control_points.ndimension() == 3
|
||
|
|
assert source_control_points.size(1) == self.num_points
|
||
|
|
assert source_control_points.size(2) == 2
|
||
|
|
batch_size = source_control_points.size(0)
|
||
|
|
|
||
|
|
Y = torch.cat([source_control_points, Variable(self.padding_matrix.expand(batch_size, 3, 2))], 1)
|
||
|
|
mapping_matrix = torch.matmul(Variable(self.inverse_kernel), Y)
|
||
|
|
source_coordinate = torch.matmul(Variable(self.target_coordinate_repr), mapping_matrix)
|
||
|
|
return source_coordinate
|
||
|
|
# phi(x1, x2) = r^2 * log(r), where r = ||x1 - x2||_2
|
||
|
|
def compute_partial_repr(self, input_points, control_points):
|
||
|
|
N = input_points.size(0)
|
||
|
|
M = control_points.size(0)
|
||
|
|
pairwise_diff = input_points.view(N, 1, 2) - control_points.view(1, M, 2)
|
||
|
|
# original implementation, very slow
|
||
|
|
# pairwise_dist = torch.sum(pairwise_diff ** 2, dim = 2) # square of distance
|
||
|
|
pairwise_diff_square = pairwise_diff * pairwise_diff
|
||
|
|
pairwise_dist = pairwise_diff_square[:, :, 0] + pairwise_diff_square[:, :, 1]
|
||
|
|
repr_matrix = 0.5 * pairwise_dist * torch.log(pairwise_dist)
|
||
|
|
# fix numerical error for 0 * log(0), substitute all nan with 0
|
||
|
|
mask = repr_matrix != repr_matrix
|
||
|
|
repr_matrix.masked_fill_(mask, 0)
|
||
|
|
return repr_matrix
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
### deside wheather further process
|
||
|
|
# point_area = cv2.contourArea(np.concatenate((biggest_angle[0].reshape(1,1,2),middle[0:3],biggest_angle[1].reshape(1,1,2),middle[9:12],biggest_angle[3].reshape(1,1,2),middle[3:6][::-1],biggest_angle[2].reshape(1,1,2),middle[6:9][::-1]),axis=0))
|
||
|
|
#### 最小外接矩形
|
||
|
|
# rect = cv2.minAreaRect(contour) # 得到最小外接矩形的（中心(x,y), (宽,高), 旋转角度）
|
||
|
|
# box = cv2.boxPoints(rect) # cv2.boxPoints(rect) for OpenCV 3.x 获取最小外接矩形的4个顶点坐标
|
||
|
|
# box = np.int0(box)
|
||
|
|
# box = box.reshape((4,1,2))
|
||
|
|
# minrect_area = cv2.contourArea(box)
|
||
|
|
# print(abs(minrect_area-point_area)/point_area)
|
||
|
|
#### 四个角点 IOU
|
||
|
|
# biggest_box = np.concatenate((biggest_angle[0,:,:].reshape(1,1,2),biggest_angle[2,:,:].reshape(1,1,2),biggest_angle[3,:,:].reshape(1,1,2),biggest_angle[1,:,:].reshape(1,1,2)),axis=0)
|
||
|
|
# biggest_mask = np.zeros_like(mask)
|
||
|
|
# # corner_area = cv2.contourArea(biggest_box)
|
||
|
|
# cv2.drawContours(biggest_mask,[biggest_box], -1, color=255, thickness=-1)
|
||
|
|
|
||
|
|
# smooth = 1e-5
|
||
|
|
# biggest_mask_ = biggest_mask > 50
|
||
|
|
# mask_ = mask > 50
|
||
|
|
# intersection = (biggest_mask_ & mask_).sum()
|
||
|
|
# union = (biggest_mask_ | mask_).sum()
|
||
|
|
# iou = (intersection + smooth) / (union + smooth)
|
||
|
|
# if iou > 0.975:
|
||
|
|
# skip = True
|
||
|
|
# else:
|
||
|
|
# skip = False
|
||
|
|
# print(iou)
|
||
|
|
# cv2.imshow('mask',cv2.resize(mask,(512,512)))
|
||
|
|
# cv2.imshow('biggest_mask',cv2.resize(biggest_mask,(512,512)))
|
||
|
|
# cv2.waitKey(0)
|