ymj_bim_test/search.py

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
    ## 为了不贴边，x，y固定加 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
    ## 为了不贴边，x，y固定加 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)