HEAT / metrics /new_utils.py
Egrt's picture
init
424188c
raw
history blame
No virus
82.2 kB
import numpy as np
import matplotlib.pyplot as plt
import cv2
import threading
import os
import skimage
import random
import time
TWO_CORNER_MINIMUM_DISTANCE = 5
SAFE_NUM = 3
score_weights = (1., 2., 100.)
#########################################################################################
################################# General Functions #####################################
#########################################################################################
def swap_two_corner_place(corners, edges, id1, id2):
for edge_i in range(edges.shape[0]):
if edges[edge_i, 0] == id1:
edges[edge_i, 0] = id2
elif edges[edge_i, 0] == id2:
edges[edge_i, 0] = id1
if edges[edge_i, 1] == id1:
edges[edge_i, 1] = id2
elif edges[edge_i, 1] == id2:
edges[edge_i, 1] = id1
temp = corners[id1].copy()
corners[id1] = corners[id2]
corners[id2] = temp
return corners, edges
def get_neighbor_corner_id(corner_id, edges):
where = np.where(edges == corner_id)
return edges[where[0], 1 - where[1]]
def swap_two_edge_place(edges, id1, id2):
temp = edges[id1].copy()
edges[id1] = edges[id2]
edges[id2] = temp
return edges
def degree_of_three_corners(cornerA, cornerB, cornerM):
# cornerM is middle corner
AM_length = l2_distance(cornerA, cornerM)
BM_length = l2_distance(cornerB, cornerM)
dot = np.dot((cornerA[0] - cornerM[0], cornerA[1] - cornerM[1]),
(cornerB[0] - cornerM[0], cornerB[1] - cornerM[1]))
cos = dot / (AM_length + 1e-8) / (BM_length + 1e-8)
cos = min(1, max(-1, cos))
degree = np.arccos(cos)
return degree / np.pi * 180
def sort_graph(corners, edges):
corners = corners.copy()
edges = edges.copy()
for corner_i in range(corners.shape[0]):
min_id = -1
min_pos = corners[corner_i]
for corner_j in range(corner_i + 1, corners.shape[0]):
if (corners[corner_j, 0] < min_pos[0]) or \
(corners[corner_j, 0] == min_pos[0] and corners[corner_j, 1] < min_pos[1]):
min_pos = corners[corner_j]
min_id = corner_j
if min_id != -1:
corners, edges = swap_two_corner_place(corners, edges, corner_i, min_id)
for edge_i in range(edges.shape[0]):
if edges[edge_i, 0] > edges[edge_i, 1]:
temp = edges[edge_i, 0]
edges[edge_i, 0] = edges[edge_i, 1]
edges[edge_i, 1] = temp
for edge_i in range(edges.shape[0]):
min_id = -1
min_pos = edges[edge_i]
for edge_j in range(edge_i + 1, edges.shape[0]):
if (edges[edge_j, 0] < min_pos[0]) or \
(edges[edge_j, 0] == min_pos[0] and edges[edge_j, 1] < min_pos[1]):
min_pos = edges[edge_j]
min_id = edge_j
if min_id != -1:
edges = swap_two_edge_place(edges, edge_i, min_id)
return corners, edges
def IOU(maskA, maskB):
return np.logical_and(maskA, maskB).sum() / np.logical_or(maskA, maskB).sum()
def render(corners, edges, render_pad=0, edge_linewidth=2, corner_size=3, scale=1.):
size = int(256 * scale)
mask = np.ones((2, size, size)) * render_pad
corners = np.round(corners.copy() * scale).astype(np.int)
for edge_i in range(edges.shape[0]):
a = edges[edge_i, 0]
b = edges[edge_i, 1]
mask[0] = cv2.line(mask[0], (int(corners[a, 1]), int(corners[a, 0])),
(int(corners[b, 1]), int(corners[b, 0])), 1.0, thickness=edge_linewidth)
for corner_i in range(corners.shape[0]):
mask[1] = cv2.circle(mask[1], (int(corners[corner_i, 1]), int(corners[corner_i, 0])), corner_size, 1.0, -1)
return mask
def patch_samples(edge_num, batch_size):
num = edge_num // batch_size
patchs = []
for i in range(num):
patchs.append([i * batch_size + j for j in range(batch_size)])
if edge_num % batch_size != 0:
patchs.append([j for j in range(batch_size * num, edge_num)])
return patchs
def l2_distance(x1, x2):
return np.sqrt((x1[0] - x2[0]) ** 2 + (x1[1] - x2[1]) ** 2)
def triangle_region(A, B, C):
l1 = np.linalg.norm(np.array(A) - np.array(B))
l2 = np.linalg.norm(np.array(A) - np.array(C))
l3 = np.linalg.norm(np.array(B) - np.array(C))
p = (l1 + l2 + l3) / 2
area = np.sqrt(np.abs(p * (p - l1) * (p - l2) * (p - l3)))
return area
def remove_intersection_and_duplicate(corners, edges, name):
over_all_flag = False
ori_corners = corners.copy()
ori_edges = edges.copy()
while True:
flag = False
for edge_i in range(edges.shape[0]):
for edge_j in range(edge_i + 1, edges.shape[0]):
corner11 = corners[edges[edge_i, 0]]
corner12 = corners[edges[edge_i, 1]]
corner21 = corners[edges[edge_j, 0]]
corner22 = corners[edges[edge_j, 1]]
y1 = corner11[0]
x1 = corner11[1]
y2 = corner12[0]
x2 = corner12[1]
a1 = y1 - y2
b1 = x2 - x1
c1 = x1 * y2 - x2 * y1
flag1 = (a1 * corner21[1] + b1 * corner21[0] + c1) * (a1 * corner22[1] + b1 * corner22[0] + c1)
y1 = corner21[0]
x1 = corner21[1]
y2 = corner22[0]
x2 = corner22[1]
a2 = y1 - y2
b2 = x2 - x1
c2 = x1 * y2 - x2 * y1
flag2 = (a2 * corner11[1] + b2 * corner11[0] + c2) * (a2 * corner12[1] + b2 * corner12[0] + c2)
if flag1 < -1e-5 and flag2 < -1e-5:
# intersection!
over_all_flag = True
flag = True
new_x = (c2 * b1 - c1 * b2) / (a1 * b2 - a2 * b1)
new_y = (a2 * c1 - a1 * c2) / (a1 * b2 - a2 * b1)
temp_d = 3
temp_id = -1
if l2_distance((new_y, new_x), corner11) < temp_d:
temp_id = edges[edge_i, 0]
temp_d = l2_distance((new_y, new_x), corner11)
if l2_distance((new_y, new_x), corner12) < temp_d:
temp_id = edges[edge_i, 1]
temp_d = l2_distance((new_y, new_x), corner12)
if l2_distance((new_y, new_x), corner21) < temp_d:
temp_id = edges[edge_j, 0]
temp_d = l2_distance((new_y, new_x), corner21)
if l2_distance((new_y, new_x), corner22) < temp_d:
temp_id = edges[edge_j, 1]
temp_d = l2_distance((new_y, new_x), corner22)
if temp_id != -1:
if edges[edge_i, 0] != temp_id and edges[edge_i, 1] != temp_id:
tt = edges[edge_i, 0]
edges[edge_i, 0] = temp_id
edges = np.append(edges, np.array([(temp_id, tt)]), 0)
if edges[edge_j, 0] != temp_id and edges[edge_j, 1] != temp_id:
tt = edges[edge_j, 0]
edges[edge_j, 0] = temp_id
edges = np.append(edges, np.array([(temp_id, tt)]), 0)
else:
corners = np.append(corners, np.array([(new_y, new_x)]), 0)
edge_id1 = edges[edge_i, 1]
edge_id2 = edges[edge_j, 1]
edges[edge_i, 1] = corners.shape[0] - 1
edges[edge_j, 1] = corners.shape[0] - 1
edges = np.append(edges, np.array([(edge_id1, corners.shape[0] - 1)]), 0)
edges = np.append(edges, np.array([(edge_id2, corners.shape[0] - 1)]), 0)
break
if flag:
break
if flag:
continue
break
# remove duplicate and zero degree
graph = Graph(np.round(corners), edges)
for corner_i in reversed(range(len(graph.getCorners()))):
corner_ele1 = graph.getCorners()[corner_i]
for corner_j in reversed(range(corner_i)):
corner_ele2 = graph.getCorners()[corner_j]
if l2_distance(corner_ele1.x, corner_ele2.x) < 3:
connected_edge = graph.getEdgeConnected(corner_ele1)
for edge_ele in connected_edge:
if edge_ele.x[0] == corner_ele1:
another = edge_ele.x[1]
else:
another = edge_ele.x[0]
if another == corner_ele2:
graph.remove(edge_ele)
edge_ele.x = (another, corner_ele2)
graph.remove(corner_ele1)
for corner_ele in graph.getCorners():
if graph.getCornerDegree(corner_ele) == 0:
graph.remove(corner_ele)
corners = graph.getCornersArray()
edges = graph.getEdgesArray()
# if over_all_flag:
# plt.subplot(121)
# ori = render(ori_corners, ori_edges, edge_linewidth=1, corner_size=1)
# temp = np.concatenate((ori.transpose((1,2,0)), np.zeros((ori.shape[1],ori.shape[2],1))),2)
# plt.imshow(temp)
# plt.subplot(122)
# new_ = render(corners, edges, edge_linewidth=1, corner_size=1)
# temp = np.concatenate((new_.transpose((1,2,0)), np.zeros((new_.shape[1],new_.shape[2],1))),2)
# plt.imshow(temp)
# plt.show()
return corners, edges
def get_two_edge_intersection_location(corner11, corner12, corner21, corner22):
y1 = corner11[0]
x1 = corner11[1]
y2 = corner12[0]
x2 = corner12[1]
a1 = y1 - y2
b1 = x2 - x1
c1 = x1 * y2 - x2 * y1
y1 = corner21[0]
x1 = corner21[1]
y2 = corner22[0]
x2 = corner22[1]
a2 = y1 - y2
b2 = x2 - x1
c2 = x1 * y2 - x2 * y1
l = a1 * b2 - a2 * b1
if l == 0:
l = 1e-5
new_x = (c2 * b1 - c1 * b2) / l
new_y = (a2 * c1 - a1 * c2) / l
return round(new_y), round(new_x)
def get_distance_of_corner_and_edge(corner1, corner2, corner):
x = corner[0]
y = corner[1]
x1 = corner1[0]
y1 = corner1[1]
x2 = corner2[0]
y2 = corner2[1]
cross = (x2 - x1) * (x - x1) + (y2 - y1) * (y - y1)
if cross <= 0:
# dist to corner1
return np.linalg.norm((x - x1, y - y1))
d2 = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1)
if cross >= d2:
# dist to corner2
return np.linalg.norm((x - x2, y - y2))
r = cross / d2
px = x1 + (x2 - x1) * r
py = y1 + (y2 - y1) * r
return np.linalg.norm((x - px, y - py))
#########################################################################################
################################# Dataset Functions #####################################
#########################################################################################
def EuclideanDistance(A, B):
BT = B.transpose()
vecProd = np.dot(A, BT)
SqA = A ** 2
sumSqA = np.matrix(np.sum(SqA, axis=1))
sumSqAEx = np.tile(sumSqA.transpose(), (1, vecProd.shape[1]))
SqB = B ** 2
sumSqB = np.sum(SqB, axis=1)
sumSqBEx = np.tile(sumSqB, (vecProd.shape[0], 1))
SqED = sumSqBEx + sumSqAEx - 2 * vecProd
SqED[SqED < 0] = 0.0
ED = np.sqrt(SqED)
return ED
def samedirection(conv_corner_id, gt_corner_id, conv_corners, gt_corners, conv_edges, gt_edges):
# degree
if np.where(conv_edges == conv_corner_id)[0].shape[0] != np.where(gt_edges == gt_corner_id)[0].shape[0]:
return False
# direction
place = np.where(conv_edges == conv_corner_id)
neighbor_id = conv_edges[place[0], 1 - place[1]]
distance = conv_corners[conv_corner_id] - conv_corners[neighbor_id]
direction = np.arctan2(distance[:, 0], distance[:, 1]) * 180 / np.pi / 15
direction = (direction + 24) % 24
conv_dir = np.sort(direction)
place = np.where(gt_edges == gt_corner_id)
neighbor_id = gt_edges[place[0], 1 - place[1]]
distance = gt_corners[gt_corner_id] - gt_corners[neighbor_id]
direction = np.arctan2(distance[:, 0], distance[:, 1]) * 180 / np.pi / 15
direction = (direction + 24) % 24
gt_dir = np.sort(direction)
conv_dir = list(conv_dir)
gt_dir = list(gt_dir)
for angle in gt_dir:
temp = sorted(conv_dir, key=lambda x: min(np.abs(x - angle), 24 - np.abs(x - angle)))
if min(np.abs(temp[0] - angle), 24 - np.abs(temp[0] - angle)) <= 1.3:
conv_dir.remove(temp[0])
else:
return False
return True
def simplify_gt(gt_match_location, gt_corner, gt_edge):
graph = Graph(np.round(gt_corner), gt_edge)
for idx, corner in enumerate(graph.getCorners()):
# use score to store the matching info
corner.store_score(gt_match_location[idx])
for idx, corner in enumerate(graph.getCorners()):
if corner.get_score() is None:
connected_edges = graph.getEdgeConnected(corner)
neighbor_corners = []
for edge in connected_edges:
if edge.x[0] != corner:
neighbor_corners.append(edge.x[0])
continue
if edge.x[1] != corner:
neighbor_corners.append(edge.x[1])
continue
raise BaseException()
neighbor_corners = sorted(neighbor_corners, key=lambda ele: l2_distance(ele.x, corner.x))
for neighbor_ele in neighbor_corners:
if l2_distance(neighbor_ele.x, corner.x) > 8:
break
if neighbor_ele.get_score() is None:
continue
# find the suitable neighbor that replace corner
for ele in neighbor_corners:
if ele == neighbor_ele:
continue
graph.add_edge(ele, neighbor_ele)
neighbor_ele.x = (0.7 * neighbor_ele.x[0] + 0.3 * corner.x[0],
0.7 * neighbor_ele.x[1] + 0.3 * corner.x[1])
graph.remove(corner)
break
return graph.getCornersArray(), graph.getEdgesArray()
def get_wrong_corners(corners, gt_corners, edges, gt_edges):
corners = corners.copy()
gt_corners = gt_corners.copy()
edges = edges.copy()
gt_edges = gt_edges.copy()
dist_matrix = EuclideanDistance(gt_corners, corners)
assigned_id = set()
gt_match_same_degree = []
gt_match_location = []
for gt_i in range(gt_corners.shape[0]):
sort_id = np.argsort(dist_matrix[gt_i]).__array__()[0]
flag = True
for id_ in sort_id:
if dist_matrix[gt_i, id_] > 7:
break
temete = samedirection(id_, gt_i, corners, gt_corners, edges, gt_edges)
if temete == False:
break
elif id_ not in assigned_id:
assigned_id.add(id_)
gt_match_same_degree.append(id_)
flag = False
break
if flag:
gt_match_same_degree.append(None)
matched = []
gt_match_location = [None for _ in range(gt_corners.shape[0])]
for gt_i in sorted(list(range(gt_corners.shape[0])), key=lambda i: np.min(dist_matrix[i])):
sort_id = np.argsort(dist_matrix[gt_i]).__array__()[0]
if dist_matrix[gt_i, sort_id[0]] > 7:
gt_match_location[gt_i] = None
else:
for c_i in sort_id:
if c_i in matched:
continue
if dist_matrix[gt_i, c_i] > 7:
gt_match_location[gt_i] = None
break
else:
gt_match_location[gt_i] = c_i
matched.append(c_i)
break
return set(range(corners.shape[0])) - assigned_id, gt_match_same_degree, gt_match_location
def get_wrong_edges(corners, gt_corners, edges, gt_edges, gt_match):
edges = edges.copy()
gt_edges = gt_edges.copy()
all_possible_good_edges = []
for edge_i in range(gt_edges.shape[0]):
if gt_match[gt_edges[edge_i, 0]] is None or gt_match[gt_edges[edge_i, 1]] is None:
continue
all_possible_good_edges.append((gt_match[gt_edges[edge_i, 0]], gt_match[gt_edges[edge_i, 1]]))
false_edge_id = []
for edge_i in range(edges.shape[0]):
id1 = edges[edge_i][0]
id2 = edges[edge_i][1]
if (id1, id2) not in all_possible_good_edges and (id2, id1) not in all_possible_good_edges:
false_edge_id.append(edge_i)
continue
return false_edge_id
def get_corner_bin_map(corners, corner_list_for_each_bin, bin_size=10):
bin_map = np.zeros((bin_size, 256, 256))
for bin_i in range(bin_size):
bin_map[bin_i] = render(corners[corner_list_for_each_bin[bin_i]], np.array([]), render_pad=0)[1]
return bin_map
#########################################################################################
################################ Searching Functions ####################################
#########################################################################################
def visualization(candidate, show=True):
corners = candidate.graph.getCornersArray()
edges = candidate.graph.getEdgesArray()
mask = render(corners, edges)
mask = np.transpose(np.concatenate((mask, np.zeros((1, 256, 256))), 0), (1, 2, 0))
plt.imshow(mask)
if show:
plt.show()
def check_intersection(edge1, edge2):
corner11 = edge1.x[0].x
corner12 = edge1.x[1].x
corner21 = edge2.x[0].x
corner22 = edge2.x[1].x
y1 = corner11[0]
x1 = corner11[1]
y2 = corner12[0]
x2 = corner12[1]
a = y1 - y2
b = x2 - x1
c = x1 * y2 - x2 * y1
flag1 = (a * corner21[1] + b * corner21[0] + c) * (a * corner22[1] + b * corner22[0] + c)
y1 = corner21[0]
x1 = corner21[1]
y2 = corner22[0]
x2 = corner22[1]
a = y1 - y2
b = x2 - x1
c = x1 * y2 - x2 * y1
flag2 = (a * corner11[1] + b * corner11[0] + c) * (a * corner12[1] + b * corner12[0] + c)
if flag1 < -1e-6 and flag2 < -1e-6:
return True
return False
def adding_a_corner_by_triangle_operation(candidate):
new_candidates = []
name = candidate.name
gt_mask = region_cache.get_region(name)
gt_mask = gt_mask > 0.4
gt_mask_grow = cv2.dilate(gt_mask.astype(np.float64), np.ones((3, 3), np.uint8), iterations=6) > 0
# get the current candidate region mask
conv_mask = render(corners=candidate.graph.getCornersArray(), edges=candidate.graph.getEdgesArray(),
render_pad=0, edge_linewidth=1)[0]
conv_mask = 1 - conv_mask
conv_mask = conv_mask.astype(np.uint8)
labels, region_mask = cv2.connectedComponents(conv_mask, connectivity=4)
background_label = region_mask[0, 0]
all_masks = []
for region_i in range(1, labels):
if region_i == background_label:
continue
the_region = region_mask == region_i
if the_region.sum() < 20:
continue
all_masks.append(the_region)
candidate_mask = (np.sum(all_masks, 0) + (1 - conv_mask)) > 0
final_mask = np.logical_xor(gt_mask_grow, np.logical_and(candidate_mask, gt_mask_grow))
for corner_i in range(random.randint(0, 16), 256, 16):
for corner_j in range(random.randint(0, 16), 256, 16):
if candidate.addable((corner_i, corner_j)):
if final_mask[corner_i, corner_j] == True: # inside the region
new_corner = Element((corner_i, corner_j))
new_candidate = candidate.generate_new_candidate_add_a_corner(new_corner)
new_graph = new_candidate.graph
corners = new_graph.getCorners()
# find two suitable existed corners to make into a triangle (no intersection and no colinear)
for id_A in range(len(corners)):
ele_A = corners[id_A]
if ele_A == new_corner:
continue
for id_B in range(id_A + 1, len(corners)):
ele_B = corners[id_B]
if ele_B == new_corner:
continue
if new_graph.has_edge(new_corner, ele_A) is not None:
raise BaseException('should not have edge in this case')
if new_graph.has_edge(new_corner, ele_B) is not None:
raise BaseException('should not have edge in this case')
temp_edge1 = Element((new_corner, ele_A))
temp_edge2 = Element((new_corner, ele_B))
# check if addable
if new_candidate.addable(temp_edge1) is False:
continue
if new_candidate.addable(temp_edge2) is False:
continue
# avoid intersection
if new_graph.checkIntersectionEdge(temp_edge1):
continue
if new_graph.checkIntersectionEdge(temp_edge2):
continue
# avoid too small triangle
if triangle_region(new_corner.x, ele_A.x, ele_B.x) < 20:
continue
### avoid colinear edge (only when fold case)
# for edge1
neighbor_edges = new_graph.getEdgeConnected(temp_edge1)
flag_ = True
for neighbor in neighbor_edges:
if new_corner in neighbor.x:
raise BaseException('new corner should not in any edge')
elif ele_A in neighbor.x:
shared_corner = ele_A
else:
raise BaseException('error.')
two_neighbor = {neighbor.x[0], neighbor.x[1], ele_A, new_corner}
two_neighbor.remove(shared_corner)
assert len(two_neighbor) == 2
two_neighbor = tuple(two_neighbor)
line1 = np.array(shared_corner.x) - np.array(two_neighbor[0].x)
line2 = np.array(shared_corner.x) - np.array(two_neighbor[1].x)
cos = np.dot(line1, line2) / (np.linalg.norm(line1) * np.linalg.norm(line2))
cos = min(1, max(-1, cos))
if np.arccos(cos) < np.pi / 9: # 20 degree
flag_ = False
break
if flag_ is False:
continue
# for edge2
neighbor_edges = new_graph.getEdgeConnected(temp_edge2)
flag_ = True
for neighbor in neighbor_edges:
if new_corner in neighbor.x:
raise BaseException('new corner should not in any edge')
elif ele_B in neighbor.x:
shared_corner = ele_B
else:
raise BaseException('error.')
two_neighbor = {neighbor.x[0], neighbor.x[1], ele_B, new_corner}
two_neighbor.remove(shared_corner)
assert len(two_neighbor) == 2
two_neighbor = tuple(two_neighbor)
line1 = np.array(shared_corner.x) - np.array(two_neighbor[0].x)
line2 = np.array(shared_corner.x) - np.array(two_neighbor[1].x)
cos = np.dot(line1, line2) / (np.linalg.norm(line1) * np.linalg.norm(line2))
cos = min(1, max(-1, cos))
if np.arccos(cos) < np.pi / 9: # 20 degree
flag_ = False
break
if flag_ is False:
continue
# make new candidate
try:
new_ = new_candidate.generate_new_candidate_add_an_edge(new_corner, ele_A)
new_ = new_.generate_new_candidate_add_an_edge(new_corner, ele_B)
new_candidates.append(new_)
except:
continue
# plt.subplot(151)
# visualization(candidate, show=False)
# plt.subplot(152)
# plt.imshow(final_mask)
# plt.subplot(153)
# plt.imshow(candidate_mask)
# plt.subplot(154)
# plt.imshow(gt_mask_grow)
# plt.subplot(155)
# visualization(new_, show=False)
# plt.show()
return new_candidates
def adding_an_edge_from_new_corner_operation(candidate):
new_candidates = []
name = candidate.name
gt_mask = region_cache.get_region(name)
gt_mask = gt_mask > 0.4
gt_mask_grow = cv2.dilate(gt_mask.astype(np.float64), np.ones((3, 3), np.uint8), iterations=6) > 0
# get the current candidate region mask
conv_mask = render(corners=candidate.graph.getCornersArray(), edges=candidate.graph.getEdgesArray(),
render_pad=0, edge_linewidth=1)[0]
conv_mask = 1 - conv_mask
conv_mask = conv_mask.astype(np.uint8)
labels, region_mask = cv2.connectedComponents(conv_mask, connectivity=4)
background_label = region_mask[0, 0]
all_masks = []
for region_i in range(1, labels):
if region_i == background_label:
continue
the_region = region_mask == region_i
if the_region.sum() < 20:
continue
all_masks.append(the_region)
candidate_mask = (np.sum(all_masks, 0) + (1 - conv_mask)) > 0
final_mask = np.logical_xor(gt_mask_grow, np.logical_and(candidate_mask, gt_mask_grow))
for corner_i in range(random.randint(0, 16), 256, 16):
for corner_j in range(random.randint(0, 16), 256, 16):
if candidate.addable((corner_i, corner_j)):
if final_mask[corner_i, corner_j] == True:
# inside the region
new_corner = Element((corner_i, corner_j))
new_candidate = candidate.generate_new_candidate_add_a_corner(new_corner)
new_graph = new_candidate.graph
corners = new_graph.getCorners()
# find a suitable existed corner that can make
# a new edge with new_corner (no intersection and colinear)
for corner_ele in corners:
if corner_ele == new_corner:
continue
if new_graph.has_edge(new_corner, corner_ele) is not None:
raise BaseException('should not have edge in this case')
temp_edge = Element((new_corner, corner_ele))
# check if addable
if new_candidate.addable(temp_edge) is False:
continue
# avoid intersection
if new_graph.checkIntersectionEdge(temp_edge):
continue
# avoid colinear edge
neighbor_edges = new_graph.getEdgeConnected(temp_edge)
flag_ = True
for neighbor in neighbor_edges:
if new_corner in neighbor.x:
raise BaseException('new corner should not in any edge')
elif corner_ele in neighbor.x:
shared_corner = corner_ele
else:
raise BaseException('error.')
two_neighbor = {neighbor.x[0], neighbor.x[1], corner_ele, new_corner}
two_neighbor.remove(shared_corner)
assert len(two_neighbor) == 2
two_neighbor = tuple(two_neighbor)
line1 = np.array(shared_corner.x) - np.array(two_neighbor[0].x)
line2 = np.array(shared_corner.x) - np.array(two_neighbor[1].x)
cos = np.dot(line1, line2) / (np.linalg.norm(line1) * np.linalg.norm(line2))
cos = min(1, max(-1, cos))
if np.arccos(cos) < np.pi / 9: # 20 degree
flag_ = False
break
if flag_ is False:
continue
# make new candidate
try:
new_ = new_candidate.generate_new_candidate_add_an_edge(new_corner, corner_ele)
new_candidates.append(new_)
except:
continue
return new_candidates
def removing_a_corner_operation(candidate):
new_candidates = []
graph = candidate.graph
corners = graph.getCorners()
for the_corner in corners:
if candidate.removable(the_corner):
try:
new_ = candidate.generate_new_candidate_remove_a_corner(the_corner)
new_candidates.append(new_)
except:
continue
return new_candidates
def removing_a_colinear_corner_operation(candidate):
new_candidates = []
graph = candidate.graph
corners = graph.getCorners()
for the_corner in corners:
if candidate.removable(the_corner): # NO NEED TO CHECK IF COLINEAR and graph.checkColinearCorner(the_corner):
try:
new_ = candidate.generate_new_candidate_remove_a_colinear_corner(the_corner)
if new_.graph.checkIntersectionEdge():
continue
new_candidates.append(new_)
except:
continue
return new_candidates
def adding_an_edge_operation(candidate):
new_candidates = []
graph = candidate.graph
corners = graph.getCorners()
for corner_i in range(len(corners)):
cornerA = corners[corner_i]
for corner_j in range(corner_i + 1, len(corners)):
cornerB = corners[corner_j]
if graph.has_edge(cornerA, cornerB) is not None:
continue
temp_edge = Element((cornerA, cornerB))
# check if addable (not in existed_before dict)
if candidate.addable(temp_edge) is False:
continue
if graph.checkIntersectionEdge(temp_edge):
continue
# avoid adding a colinear edge
neighbor_edges = graph.getEdgeConnected(temp_edge)
flag_ = True
for neighbor in neighbor_edges:
if cornerA in neighbor.x:
shared_corner = cornerA
elif cornerB in neighbor.x:
shared_corner = cornerB
else:
raise BaseException('error.')
two_neighbor = {neighbor.x[0], neighbor.x[1], cornerA, cornerB}
two_neighbor.remove(shared_corner)
assert len(two_neighbor) == 2
two_neighbor = tuple(two_neighbor)
line1 = np.array(shared_corner.x) - np.array(two_neighbor[0].x)
line2 = np.array(two_neighbor[1].x) - np.array(shared_corner.x)
cos = np.dot(line1, line2) / (np.linalg.norm(line1) * np.linalg.norm(line2))
cos = min(1, max(-1, cos))
if np.arccos(cos) < np.pi / 18 or np.arccos(cos) > np.pi - np.pi / 18: # 10 degree
flag_ = False
break
if flag_ is False:
continue
# make new candidate
try:
new_ = candidate.generate_new_candidate_add_an_edge(cornerA, cornerB)
new_candidates.append(new_)
except:
continue
return new_candidates
def removing_an_edge_operation(candidate):
new_candidates = []
graph = candidate.graph
edges = graph.getEdges()
for edge_ele in edges:
if candidate.removable(edge_ele):
try:
new_ = candidate.generate_new_candidate_remove_an_edge(edge_ele)
new_candidates.append(new_)
except:
continue
return new_candidates
def adding_an_edge_from_gt(candidate, gt_data):
new_candidates = []
corners_array = candidate.graph.getCornersArray()
edges_array = candidate.graph.getEdgesArray()
gt_corners = gt_data['corners'].copy()
gt_edges = gt_data['edges'].copy()
_, _, map_same_location = get_wrong_corners(
corners_array, gt_corners, edges_array, gt_edges)
gt_corners, gt_edges = simplify_gt(map_same_location, gt_corners, gt_edges)
_, _, map_same_location = get_wrong_corners(
corners_array, gt_corners, edges_array, gt_edges)
for corner_i in range(gt_corners.shape[0]):
if map_same_location[corner_i] is None:
# doesn't exist in candidate
neighbor_id = get_neighbor_corner_id(corner_i, gt_edges)
for corner_j in neighbor_id:
if map_same_location[corner_j] is not None:
# exist corner in candidate that maps neighbor corner
new_candidate = candidate.copy()
new_corner = Element(
(
int(np.round(gt_corners[corner_i, 0])), int(np.round(gt_corners[corner_i, 1]))
)
)
if new_candidate.addable(new_corner) is False:
continue
# new corner can be too close to an edge
flag = False
for edge_ele in new_candidate.graph.getEdges():
if get_distance_of_corner_and_edge(edge_ele.x[0].x, edge_ele.x[1].x, new_corner.x) < 7:
flag = True
break
if flag:
continue
new_corner = new_candidate.addCorner(new_corner)
neighbor_index = map_same_location[corner_j]
neighbor_corner = new_candidate.graph.getCorners()[neighbor_index]
new_edge = new_candidate.addEdge(new_corner, neighbor_corner)
if new_candidate.graph.checkIntersectionEdge(new_edge):
continue
new_candidates.append(new_candidate)
return new_candidates
def adding_a_corner_from_two_edges_extension(candidate):
new_candidates = []
graph = candidate.graph
edges = candidate.graph.getEdges()
for edge_i in range(len(edges)):
for edge_j in range(edge_i + 1, len(edges)):
edgeA = edges[edge_i]
edgeB = edges[edge_j]
if graph.isNeighbor(edgeA, edgeB):
continue
intersection_loc = get_two_edge_intersection_location(edgeA.x[0].x, edgeA.x[1].x, edgeB.x[0].x,
edgeB.x[1].x)
if intersection_loc[0] >= 255 or intersection_loc[1] >= 255 or \
intersection_loc[0] <= 0 or intersection_loc[1] <= 0:
continue
# intersection point can not be too close to an edge
flag = False
for edge_ele in graph.getEdges():
if get_distance_of_corner_and_edge(edge_ele.x[0].x, edge_ele.x[1].x, intersection_loc) < 7:
flag = True
break
if flag:
continue
new_candidate = candidate.copy()
new_graph = new_candidate.graph
new_edgeA = new_graph.getRealElement(edgeA)
new_edgeB = new_graph.getRealElement(edgeB)
new_corner = Element(intersection_loc)
if new_candidate.addable(new_corner) is False:
continue
new_corner = new_candidate.addCorner_v2(new_corner)
# get cornerA and cornerB from edgeA, edgeB
if l2_distance(new_corner.x, new_edgeA.x[0].x) < l2_distance(new_corner.x, new_edgeA.x[1].x):
cornerA = new_edgeA.x[0]
else:
cornerA = new_edgeA.x[1]
if l2_distance(new_corner.x, new_edgeB.x[0].x) < l2_distance(new_corner.x, new_edgeB.x[1].x):
cornerB = new_edgeB.x[0]
else:
cornerB = new_edgeB.x[1]
# new edge can not be too short
if l2_distance(cornerA.x, new_corner.x) < 7:
continue
if l2_distance(cornerB.x, new_corner.x) < 7:
continue
# new intersection cannot be too flat
if degree_of_three_corners(cornerA.x, cornerB.x, new_corner.x) > 165:
continue
flag = False
for edge_ele in new_graph.getEdges():
if new_corner in edge_ele.x and cornerA in edge_ele.x:
flag = True
break
if edge_ele.x[0] not in (new_corner, cornerA):
l = get_distance_of_corner_and_edge(new_corner.x, cornerA.x, edge_ele.x[0].x)
if l <= 7:
flag = True
break
if edge_ele.x[1] not in (new_corner, cornerA):
l = get_distance_of_corner_and_edge(new_corner.x, cornerA.x, edge_ele.x[1].x)
if l <= 7:
flag = True
break
if flag:
continue
add_edgeA = new_candidate.addEdge(new_corner, cornerA)
if new_graph.checkIntersectionEdge(add_edgeA):
continue
flag = False
for edge_ele in new_graph.getEdges():
if new_corner in edge_ele.x and cornerB in edge_ele.x:
flag = True
break
if edge_ele.x[0] not in (new_corner, cornerB):
l = get_distance_of_corner_and_edge(new_corner.x, cornerB.x, edge_ele.x[0].x)
if l <= 7:
flag = True
break
if edge_ele.x[1] not in (new_corner, cornerB):
l = get_distance_of_corner_and_edge(new_corner.x, cornerB.x, edge_ele.x[1].x)
if l <= 7:
flag = True
break
if flag:
continue
add_edgeB = new_candidate.addEdge(new_corner, cornerB)
if new_graph.checkIntersectionEdge(add_edgeB):
continue
# make real new candidate
# new_candidate = candidate.copy()
# new_graph = new_candidate.graph
# new_corner = Element(intersection_loc)
# new_corner = new_graph.add_corner_v2(new_corner)
# new_candidate = new_candidate.generate_new_candidate_add_an_edge(new_corner, cornerA)
# new_candidate = new_candidate.generate_new_candidate_add_an_edge(new_corner, cornerB)
new_candidates.append(new_candidate)
return new_candidates
def adding_a_corner_from_parallel(candidate):
new_candidates = []
graph = candidate.graph
edges = candidate.graph.getEdges()
for edge_i in range(len(edges)):
for edge_j in range(edge_i + 1, len(edges)):
edgeA = edges[edge_i]
edgeB = edges[edge_j]
# get intersection loc
if graph.isNeighbor(edgeA, edgeB):
shared_corner = edgeA.x[0] if edgeA.x[0] in edgeB.x else edgeA.x[1]
intersection_loc = shared_corner.x
else:
intersection_loc = get_two_edge_intersection_location(
edgeA.x[0].x, edgeA.x[1].x, edgeB.x[0].x, edgeB.x[1].x)
if intersection_loc[0] >= 255 or intersection_loc[1] >= 255 or \
intersection_loc[0] <= 0 or intersection_loc[1] <= 0:
continue
# get another two loc
locA = edgeA.x[1].x if \
l2_distance(edgeA.x[0].x, intersection_loc) < l2_distance(edgeA.x[1].x, intersection_loc) else \
edgeA.x[0].x
locB = edgeB.x[1].x if \
l2_distance(edgeB.x[0].x, intersection_loc) < l2_distance(edgeB.x[1].x, intersection_loc) else \
edgeB.x[0].x
# get new loc
new_loc = (locA[0] + locB[0] - intersection_loc[0], locA[1] + locB[1] - intersection_loc[1])
if new_loc[0] >= 255 or new_loc[1] >= 255 or \
new_loc[0] <= 0 or new_loc[1] <= 0:
continue
new_corner = Element(new_loc)
new_candidate = candidate.copy()
new_graph = new_candidate.graph
edgeA = new_graph.getRealElement(edgeA)
edgeB = new_graph.getRealElement(edgeB)
if new_candidate.addable(new_corner) is False:
continue
new_corner = new_candidate.addCorner_v2(new_corner)
# get cornerA and cornerB from edgeA, edgeB
cornerA = edgeA.x[1] if l2_distance(edgeA.x[0].x, intersection_loc) < l2_distance(edgeA.x[1].x,
intersection_loc) \
else edgeA.x[0]
cornerB = edgeB.x[1] if l2_distance(edgeB.x[0].x, intersection_loc) < l2_distance(edgeB.x[1].x,
intersection_loc) \
else edgeB.x[0]
# new edge can not be too short
if l2_distance(cornerA.x, new_corner.x) < 12:
continue
if l2_distance(cornerB.x, new_corner.x) < 12:
continue
flag = False
for edge_ele in new_graph.getEdges():
if new_corner in edge_ele.x and cornerA in edge_ele.x:
flag = True
break
if edge_ele.x[0] not in (new_corner, cornerA):
l = get_distance_of_corner_and_edge(new_corner.x, cornerA.x, edge_ele.x[0].x)
if l <= 7:
flag = True
break
if edge_ele.x[1] not in (new_corner, cornerA):
l = get_distance_of_corner_and_edge(new_corner.x, cornerA.x, edge_ele.x[1].x)
if l <= 7:
flag = True
break
if flag:
continue
add_edgeA = new_candidate.addEdge(new_corner, cornerA)
if new_graph.checkIntersectionEdge(add_edgeA):
continue
flag = False
for edge_ele in new_graph.getEdges():
if new_corner in edge_ele.x and cornerB in edge_ele.x:
flag = True
break
if edge_ele.x[0] not in (new_corner, cornerB):
l = get_distance_of_corner_and_edge(new_corner.x, cornerB.x, edge_ele.x[0].x)
if l <= 7:
flag = True
break
if edge_ele.x[1] not in (new_corner, cornerB):
l = get_distance_of_corner_and_edge(new_corner.x, cornerB.x, edge_ele.x[1].x)
if l <= 7:
flag = True
break
if flag:
continue
add_edgeB = new_candidate.addEdge(new_corner, cornerB)
if new_graph.checkIntersectionEdge(add_edgeB):
continue
new_candidates.append(new_candidate)
return new_candidates
def adding_a_orthogonal_edge(candidate):
new_candidates = []
graph = candidate.graph
edges = candidate.graph.getEdges()
for edge in edges:
cornerA = edge.x[0]
cornerB = edge.x[1]
# get orthogonal direction
dir_ = (cornerA.x[1] - cornerB.x[1], cornerB.x[0] - cornerA.x[0])
for the_corner in edge.x:
temp_orth_loc = (the_corner.x[0] - dir_[0], the_corner.x[1] - dir_[1])
for inter_edge in edges:
if inter_edge == edge:
continue
if the_corner in inter_edge.x:
continue
intersection_loc = get_two_edge_intersection_location(
the_corner.x, temp_orth_loc, inter_edge.x[0].x, inter_edge.x[1].x
)
if intersection_loc[0] >= 255 or intersection_loc[1] >= 255 or \
intersection_loc[0] <= 0 or intersection_loc[1] <= 0:
continue
if np.dot((inter_edge.x[0].x[0] - intersection_loc[0], inter_edge.x[0].x[1] - intersection_loc[1]),
(inter_edge.x[1].x[0] - intersection_loc[0], inter_edge.x[1].x[1] - intersection_loc[1])) > 0:
# which means the intersection is not inside inter_edge but at the edge extension
continue
if l2_distance(intersection_loc, inter_edge.x[0].x) < 5 or \
l2_distance(intersection_loc, inter_edge.x[1].x) < 5:
continue
# no thin degree with neighbor edge
flag = False
neighbor_corners = graph.getNeighborCorner(the_corner)
for corner_ele in neighbor_corners:
if corner_ele in edge.x:
continue
if degree_of_three_corners(corner_ele.x, intersection_loc, the_corner.x) < 15:
flag = True
break
if degree_of_three_corners(corner_ele.x, intersection_loc, the_corner.x) > 165:
flag = True
break
if flag:
continue
new_candidate = candidate.copy()
new_graph = new_candidate.graph
new_corner = Element(intersection_loc)
if new_candidate.addable(new_corner) is False:
continue
new_corner = new_candidate.addCorner_v2(new_corner)
# new edge can not be too short
if l2_distance(new_corner.x, the_corner.x) < 7:
continue
add_edge = new_candidate.addEdge(new_corner, new_graph.getRealElement(the_corner))
if new_graph.checkIntersectionEdge(add_edge):
continue
new_candidates.append(new_candidate)
return new_candidates
class _thread(threading.Thread):
def __init__(self, threadID, name, candidate, lock, result_list, func):
threading.Thread.__init__(self)
self.threadID = threadID
self.name = name
self.candidate = candidate
self.lock = lock
self.result_list = result_list
self.func = func
def run(self):
print('running id: ', self.name)
start_time = time.time()
candidates = self.func(self.candidate)
print('test: =================================', self.name, len(candidates))
self.lock.acquire()
self.result_list.extend(candidates)
self.lock.release()
print(self.name, "spend time: {}s".format(time.time() - start_time))
def candidate_enumerate_training(candidate, gt):
new_candidates = []
# remove a corner
try:
new_ = removing_a_corner_operation(candidate)
if len(new_) > 0:
new_candidates.append(random.choice(new_))
except:
print('something wrong with remove a corner !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
# remove a colinear corner
try:
new_ = removing_a_colinear_corner_operation(candidate)
if len(new_) > 0:
new_candidates.append(random.choice(new_))
except:
print('something wrong with remove a colinear corner !!!!!!!!!!!!!!!!!!!!!!!!!!!')
# remove an edge
try:
new_ = removing_an_edge_operation(candidate)
if len(new_) > 0:
new_candidates.append(random.choice(new_))
except:
print('something wrong with remove an edge !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
# add an edge from existed corner
try:
new_ = adding_an_edge_operation(candidate)
if len(new_) > 0:
new_candidates.append(random.choice(new_))
except:
print('something wrong with add an edge from existed corner !!!!!!!!!!!!!!!!!!!!')
# add a corner from two edges
try:
new_ = adding_a_corner_from_two_edges_extension(candidate)
if len(new_) > 0:
new_candidates.append(random.choice(new_))
except:
print('something wrong with add a corner from two edges !!!!!!!!!!!!!!!!!!!!!!!!')
try:
new_ = adding_a_corner_from_parallel(candidate)
if len(new_) > 0:
new_candidates.append(random.choice(new_))
except:
print('something wrong with add a corner from parallel !!!!!!!!!!!!!!!!!!!!!!!!')
# add an edge from gt
try:
new_ = adding_an_edge_from_gt(candidate, gt)
if len(new_) > 0:
new_candidates.append(random.choice(new_))
except:
print('something wrong with add an edge from gt !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
# add a orthogonal edge
try:
new_ = adding_a_orthogonal_edge(candidate)
if len(new_) > 0:
new_candidates.append(random.choice(new_))
except:
print('something wrong with add a orthogonal edge !!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
return new_candidates
def candidate_enumerate(candidate):
new_candidates = []
new_candidates.extend(removing_a_corner_operation(candidate))
new_candidates.extend(removing_a_colinear_corner_operation(candidate))
new_candidates.extend(removing_an_edge_operation(candidate))
new_candidates.extend(adding_an_edge_operation(candidate))
new_candidates.extend(adding_a_corner_from_two_edges_extension(candidate))
new_candidates.extend(adding_a_corner_from_parallel(candidate))
new_candidates.extend(adding_a_orthogonal_edge(candidate))
return new_candidates
def candidate_enumerate_thread(candidate):
new_candidates = []
lock = threading.Lock()
thread1 = _thread(1, 'remove_a_corner', candidate, lock, new_candidates, removing_a_corner_operation)
thread2 = _thread(2, 'remove_a_colinear_corner', candidate, lock, new_candidates,
removing_a_colinear_corner_operation)
thread3 = _thread(3, 'add_an_edge', candidate, lock, new_candidates, adding_an_edge_operation)
thread4 = _thread(4, 'remove_an_edge', candidate, lock, new_candidates, removing_an_edge_operation)
thread1.start()
thread2.start()
thread3.start()
thread4.start()
threads = []
threads.append(thread1)
threads.append(thread2)
threads.append(thread3)
threads.append(thread4)
for t in threads:
t.join()
return new_candidates
def reduce_duplicate_candidate(candidates):
i = 0
while i < len(candidates):
for j in reversed(range(i + 1, len(candidates))):
if candidates[i].equal(candidates[j]):
del candidates[j]
i = i + 1
return candidates
def save_candidate_image(candidate, base_path, base_name):
corners = candidate.graph.getCornersArray()
edges = candidate.graph.getEdgesArray()
# graph svg
svg = svg_generate(corners, edges, base_name, samecolor=True)
svg.saveas(os.path.join(base_path, base_name + '.svg'))
# corner image
temp_mask = np.zeros((256, 256))
for ele in candidate.graph.getCorners():
if ele.get_score() < 0:
temp_mask = cv2.circle(temp_mask, ele.x[::-1], 3, 1, -1)
fig = plt.figure(frameon=False)
fig.set_size_inches(1, 1)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.set_axis_off()
fig.add_axes(ax)
ax.imshow(temp_mask, aspect='auto')
fig.savefig(os.path.join(base_path, base_name + '_corner.png'), dpi=256)
# edges image
temp_mask = np.zeros((256, 256))
for ele in candidate.graph.getEdges():
if ele.get_score() < 0:
A = ele.x[0]
B = ele.x[1]
temp_mask = cv2.line(temp_mask, A.x[::-1], B.x[::-1], 1, thickness=1)
ax.imshow(temp_mask, aspect='auto')
fig.savefig(os.path.join(base_path, base_name + '_edge.png'), dpi=256)
# region no need fig
plt.close()
#########################################################################################
###################################### Class ############################################
#########################################################################################
class Element:
def __init__(self, x, safe_count=0):
assert type(x) is tuple
assert type(x[0]) == int or type(x[0]) == Element
assert type(x[1]) == int or type(x[1]) == Element
self.x = x
self.__score = None
self.safe_count = safe_count
def store_score(self, score):
self.__score = score
def get_score(self):
return self.__score
def equal(self, ele):
if type(self.x[0]) != type(ele.x[0]):
return False
if type(self.x[0]) == int:
# corner
return True if self.x[0] == ele.x[0] and self.x[1] == ele.x[1] else False
if type(self.x[0]) == Element:
# edge
if self.x[0].equal(ele.x[0]) and self.x[1].equal(ele.x[1]):
return True
if self.x[1].equal(ele.x[0]) and self.x[0].equal(ele.x[1]):
return True
return False
raise BaseException('no implement type')
class regionCache():
def __init__(self, datapath):
self.cache = {}
self.datapath = datapath
def get_region(self, name):
if name in self.cache.keys():
return self.cache[name]
gt_mask = np.load(os.path.join(self.datapath, name + '.npy'))
if len(self.cache) == 5:
self.cache.pop(list(self.cache.keys())[0])
self.cache[name] = gt_mask
return gt_mask
class imgCache():
def __init__(self, datapath):
self.cache = {}
self.datapath = datapath
def get_image(self, name):
if name in self.cache.keys():
return self.cache[name]
img = skimage.img_as_float(plt.imread(os.path.join(self.datapath, 'rgb', name + '.jpg')))
if len(self.cache) == 5:
self.cache.pop(list(self.cache.keys())[0])
self.cache[name] = img
return img
class Graph:
def __init__(self, corners, edges):
corners, edges = sort_graph(corners, edges)
self.__corners = []
for corner_i in range(corners.shape[0]):
self.__corners.append(
Element(
tuple(
(int(corners[corner_i, 0]), int(corners[corner_i, 1]))
)
)
)
self.__edges = []
for edge_i in range(edges.shape[0]):
self.__edges.append(Element((self.__corners[edges[edge_i, 0]], self.__corners[edges[edge_i, 1]])))
self.__regions = []
self.__regions.append(Element((0, 0))) # we use entire region here
@classmethod
def initialFromTuple(cls, corners, edges):
edge_index = []
for item in edges:
a = corners.index(item[0])
b = corners.index(item[1])
edge_index.append((a, b))
edge_index = np.array(edge_index)
corners = np.array(corners)
return cls(corners, edge_index)
def store_score(self, corner_score=None, edge_score=None, region_score=None):
'''
:param corner_score: np array size: len(corners)
:param edge_score: np array size: len(edges)
:param region_score: np.array size: len(regions)
:return:
'''
if corner_score is not None:
for idx, element in enumerate(self.__corners):
element.store_score(corner_score[idx])
if edge_score is not None:
for idx, element in enumerate(self.__edges):
element.store_score(edge_score[idx])
if region_score is not None:
for idx, element in enumerate(self.__regions):
element.store_score(region_score[idx])
return
def getCornersArray(self):
c = []
for ele in self.__corners:
c.append(ele.x)
return np.array(c)
def getEdgesArray(self):
c = []
for ele in self.__edges:
corner1 = ele.x[0]
corner2 = ele.x[1]
idx1 = self.__corners.index(corner1)
idx2 = self.__corners.index(corner2)
c.append([idx1, idx2])
return np.array(c)
def getCorners(self):
return self.__corners
def getRegions(self):
return self.__regions
def getEdges(self):
return self.__edges
def graph_score(self):
corner_score = 0
for ele in self.__corners:
corner_score += ele.get_score()
edge_score = 0
for ele in self.__edges:
edge_score += ele.get_score()
region_score = 0
for ele in self.__regions:
region_score += ele.get_score()
return score_weights[0] * corner_score + score_weights[1] * edge_score + score_weights[2] * region_score
def corner_score(self):
corner_score = 0
for ele in self.__corners:
corner_score += ele.get_score()
return corner_score
def edge_score(self):
edge_score = 0
for ele in self.__edges:
edge_score += ele.get_score()
return edge_score
def region_score(self):
region_score = 0
for ele in self.__regions:
region_score += ele.get_score()
return region_score
def remove(self, ele):
'''
:param ele: remove eles as well as some other related elements
:return: set() of removed elements
'''
# corner
removed = set()
if ele in self.__corners:
self.__corners.remove(ele)
removed.add(ele)
# remove edge that has the corner
for idx in reversed(range(len(self.__edges))):
edge_ele = self.__edges[idx]
if ele in edge_ele.x:
removed = removed.union(self.remove(edge_ele))
# edge
elif ele in self.__edges:
self.__edges.remove(ele)
removed.add(ele)
corner1 = ele.x[0]
corner2 = ele.x[1]
if corner1.safe_count == 0:
# can be delete
_count = 0
for edge_ele in self.__edges:
if corner1 in edge_ele.x:
_count += 1
if _count == 0:
removed = removed.union(self.remove(corner1))
if corner2.safe_count == 0:
# can be delete
_count = 0
for edge_ele in self.__edges:
if corner2 in edge_ele.x:
_count += 1
if _count == 0:
removed = removed.union(self.remove(corner2))
return removed
def has_edge(self, ele1, ele2):
"""
:param ele1: corner1
:param ele2: corner2
:return: edge or none
"""
for edge_ele in self.__edges:
if ele1 in edge_ele.x and ele2 in edge_ele.x:
return edge_ele
return None
def add_edge(self, ele1, ele2):
temp = self.has_edge(ele1, ele2)
if temp is not None:
temp.safe_count = SAFE_NUM
return temp
new_ele = Element((ele1, ele2), safe_count=SAFE_NUM)
self.__edges.append(new_ele)
return new_ele
def add_corner(self, ele):
for corner in self.__corners:
if corner.x == ele.x:
corner.safe_count = SAFE_NUM
return corner
ele.safe_count = SAFE_NUM
self.__corners.append(ele)
return ele
def add_corner_v2(self, ele):
# if new corner is near a existed corner, return the existed corner
# if new corner is on an edge, split edge
for corner in self.__corners:
if l2_distance(corner.x, ele.x) < 5:
corner.safe_count = SAFE_NUM
return corner
min_d = 256
the_edge = None
for edge in self.__edges:
temp = get_distance_of_corner_and_edge(edge.x[0].x, edge.x[1].x, ele.x)
if temp < min_d:
min_d = temp
the_edge = edge
if min_d < 3:
# split edge
corner1 = the_edge.x[0]
corner2 = the_edge.x[1]
new_ele = Element((corner1, ele), safe_count=the_edge.safe_count)
self.__edges.append(new_ele)
new_ele = Element((corner2, ele), safe_count=the_edge.safe_count)
self.__edges.append(new_ele)
self.__edges.remove(the_edge)
ele.safe_count = SAFE_NUM
self.__corners.append(ele)
return ele
def checkColinearCorner(self, ele):
if self.getCornerDegree(ele) != 2:
return False
edge_in = []
for edge_ele in self.__edges:
if ele in edge_ele.x:
edge_in.append(edge_ele)
if len(edge_in) == 2:
break
two_neighbor = {edge_in[0].x[0], edge_in[0].x[1], edge_in[1].x[0], edge_in[1].x[1]}
two_neighbor.remove(ele)
two_neighbor = tuple(two_neighbor)
if self.has_edge(two_neighbor[0], two_neighbor[1]) is not None:
return False
line1 = np.array(ele.x) - np.array(two_neighbor[0].x)
line2 = np.array(two_neighbor[1].x) - np.array(ele.x)
cos = np.dot(line1, line2) / (np.linalg.norm(line1) * np.linalg.norm(line2))
cos = min(1, max(-1, cos))
if np.arccos(cos) < np.pi / 9: # 20 degree
return True
return False
def checkIntersectionEdge(self, ele=None):
if ele is None:
for edge_i in range(len(self.__edges)):
for edge_j in range(edge_i + 1, len(self.__edges)):
if check_intersection(self.__edges[edge_i], self.__edges[edge_j]):
return True
return False
for edge_ele in self.__edges:
if ele == edge_ele:
continue
if check_intersection(edge_ele, ele):
return True
return False
def getCornerDegree(self, ele):
degree = 0
for edge_ele in self.__edges:
if ele in edge_ele.x:
degree += 1
return degree
def getEdgeConnected(self, ele):
out_ = set()
if type(ele.x[0]) == int:
# corner
for edge_ele in self.__edges:
if ele in edge_ele.x:
out_.add(edge_ele)
return out_
if type(ele.x[0]) == Element:
# Edge
out_ = out_.union(self.getEdgeConnected(ele.x[0]))
out_ = out_.union(self.getEdgeConnected(ele.x[1]))
if ele in out_:
out_.remove(ele)
return out_
def getNeighborCorner(self, ele):
out_ = set()
for edge_ele in self.__edges:
if ele == edge_ele.x[0]:
out_.add(edge_ele.x[1])
if ele == edge_ele.x[1]:
out_.add(edge_ele.x[0])
return out_
def getRealElement(self, ele):
# edge
if type(ele.x[0]) == Element:
for e in self.__edges:
if (e.x[0].x == ele.x[0].x and e.x[1].x == ele.x[1].x) or \
(e.x[1].x == ele.x[0].x and e.x[0].x == ele.x[1].x):
return e
raise BaseException("no same edge exists.")
# corner
elif type(ele.x[0]) == int:
for c in self.__corners:
if c.x == ele.x:
return c
raise BaseException("no same corner exists.")
def copy(self):
corners = self.getCornersArray()
edges = self.getEdgesArray()
new_graph = Graph(corners, edges)
for idx, ele in enumerate(self.__corners):
new_graph.__corners[idx].store_score(self.__corners[idx].get_score())
for idx, ele in enumerate(self.__edges):
new_graph.__edges[idx].store_score(self.__edges[idx].get_score())
for idx, ele in enumerate(self.__regions):
new_graph.__regions[idx].store_score(self.__regions[idx].get_score)
return new_graph
def update_safe_count(self):
for ele in self.__corners:
if ele.safe_count > 0:
ele.safe_count -= 1
for ele in self.__edges:
if ele.safe_count > 0:
ele.safe_count -= 1
def isNeighbor(self, element1, element2):
'''
:param element1:
:param element2:
:return: True / False
'''
if element1 == element2:
return False
if type(element1.x[0]) != type(element2.x[0]):
# corner and edge
return False
if type(element1.x[0]) == int:
# both are corner type
for edge_ele in self.__edges:
if edge_ele.x[0] == element1 and edge_ele.x[1] == element2:
return True
if edge_ele.x[0] == element2 and edge_ele.x[1] == element1:
return True
return False
if type(element1.x[0]) == Element:
# both are edge type
if len({element1.x[0], element1.x[1], element2.x[0], element2.x[1]}) < 4:
return True
return False
def equal(self, graph):
if len(self.__corners) != len(graph.__corners) or \
len(self.__edges) != len(graph.__edges):
return False
for corner_i in range(len(self.__corners)):
if self.__corners[corner_i].equal(graph.__corners[corner_i]) is False:
return False
for edge_i in range(len(self.__edges)):
if self.__edges[edge_i].equal(graph.__edges[edge_i]) is False:
return False
return True
class Candidate:
def __init__(self, graph, name, corner_existed_before, edge_existed_before):
'''
:param graph: Class graph
:param name: string, data name
:param corner_existed_before: dict {(x_i,y_i):c_1 ...} indicates counts for corresponding corners, after one search,
counts -= 1, if count == 0, remove from the set.
:param edge_existed_before: dict {((x_i1,y_i1),(x_i2,y_i2)):ci}
'''
self.graph = graph
self.name = name
self.corner_existed_before = corner_existed_before
self.edge_existed_before = edge_existed_before
@classmethod
def initial(cls, graph, name):
return cls(graph, name, {}, {})
def update(self):
# all the existed before elements count - 1
for key in self.corner_existed_before.keys():
self.corner_existed_before[key] -= 1
for key in self.edge_existed_before.keys():
self.edge_existed_before[key] -= 1
# check if some need to remove from existed before set
for key in list(self.corner_existed_before.keys()):
if self.corner_existed_before[key] == 0:
self.corner_existed_before.pop(key)
for key in list(self.edge_existed_before.keys()):
if self.edge_existed_before[key] == 0:
self.edge_existed_before.pop(key)
# update graph
self.graph.update_safe_count()
def copy(self):
corner_existed_before = self.corner_existed_before.copy()
edge_existed_before = self.edge_existed_before.copy()
new_graph = self.graph.copy()
return Candidate(new_graph, self.name, corner_existed_before, edge_existed_before)
def removable(self, ele):
'''
:param x: input is element
:return:
'''
assert type(ele) == Element
# edge
return True if ele.safe_count == 0 else False
def addable(self, ele):
if type(ele) == Element:
if type(ele.x[0]) == Element:
# edge
for edge in self.graph.getEdges():
c1 = edge.x[0]
c2 = edge.x[1]
if (ele.x[0].x == c1.x and ele.x[1].x == c2.x) or \
(ele.x[1].x == c1.x and ele.x[0].x == c2.x):
# already existed
return False
corner1_loc = ele.x[0].x
corner2_loc = ele.x[1].x
if (corner1_loc, corner2_loc) in self.edge_existed_before.keys() or \
(corner2_loc, corner1_loc) in self.edge_existed_before.keys():
return False
return True
else:
# corner
for corner in self.graph.getCorners():
if l2_distance(ele.x, corner.x) < TWO_CORNER_MINIMUM_DISTANCE:
# already existed
return False
if ele.x in self.corner_existed_before.keys():
return False
return True
else: # (x,y) or ((x1,y1),(x2,y2))
if type(ele[0]) == tuple:
# edge
corner1_loc = ele[0]
corner2_loc = ele[1]
for edge in self.graph.getEdges():
c1 = edge.x[0]
c2 = edge.x[1]
if (corner1_loc == c1.x and corner2_loc == c2.x) or \
(corner2_loc == c1.x and corner1_loc == c2.x):
# already existed
return False
if (corner1_loc, corner2_loc) in self.edge_existed_before.keys() or \
(corner2_loc, corner1_loc) in self.edge_existed_before.keys():
return False
return True
else:
# corner
for corner in self.graph.getCorners():
if l2_distance(ele, corner.x) < TWO_CORNER_MINIMUM_DISTANCE:
# already existed
return False
if ele in self.corner_existed_before.keys():
return False
return True
def addCorner(self, ele):
if ele.x in self.corner_existed_before.keys():
raise BaseException('cannot add the corner')
new_ele = self.graph.add_corner(ele) # possible changed
return new_ele
def addCorner_v2(self, ele):
if ele.x in self.corner_existed_before.keys():
raise BaseException('cannot add the corner')
new_ele = self.graph.add_corner_v2(ele)
return new_ele
def addEdge(self, ele1, ele2):
corner1 = ele1
corner2 = ele2
assert corner1 in self.graph.getCorners()
assert corner2 in self.graph.getCorners()
if (corner1.x, corner2.x) in self.edge_existed_before.keys() or \
(corner2.x, corner1.x) in self.edge_existed_before.keys():
raise BaseException('cannot add the edge')
new_ele = self.graph.add_edge(corner1, corner2)
return new_ele
def removeCorner(self, ele):
if ele.x in self.corner_existed_before.keys():
raise BaseException('already existed.')
self.corner_existed_before[ele.x] = SAFE_NUM
def removeEdge(self, ele):
corner1 = ele.x[0]
corner2 = ele.x[1]
loc1 = corner1.x
loc2 = corner2.x
if (loc1[0] > loc2[0]) or (loc1[0] == loc2[0] and loc1[1] > loc2[1]):
loc1 = corner2.x
loc2 = corner1.x
if (loc1, loc2) in self.edge_existed_before.keys():
raise BaseException('already existed.')
self.edge_existed_before[(loc1, loc2)] = SAFE_NUM
def generate_new_candidate_remove_a_colinear_corner(self, ele):
# need to check if ele is a colinear corner before
new_candidate = self.copy()
new_graph = new_candidate.graph
ele = new_graph.getRealElement(ele)
# find two neighbor corners
temp = set()
for element in new_graph.getEdgeConnected(ele):
# edge
if type(element.x[0]) == Element:
temp.add(element.x[0])
temp.add(element.x[1])
temp.remove(ele)
temp = tuple(temp)
assert len(temp) == 2
# add edge to two neighbor corners
# (add before remove, in case the neighbor corners will be removed by zero degree)
# special case no need to check existed_before, instead remove if in existed_before dict
added = new_graph.add_edge(temp[0], temp[1])
if (temp[0].x, temp[1].x) in self.edge_existed_before.keys():
self.edge_existed_before.pop((temp[0].x, temp[1].x))
if (temp[1].x, temp[0].x) in self.edge_existed_before.keys():
self.edge_existed_before.pop((temp[1].x, temp[0].x))
# remove
removed = new_graph.remove(ele)
# add removed elements into existed before
for element in removed:
# edge
if type(element.x[0]) == Element:
new_candidate.removeEdge(element)
# corner
elif type(element.x[0]) == int:
new_candidate.removeCorner(element)
else:
raise BaseException('wrong type.')
# modify scores that need to be recounted
# all corners are recounted
for element in new_graph.getCorners():
element.store_score(None)
# edges that are neighbors to the removed edges OR new edges will be recounted
for element in new_graph.getEdges():
for modified_ele in removed.union({added}):
if new_graph.isNeighbor(element, modified_ele):
element.store_score(None)
break
# all regions are recounted
for element in new_graph.getRegions():
element.store_score(None)
return new_candidate
def generate_new_candidate_remove_a_corner(self, ele):
# need to check if ele is removable before call this method
new_candidate = self.copy()
new_graph = new_candidate.graph
ele = new_graph.getRealElement(ele)
removed = new_graph.remove(ele)
# add removed elements into existed before
for element in removed:
# edge
if type(element.x[0]) == Element:
corner1 = element.x[0]
corner2 = element.x[1]
loc1 = corner1.x
loc2 = corner2.x
if (loc1[0] > loc2[0]) or (loc1[0] == loc2[0] and loc1[1] > loc2[1]):
loc1 = corner2.x
loc2 = corner1.x
if (loc1, loc2) in self.edge_existed_before.keys():
raise BaseException('already existed.')
new_candidate.edge_existed_before[(loc1, loc2)] = SAFE_NUM
# corner
elif type(element.x[0]) == int:
if element.x in self.corner_existed_before.keys():
raise BaseException('already existed.')
new_candidate.corner_existed_before[element.x] = SAFE_NUM
else:
raise BaseException('wrong type.')
# modify scores that need to be recounted
# all corners are recounted
for element in new_graph.getCorners():
element.store_score(None)
# edges that are neighbors to the removed edges will be recounted
for element in new_graph.getEdges():
for removed_ele in removed:
if new_graph.isNeighbor(element, removed_ele):
element.store_score(None)
break
# all regions are recounted
for element in new_graph.getRegions():
element.store_score(None)
return new_candidate
def generate_new_candidate_add_an_edge(self, ele1, ele2):
# need to check addable before call this method
new_candidate = self.copy()
new_graph = new_candidate.graph
ele1 = new_graph.getRealElement(ele1)
ele2 = new_graph.getRealElement(ele2)
# add edge
new_ele = new_candidate.addEdge(ele1, ele2)
# modify scores that need to be recounted
# all corners are recounted
for element in new_graph.getCorners():
element.store_score(None)
# edges that are neighbors to the added edges will be recounted
for element in new_graph.getEdges():
if new_graph.isNeighbor(element, new_ele):
element.store_score(None)
# all regions are recounted
for element in new_graph.getRegions():
element.store_score(None)
return new_candidate
def generate_new_candidate_remove_an_edge(self, ele):
# need to check if ele is removable before call this method
new_candidate = self.copy()
new_graph = new_candidate.graph
ele = new_graph.getRealElement(ele)
removed = new_graph.remove(ele)
# add removed elements into existed before
for element in removed:
# edge
if type(element.x[0]) == Element:
corner1 = element.x[0]
corner2 = element.x[1]
loc1 = corner1.x
loc2 = corner2.x
if (loc1[0] > loc2[0]) or (loc1[0] == loc2[0] and loc1[1] > loc2[1]):
loc1 = corner2.x
loc2 = corner1.x
if (loc1, loc2) in self.edge_existed_before.keys():
raise BaseException('already existed.')
new_candidate.edge_existed_before[(loc1, loc2)] = SAFE_NUM
# corner
elif type(element.x[0]) == int:
if element.x in self.corner_existed_before.keys():
raise BaseException('already existed.')
new_candidate.corner_existed_before[element.x] = SAFE_NUM
else:
raise BaseException('wrong type.')
# modify scores that need to be recounted
# all corners are recounted
for element in new_graph.getCorners():
element.store_score(None)
# edges that are neighbors to the removed edges will be recounted
for element in new_graph.getEdges():
for removed_ele in removed:
if new_graph.isNeighbor(element, removed_ele):
element.store_score(None)
break
# all regions are recounted
for element in new_graph.getRegions():
element.store_score(None)
return new_candidate
def generate_new_candidate_add_a_new_triangle(self, ele_new, ele1, ele2):
# this method is to add a new corner as well as two new edges into the graph
# need to check addable of ele_new before call this method
new_candidate = self.copy()
new_graph = new_candidate.graph
ele1 = new_graph.getRealElement(ele1)
ele2 = new_graph.getRealElement(ele2)
# add corner
ele_new = new_candidate.addCorner(ele_new) # ele_new possible change
# no score need to be recounted in current situation
# add two_new edge (ele1, ele_new) and (ele2, ele_new)
new_candidate = new_candidate.generate_new_candidate_add_an_edge(ele_new, ele1)
new_candidate = new_candidate.generate_new_candidate_add_an_edge(ele_new, ele2)
return new_candidate
def generate_new_candidate_add_a_corner(self, ele):
# need to check addable of ele before call this method
new_candidate = self.copy()
new_graph = new_candidate.graph
# add corner
ele = new_candidate.addCorner(ele)
# modify scores that need to be recounted
# all corners are recounted
for element in new_graph.getCorners():
element.store_score(None)
# no edge need to be recounted
# all regions are recounted
for element in new_graph.getRegions():
element.store_score(None)
return new_candidate
def equal(self, candidate):
return self.graph.equal(candidate.graph)