ymj_bim_test/search.py
2025-02-28 19:19:41 +08:00

592 lines
24 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_rectangle
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 # y坐标最大值
max_y_idx = -1 # y坐标最大值的索引
max_x = -999999 # x坐标最大值
max_x_idx = -1 # x坐标最大值的索引
# 遍历,找到矩形x坐标最大值和矩形y坐标最大值
for i, param in enumerate(params):
# x中心点坐标加上宽度的一半 = 当前矩形的x最大坐标
if param["x"] + param["w"] / 2 > max_x:
# 如果是最大的x坐标就更新
max_x = param["x"] + param["w"] / 2
max_x_idx = i
# y中心点坐标加上高度的一半 = 当前矩形的y最大坐标
if param["center"] + param["h"] / 2 > max_y:
# 如果是最大的y坐标就更新
max_y = param["center"] + param["h"] / 2
max_y_idx = i
padding_value = 0 # 内边距,避免整个bim图片贴着边缘展示
bim_width = int(max_x) + padding_value
print(f"[bim_width] ====== [{bim_width}]")
bim_height = int(max_y) + padding_value
print(f"[bim_height] ====== [{bim_height}]")
bim_channels = 3
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,255,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
"""
判断点是否在区域内部
True 在
False 不在
如果没有提供区域,不做判断,返回True
"""
def is_inside_roi(point, roi_w, roi_h):
if (roi_w != None and roi_h != None):
x = point[0]
y = point[1]
if (x <= 0 or y <= 0 or y >= roi_h or x >= roi_w ):
# 不在区域内部
return False
else:
# 在区域内部
return True
else:
# 没有提供区域,不做判断,默认返回True
return True
def gen_points_from_params(params,roi_w=None,roi_h=None):
# 依据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
bim_height = int(max_y)
points = []
for i, param in enumerate(params):
# 排序点 按照上左、上右、下右、下左的顺时针顺序
if (roi_w == None and roi_h == None):
cal_c1c2c3c4(param, bim_height)
# 过滤点,把在roi区域之外的点全部过滤掉.
if(is_inside_roi([param['x1'], param['y1']], roi_w, roi_h)):
points.append([param['x1'], param['y1'], i])
if(is_inside_roi([param['x2'], param['y2']], roi_w, roi_h)):
points.append([param['x2'], param['y2'], i])
if(is_inside_roi([param['x3'], param['y3']], roi_w, roi_h)):
points.append([param['x3'], param['y3'], i])
if(is_inside_roi([param['x4'], param['y4']], roi_w, roi_h)):
points.append([param['x4'], param['y4'], i])
if(is_inside_roi([param['x_center'], param['y_center']], roi_w, roi_h)):
points.append([param['x_center'], param['y_center'], i])
if (roi_w != None and roi_h != None):
print(f"[经区域过滤之后的点数一共为] ====== [{len(points)}]")
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)
original_rectangle = read_from_json("data_sub/test_1/data_sub.json")
# 过滤矩形
# cnts 过滤之后的矩形
# sub_im 裁剪之后的图像
cnts,sub_im = filter_rectangle("data_sub/test_1/wide_image.png", original_rectangle)
sub_zero = np.zeros_like(sub_im)
for contour in cnts:
# x, y, w, h = cv2.boundingRect(contour)
x = contour["x"]
y = contour["y"]
w = contour["width"]
h = contour["height"]
# 由于定位框大小大于预埋件大小,因此这里需要做缩放处理
kh = int(h * 0.01) # roi1
kw = int(w * 0.01) # roi1
x += int(kw)
y += int(kh)
w -= int(kw)
h -= int(kh)
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
# for i, p in enumerate(_sub_sort_params):
# # 画点
# cv2.circle(sub_im, (p["x_center"], p["y_center"]), 2, (0, 255, 255), -1)
# cv2.imshow("_sub_sort_params", sub_im)
# cv2.waitKey(0)
pts = gen_points_from_params(_sub_sort_params,sub_roi_width,sub_roi_height)
# # 测试,画出所有的pts
for i, p in enumerate(pts):
# 画点
cv2.circle(sub_im, (p[0], p[1]), 2, (0, 255, 255), -1)
# 写编号
cv2.putText(sub_im, str(i), (p[0], p[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
cv2.imshow("sub_im_points", sub_im)
# cv2.waitKey(0)
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(f"[所有点到该预埋件左上点的个数] ====== [{count}]")
print(f"[所有点到该预埋件左上点的平均极半径] ====== [{sum_r}]")
print(f"[所有点到该预埋件左上点的平均极角] ====== [{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.4
elif abs(tmp_count - count) <= 20: score = 0.3
elif abs(tmp_count - count) <= 30: score = 0.2
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.35
print("score=======", str(score))
if score > 0.6: #????
cartesian_points1 = polar_to_cartesian(np.asarray(tmp_polar_list)) # bim上的坐标
cartesian_points2 = polar_to_cartesian(np.asarray(polar_list)) # sub上的坐标
sc1 = compute_shape_context(cartesian_points1)
sc2 = compute_shape_context(cartesian_points2)
match_score = match_shapes(sc1, sc2)
print("score>0.6")
print(f"[score] ====== [{score}]")
print(f"[match_score] ====== [{match_score}]")
print(f"[tmp_count] ====== [{tmp_count}]")
print(f"[tmp_sum_r] ====== [{tmp_sum_r}]")
print(f"[tmp_sum_theta] ====== [{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
print(f"[start_point] ====== [{param["start_point"]}]")
print(f"[end_point] ====== [{param["end_point"]}]")
rst_params.append(param)
max_score = -99999
max_index = -1
for i, param in enumerate(rst_params):
score = param["score"]
match_score = param["match_score"]
if abs(score[0] / match_score) > max_score:
max_score = abs(score[0] / match_score)
max_index = i
bim_im = cv2.rectangle(bim_im,rst_params[max_index]["start_point"], rst_params[max_index]["end_point"], 100 * (i + 1), 50)
# # 预埋件匹配
# for i, param in enumerate(rst_params):
# score = param["score"]
# match_score = param["match_score"]
# print(f"[match_score] ====== [{match_score}]")
#
# id = param["code"]
#
# if min_match_score == match_score or match_score < 5.0: #????
# # 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}")
# print(f"[start_point] ====== [{param["start_point"]}]")
# print(f"[end_point] ====== [{param["end_point"]}]")
#
# bim_im = cv2.rectangle(bim_im, param["start_point"], param["end_point"], 100 * (i + 1), 50)
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)