ymj_bim_test/search.py
2025-02-27 17:35:39 +08:00

521 lines
20 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import cv2
from utils import filter_points
print(cv2.__version__) # 4.9.0
import json
import math
import copy
import numpy as np
import time
from scipy.spatial import distance
from scipy.optimize import linear_sum_assignment
def get_params(num_param, start, end):
return copy.deepcopy(num_param[start:end+1])
def parmas_to_num(text_param):
for item in text_param:
# item['center'] = int(float(item['center']) * 1000)
item['center'] = int(float(item['center']))
item['x'] = int(float(item['x']))
item['w'] = int(item['w'])
item['h'] = int(item['h'])
return text_param
def parmas_to_text(num_param):
for item in num_param:
item['center'] = str(item['center'])
item['x'] = str(item['x'])
item['w'] = str(item['w'])
item['h'] = str(item['h'])
item['x1'] = str(item['x1'])
item['y1'] = str(item['y1'])
item['x2'] = str(item['x2'])
item['y2'] = str(item['y2'])
item['x3'] = str(item['x3'])
item['y3'] = str(item['y3'])
item['x4'] = str(item['x4'])
item['y4'] = str(item['y4'])
item['x_center'] = str(item['x_center'])
item['y_center'] = str(item['y_center'])
return num_param
def sort_params(params):
sorted_params = sorted(params, key=lambda item: (item['center'], item['x']))
return sorted_params
def print_params(sort_params):
for param in sort_params:
print(param["center"], param["x"], param["w"], param["h"])
def print_path(search_path):
for path in search_path:
print(path[0], path[1])
def search_path(sort_params):
searchPath = []
for i in range(len(sort_params) - 1):
(r, theta) = cartesian_to_polar(sort_params[i]["x"], sort_params[i]["center"],
sort_params[i + 1]["x"], sort_params[i + 1]["center"])
searchPath.append([r, theta])
return searchPath
def normalize_params_and_path(sort_params, search_path, index=0):
base = sort_params[index]["h"]
for param in sort_params:
param['center'] /= base
param['x'] /= base
param['w'] /= base
param['h'] /= base
param['w/h'] = param['w'] / param['h']
if search_path != None:
for path in search_path:
path[0] /= base
# path[1] /= base
return sort_params, search_path
def read_from_json(file_path):
with open(file_path, 'r') as f:
loaded_array = json.load(f)
return loaded_array
def cartesian_to_polar(x1, y1, x2, y2):
dx = x2 - x1
dy = y2 - y1
r = math.sqrt(dx**2 + dy**2)
theta = math.atan2(dy, dx)
return r, theta
def calculate_second_point(x1, y1, r, theta_radians):
# theta_radians = math.radians(theta_degrees)
x2 = x1 + r * math.cos(theta_radians)
y2 = y1 + r * math.sin(theta_radians)
return x2, y2
def cal_c1c2c3c4(param, heigt):
'''
按照上左、上右、下右、下左的顺时针顺序
'''
param['x1'] = int(param['x'] - param['w'] / 2)
param['y1'] = heigt - int(param['center'] + param['h'] / 2)
param['x2'] = int(param['x'] + param['w'] / 2)
param['y2'] = heigt - int(param['center'] + param['h'] / 2)
param['x3'] = int(param['x'] + param['w'] / 2)
param['y3'] = heigt - int(param['center'] - param['h'] / 2)
param['x4'] = int(param['x'] - param['w'] / 2)
param['y4'] = heigt - int(param['center'] - param['h'] / 2)
param['x_center'] = int((param['x1'] + param['x2']) / 2)
param['y_center'] = int((param['y1'] + param['y3']) / 2)
return param
def gen_im_from_params(params, type="lines"):
# 依据bim数据生成bim图
# type: line points
## 计算bim图长和宽
max_y = -999999
max_y_idx = -1
max_x = -999999
max_x_idx = -1
for i, param in enumerate(params):
if param["x"] + param["w"] / 2 > max_x:
max_x = param["x"] + param["w"] / 2
max_x_idx = i
if param["center"] + param["h"] / 2 > max_y:
max_y = param["center"] + param["h"] / 2
max_y_idx = i
## 为了不贴边xy固定加 500
FILL_VALUE = 0
bim_width = int(max_x)
bim_height = int(max_y)
bim_width += FILL_VALUE
bim_height += FILL_VALUE
bim_channels = 1
im = np.zeros((bim_height, bim_width, bim_channels), dtype=np.uint8)
for i, param in enumerate(params):
cal_c1c2c3c4(param, bim_height)
if type == "lines":
pts = np.asarray([[param['x1'], param['y1']],
[param['x2'], param['y2']],
[param['x3'], param['y3']],
[param['x4'], param['y4']]])
cv2.polylines(im, [pts], True, (255,0,0), 8)
elif type == "points":
cv2.circle(im, (param['x1'], param['y1']), 1, 255, 20)
cv2.circle(im, (param['x2'], param['y2']), 1, 255, 20)
cv2.circle(im, (param['x3'], param['y3']), 1, 255, 20)
cv2.circle(im, (param['x4'], param['y4']), 1, 255, 20)
cv2.circle(im, (param['x'], int((param['y3'] + param['y2']) / 2 )), 1, (255,0,0), 8)
return im
def gen_points_from_params(params):
# 依据bim数据生成bim图
# type: line points
## 计算bim图长和宽
max_y = -999999
max_y_idx = -1
max_x = -999999
max_x_idx = -1
for i, param in enumerate(params):
if param["x"] + param["w"] / 2 > max_x:
max_x = param["x"] + param["w"] / 2
max_x_idx = i
if param["center"] + param["h"] / 2 > max_y:
max_y = param["center"] + param["h"] / 2
max_y_idx = i
## 为了不贴边xy固定加 500
bim_height = int(max_y)
points = []
adjacency = [] # 邻接矩阵
for i, param in enumerate(params):
cal_c1c2c3c4(param, bim_height)
points.append([param['x1'], param['y1'], i])
points.append([param['x2'], param['y2'], i])
points.append([param['x3'], param['y3'], i])
points.append([param['x4'], param['y4'], i])
# points.append([param['x'], int((param['y3'] + param['y2']) / 2 )])
points.append([param['x_center'], param['y_center'], i])
return points
def topological_similarity(adj1, adj2):
"""
计算两个拓扑结构的相似度
"""
# 计算邻接矩阵的相似度
similarity = np.sum(adj1 == adj2) / (adj1.shape[0] * adj2.shape[1])
return similarity
def find_topological_matches(points1, adj1, points2, adj2, threshold=0.8):
"""
基于拓扑结构寻找匹配点集
"""
matches = []
for i in range(len(points1)):
for j in range(len(points2)):
# 计算拓扑相似度
sim = topological_similarity(adj1[i], adj2[j])
if sim > threshold:
matches.append((i, j, sim))
# 按相似度排序
matches.sort(key=lambda x: x[2], reverse=True)
return matches
from sklearn.linear_model import RANSACRegressor
def ransac_shape_matching(points, reference_points):
model_ransac = RANSACRegressor(random_state=42)
try:
model_ransac.fit(reference_points, points[:len(reference_points)]) # 假设点数相同
except ValueError as e:
print("Error fitting the model:", e)
return []
inlier_mask = model_ransac.inlier_mask_
best_match_subset = points[inlier_mask]
return best_match_subset
def polar_to_cartesian(polar_points):
r = polar_points[:, 0]
theta = polar_points[:, 1]
x = r * np.cos(theta)
y = r * np.sin(theta)
return np.vstack((x, y)).T
def compute_shape_context(points, nbins_r=5, nbins_theta=12):
n = points.shape[0]
shape_contexts = []
r_max = np.max(distance.pdist(points))
r_edges = np.logspace(-1, np.log10(r_max), nbins_r)
theta_edges = np.linspace(-np.pi, np.pi, nbins_theta + 1)
for i in range(n):
current_point = points[i]
relative_points = points - current_point
rs = np.hypot(relative_points[:, 0], relative_points[:, 1])
thetas = np.arctan2(relative_points[:, 1], relative_points[:, 0])
H, _, _ = np.histogram2d(thetas, rs, bins=[theta_edges, r_edges])
H /= np.sum(H) + 1e-8 # 归一化
shape_contexts.append(H.flatten())
return np.array(shape_contexts)
def match_shapes(sc1, sc2):
cost_matrix = distance.cdist(sc1, sc2, metric='sqeuclidean')
row_ind, col_ind = linear_sum_assignment(cost_matrix)
return cost_matrix[row_ind, col_ind].sum()
def _sobel(image):
'''
_sobel
'''
if image.ndim > 2:
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# todo 增加几个参数 http://blog.csdn.net/sunny2038/article/details/9170013
_sobelx = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=1)
_sobely = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=1)
_sobelx = np.uint8(np.absolute(_sobelx))
_sobely = np.uint8(np.absolute(_sobely))
_sobelcombine = cv2.bitwise_or(_sobelx,_sobely)
return _sobelcombine
def _findContours(image):
'''
_findContours
http://blog.csdn.net/mokeding/article/details/20153325
'''
contours, _ = cv2.findContours(image.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
return sorted(contours, key=cv2.contourArea, reverse=True)
if __name__ == "__main__":
# 读取并处理数据
data_bim = {}
data_bim["type"] = 0
data_bim["params"] = read_from_json("data_bim.json")
data_bim["point"] = []
data_bim["params"] = parmas_to_num(data_bim["params"])
data_sub = {}
data_sub["type"] = 0
data_sub["params"] = []
# 创建测试子集
# sub_im = cv2.imread("wide_image.png")
# sub_zero = np.zeros_like(sub_im)
# _im_gray = cv2.cvtColor(sub_im, cv2.COLOR_BGR2GRAY)
# _im_gray = cv2.GaussianBlur(_im_gray, (5, 5), 0)
# _im_edge_sobel = _sobel(_im_gray)
# _, _im_thresh = cv2.threshold(_im_edge_sobel, 5, 255, cv2.THRESH_BINARY)
# cnts = _findContours(_im_thresh)
# for contour in cnts:
# x, y, w, h = cv2.boundingRect(contour)
# cv2.rectangle(sub_zero, (x, y), (x + w, y + h), (255, 255, 255), -1)
# _im_edge_sobel = _sobel(sub_zero)
# _, _im_thresh = cv2.threshold(_im_edge_sobel, 5, 255, cv2.THRESH_BINARY)
# cnts = _findContours(_im_thresh)
sub_im = cv2.imread("data_sub/test_1/wide_image.png")
sub_im2 = cv2.imread("data_sub/test_1/wide_image.png")
sub_zero = np.zeros_like(sub_im)
cnts = read_from_json("data_sub/test_1/data_sub.json")
# cnts的矩形画在sub_im 上
for i in range(len(cnts)):
p = cnts[i]
cv2.rectangle(sub_im2, (p['x'], p['y']), (p['x'] + p['width'], p['y'] + p['height']), (0, 0, 255), 2)
# 写编号
cv2.putText(sub_im2, str(i), (p['x'], p['y']), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
cv2.imshow("sub_im_before_filter", sub_im2)
# 过滤点
cnts = filter_points("data_sub/test_1/wide_image.png",cnts)
for contour in cnts:
# x, y, w, h = cv2.boundingRect(contour)
x = contour["x"]
y = contour["y"]
w = contour["width"]
h = contour["height"]
# 由于定位框大小大于预埋件大小,因此这里需要做缩放处理
k = int(h*0.2) # roi2
k = int(h*0.01) # roi1
x += k
y += k
w -= int(k) #
h -= int(k) # 2k
param = {}
param["x1"] = x
param["y1"] = y
param["x2"] = x + w
param["y2"] = y
param["x3"] = x + w
param["y3"] = y + h
param["x4"] = x
param["y4"] = y + h
param['x_center'] = int((param['x1'] + param['x2']) / 2)
param['y_center'] = int((param['y1'] + param['y3']) / 2)
param["w"] = param["x2"] - param["x1"]
param["h"] = param["y3"] - param["y1"]
param["x"] = int((param["x1"] + param["x2"]) / 2)
param['center'] = sub_im.shape[0] - int((param["y1"] + param["y3"]) / 2)
data_sub["params"].append(param)
cv2.rectangle(sub_zero, (x, y), (x + w, y + h), (0, 255, 0), 1)
bim_im = gen_im_from_params(data_bim["params"])
# cv2.namedWindow("bim", cv2.WINDOW_NORMAL)
# cv2.imshow("bim", bim_im)
cv2.imwrite("bim_im.png", bim_im)
# cv2.waitKey(0)
cv2.imshow("sub_zero", sub_zero)
cv2.waitKey(0)
# data_sub = {}
# data_sub["type"] = 0
# data_sub["params"] = get_params(data_bim["params"], 1,8)
# data_sub["point"] = []
# data_sub["params"][0]["center"] += 20
# data_sub["params"][0]["x"] += 13
# data_sub["params"][0]["w"] += 40
##################### 开始计算 ####################
# 1、计算sub的搜索路径
## 1.1 排序 y升序x升序
data_sub["params"] = sort_params(data_sub["params"])
_sub_sort_params = data_sub["params"]
start_time = time.time()
# sub_roi_height = int(max_y) #????
# sub_roi_width = int(max_x) #????
sub_roi_height = sub_im.shape[0]
sub_roi_width = sub_im.shape[1]
sub_roi_params_select_id = -1
sub_roi_w_base_len = 0 # 选中预埋件的宽度作为基础长度
sub_roi_divide_w_h = 1
polar_origin_x = 0 # 极坐标原点 x
polar_origin_y = 0 # 极坐标原点 y
start_x = 0
start_y = 0
## 1.2 选择一块完整的预埋件
for i, param in enumerate(_sub_sort_params):
if 0 < param['x1'] and param['x2'] < sub_roi_width \
and 0 < param['y1'] and param['y3'] < sub_roi_height:
sub_roi_params_select_id = i
sub_roi_w_base_len = param['x2'] - param['x1']
sub_roi_divide_w_h = param['w'] / param['h']
polar_origin_x = int(param['x1']) # 当前选择的预埋件的左上角 x 坐标
polar_origin_y = int(param['y1']) # 当前选择的预埋件的左上角 y 坐标
break
if sub_roi_params_select_id == -1 or sub_roi_w_base_len == 0 :
print("[ERROR]\t 拍摄的图像中没有完整的预埋件信息\n")
assert(0)
## 1.2.2 将其他预埋件相对于它的极坐标进行填写
for i, param in enumerate(_sub_sort_params):
if i != sub_roi_params_select_id:
param['r'], param['theta'] = cartesian_to_polar(_sub_sort_params[sub_roi_params_select_id]['x_center'],
_sub_sort_params[sub_roi_params_select_id]['y_center'],
param['x_center'],param['y_center'])
## 1.3计算所有点到该预埋件左上点的,点个数,平均极半径 和 平均极角度
sum_r, sum_theta = 0.0,0
count = 0
pts = gen_points_from_params(_sub_sort_params)
polar_list = []
for i, pt in enumerate(pts):
r, theta = cartesian_to_polar(polar_origin_x, polar_origin_y, pt[0], pt[1])
sum_r += r
sum_theta += theta
count += 1
polar_list.append([r, theta])
sum_r /= count * sub_roi_w_base_len
sum_theta /= count
print(count, sum_r, sum_theta)
# 初始化候选预埋件
candi_params = []
for i, param in enumerate(data_bim["params"]):
temp_div_w_h = param['w'] / param['h']
if temp_div_w_h / sub_roi_divide_w_h < 1.5 and temp_div_w_h / sub_roi_divide_w_h > 0.66:
candi_params.append(param)
print(f"形状筛选后还剩下:{len(candi_params)} 个候选, w / h = {sub_roi_divide_w_h}")
rst_params = []
bim_all_pts = gen_points_from_params(data_bim["params"])
bim_im = gen_im_from_params(data_bim["params"])
# sub_im = gen_im_from_params(data_sub["params"])# 需要读取
min_match_score = 999999
for i, param in enumerate(candi_params):
tmp_roi_w_base_len = param['x2'] - param['x1']
scale = tmp_roi_w_base_len / sub_roi_w_base_len
tmp_roi_width = int(scale * sub_roi_width)
tmp_roi_height = int(scale * sub_roi_height)
# 相对于bim图的坐标
tmp_roi_start_x = int(param['x1'] - scale * polar_origin_x)
tmp_roi_end_x = tmp_roi_start_x + tmp_roi_width
tmp_roi_start_y = int(param['y1'] - scale * polar_origin_y)
tmp_roi_end_y = tmp_roi_start_y + tmp_roi_height
tmp_roi_conners = [[tmp_roi_start_x, tmp_roi_start_y],
[tmp_roi_end_x, tmp_roi_start_y],
[tmp_roi_end_x, tmp_roi_end_y],
[tmp_roi_start_x, tmp_roi_end_y]]
tmp_sum_r, tmp_sum_theta = 0.0, 0
tmp_count = 0
tmp_polar_list = []
# 这里需要把选中的预埋件也提取出来
param['effective_points'] = []
for j, pt in enumerate(bim_all_pts):
if cv2.pointPolygonTest(np.asarray(tmp_roi_conners), (pt[0],pt[1]), False) > 0:
r, theta = cartesian_to_polar(param['x1'], param['y1'], pt[0], pt[1])
tmp_sum_r += r
tmp_sum_theta += theta
tmp_count += 1
tmp_polar_list.append([r, theta])
# 搜集有效点,以供后续继续处理
param['effective_points'].append(pt)
tmp_sum_r /= tmp_count * tmp_roi_w_base_len
tmp_sum_theta /= tmp_count
# 预埋件数量相差 30%,则不进行计算
# if tmp_count / count > 1.3 or tmp_count / count < 0.77: continue
if abs(tmp_count - count) == 0: score = 0.5
elif abs(tmp_count - count) <= 10: score = 0.3
elif abs(tmp_count - count) <= 20: score = 0.2
elif abs(tmp_count - count) <= 30: score = 0.1
else: score = 0.0
# else: score = (1 / abs(tmp_count - count) ) * 0.7
score += (1 - abs(tmp_sum_r - sum_r) / sub_roi_width) * 0.25
score += (1 - abs(tmp_sum_theta - sum_theta) / 3.14) * 0.25
if score > 0.6: #????
cartesian_points1 = polar_to_cartesian(np.asarray(tmp_polar_list))
cartesian_points2 = polar_to_cartesian(np.asarray(polar_list))
sc1 = compute_shape_context(cartesian_points1)
sc2 = compute_shape_context(cartesian_points2)
match_score = match_shapes(sc1, sc2)
print(score, match_score, tmp_count, tmp_sum_r, tmp_sum_theta)
if match_score < 5.0: #????
param["start_point"] = (tmp_roi_start_x, tmp_roi_start_y)
param["end_point"] = (tmp_roi_end_x, tmp_roi_end_y)
param["score"] = (score, tmp_count, tmp_sum_r*tmp_roi_w_base_len, tmp_sum_theta)
param['match_score'] = match_score
if min_match_score > match_score:
min_match_score = match_score
rst_params.append(param)
# 预埋件匹配
for i, param in enumerate(rst_params):
score = param["score"]
match_score = param["match_score"]
id = param["code"]
if min_match_score == match_score or match_score < 0.5: #????
# if True:
index_list = []
for j in range(len(param['effective_points'])):
pt = param['effective_points'][j]
if pt[2] not in index_list:
index_list.append(param['effective_points'][j][2])
# for i, param in enumerate(_sub_sort_params):
# param[]
print(f"起始预埋件ID{id},置信度为:{score[0]},点数量:{score[1]},平均长度为:{score[2]},平均角度为:{score[3]}")
print(f"match_sscore为{match_score}")
bim_im = cv2.rectangle(bim_im, param["start_point"], \
param["end_point"], 255, 6)
elapsed_time = time.time() - start_time
print(f"Execution time: {elapsed_time:.4f} seconds")
img_matches = cv2.resize(bim_im, (int(bim_im.shape[1]/6), int(bim_im.shape[0]/6)))
# sub_im = cv2.resize(sub_im, (int(sub_im.shape[1]/10), int(sub_im.shape[0]/10)))
cv2.imshow("2", img_matches)
# cnts的矩形画在sub_im 上
for i in range(len(cnts)):
p = cnts[i]
cv2.rectangle(sub_im, (p['x'], p['y']), (p['x']+p['width'],p['y']+p['height']), (0, 0, 255), 2)
# 写编号
cv2.putText(sub_im, str(i), (p['x'], p['y']), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
cv2.imshow("sub_im_after_filter", sub_im)
cv2.waitKey(0)