双目相机测距原理

sgbm-video.py

import cv2
import numpy as np
import random
import math


# 左镜头的内参，如焦距
left_camera_matrix = np.array([[516.5066236,-1.444673028,320.2950423],[0,516.5816117,270.7881873],[0.,0.,1.]])
right_camera_matrix = np.array([[511.8428182,1.295112628,317.310253],[0,513.0748795,269.5885026],[0.,0.,1.]])


# 畸变系数,K1、K2、K3为径向畸变,P1、P2为切向畸变
# left_distortion = np.array([[-0.154511565,0.325173292, 0.006934081,0.017466934, -0.340007548]])
left_distortion = np.array([[-0.046645194,0.077595167, 0.012476819,-0.000711358,0]])
# right_distortion = np.array([[-0.192887524,0.706728768, 0.004233541,0.021340116,-1.175486913]])
right_distortion = np.array([[-0.061588946,0.122384376,0.011081232,-0.000750439,0]])

# 旋转矩阵
R = np.array([[0.999911333,-0.004351508,0.012585312],
              [0.004184066,0.999902792,0.013300386],
              [-0.012641965,-0.013246549,0.999832341]])
# 平移矩阵
T = np.array([-120.3559901,-0.188953775,-0.662073075])
size = (640, 480)

R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(left_camera_matrix, left_distortion,
                                                                  right_camera_matrix, right_distortion, size, R,
                                                                  T)
# 校正查找映射表,将原始图像和校正后的图像上的点一一对应起来
left_map1, left_map2 = cv2.initUndistortRectifyMap(left_camera_matrix, left_distortion, R1, P1, size, cv2.CV_16SC2)
right_map1, right_map2 = cv2.initUndistortRectifyMap(right_camera_matrix, right_distortion, R2, P2, size, cv2.CV_16SC2)

print(Q)
# -*- coding: utf-8 -*-
# 鼠标回调函数
def onmouse_pick_points(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        threeD = param
        print('\n像素坐标 x = %d, y = %d' % (x, y))
        # print("世界坐标是：", threeD[y][x][0], threeD[y][x][1], threeD[y][x][2], "mm")
        print("世界坐标xyz 是：", threeD[y][x][0] / 1000.0, threeD[y][x][1] / 1000.0, threeD[y][x][2] / 1000.0, "m")

        distance = math.sqrt(threeD[y][x][0] ** 2 + threeD[y][x][1] ** 2 + threeD[y][x][2] ** 2)
        distance = distance / 1000.0  # mm -> m
        print("距离是：", distance, "m")



# 2 加载视频文件
# capture = cv2.VideoCapture("output.avi")
capture = cv2.VideoCapture(1)

WIN_NAME = 'Deep disp'
cv2.namedWindow(WIN_NAME, cv2.WINDOW_AUTOSIZE)

# 3 读取视频
ret, frame = capture.read()
while ret:
    # 4 ret 是否读取到了帧，读取到了则为True
    cv2.imshow("video", frame)
    ret, frame = capture.read()
    img_color = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame1 = frame[0:480, 0:640]
    frame2 = frame[0:480, 640:1280]  # 割开双目图像

    imgL = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)  # 将BGR格式转换成灰度图片
    imgR = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)

    # cv2.remap 重映射，就是把一幅图像中某位置的像素放置到另一个图片指定位置的过程。
    # 依据MATLAB测量数据重建无畸变图片
    img1_rectified = cv2.remap(imgL, left_map1, left_map2, cv2.INTER_LINEAR)
    img2_rectified = cv2.remap(imgR, right_map1, right_map2, cv2.INTER_LINEAR)

    imageL = cv2.cvtColor(img1_rectified, cv2.COLOR_GRAY2BGR)
    imageR = cv2.cvtColor(img2_rectified, cv2.COLOR_GRAY2BGR)

    cv2.imshow("img-L", imageL)
    cv2.imshow("img-R", imageR)

    # SGBM-室外
    blockSize = 8
    img_channels = 3
    stereo = cv2.StereoSGBM_create(minDisparity=1,
                                   numDisparities=64,
                                   blockSize=blockSize,
                                   P1=8 * img_channels * blockSize * blockSize,
                                   P2=32 * img_channels * blockSize * blockSize,
                                   disp12MaxDiff=-1,
                                   preFilterCap=1,
                                   uniquenessRatio=10,
                                   speckleWindowSize=100,
                                   speckleRange=100,
                                   mode=cv2.STEREO_SGBM_MODE_HH)

    disparity = stereo.compute(img1_rectified, img2_rectified)  # 计算视差

    disp = cv2.normalize(disparity, disparity, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)  # 归一化函数算法

    threeD = cv2.reprojectImageTo3D(disparity, Q, handleMissingValues=True)  # 计算三维坐标数据值
    threeD = threeD * 16

    # threeD[y][x] x:0~640; y:0~480;   !!!!!!!!!!
    cv2.setMouseCallback(WIN_NAME, onmouse_pick_points, threeD)  # 鼠标回调事件

    #
    # # 三维重建代码
    # import open3d as o3d
    #
    dis_color = disparity
    dis_color = cv2.normalize(dis_color, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    dis_color = cv2.applyColorMap(dis_color, 2)
    cv2.imshow("depth", dis_color)
    cv2.imshow("left", frame1)
    cv2.imshow(WIN_NAME, disp)  # 显示深度图的双目画面
    # 5 若键盘按下q则退出播放
    if cv2.waitKey(20) & 0xff == ord('q'):
        break

# 4 释放资源
capture.release()

# 5 关闭所有窗口
cv2.destroyAllWindows()

sgbm_img.py

'''
Author: YeHanyu
Date: 2022-08-02 11:32:58
LastEditors: YeHanyu
LastEditTime: 2022-08-09 18:09:38
FilePath: /SGBM/sgbm_img.py
Description: 

Copyright (c) 2022 by YeHanyu, All Rights Reserved. 
'''

import cv2
import numpy as np
import random
import math
# import pcl
# import pcl.pcl_visualization
import time
from scipy.optimize import leastsq
# from scipy import io

# import matplotlib.pyplot as plt
# from mpl_toolkits.mplot3d import Axes3D

camera = '720p'
gt_pred = False
v_size = 5
Dst_mode = 'Z'
leastsq_mode = True
point_x = 0
point_y = 0


if camera=='kitti':
# 左镜头的内参，如焦距
    left_camera_matrix = np.array([[718.856,0,607.1928], [0, 718.856, 185.2157],[0.,0.,1.]])
    right_camera_matrix = np.array([[718.856,0,607.1928], [0, 718.856, 185.2157],[0.,0.,1.]])
    # 畸变系数,K1、K2、K3为径向畸变,P1、P2为切向畸变
    # left_distortion = np.array([[-0.154511565,0.325173292, 0.006934081,0.017466934, -0.340007548]])
    left_distortion = np.array([[0.0,0.0,0.0,0.0,0.0]])
    # right_distortion = np.array([[-0.192887524,0.706728768, 0.004233541,0.021340116,-1.175486913]])
    right_distortion = np.array([[0.0,0.0,0.0,0.0,0.0]])
    # 旋转矩阵
    R = np.array([[1,0.0,0.0],
                [0.0,1,0.0],
                [0.0,0.0,1]])
    # 平移矩阵
    T = np.array([537.165719, 0, 0])
    size = (1241, 376)
    
elif camera=='720p':
# 左镜头的内参，如焦距
    left_camera_matrix = np.array([[730.207611356134, 0, 617.761835197854], 
                                   [0, 729.722484488817, 356.281832636247],
                                   [0., 0., 1.]])
    right_camera_matrix = np.array([[726.617693292221, 0, 615.126833857842],
                                    [0, 726.598401621709, 355.079105280920],
                                    [0.,0.,1.]])
    # 畸变系数,K1、K2、K3为径向畸变,P1、P2为切向畸变
    # left_distortion=[K1，K2，P1，P2，K3]
    # left_distortion = np.array([[-0.154511565,0.325173292, 0.006934081,0.017466934, -0.340007548]])
    left_distortion = np.array([[0.114254497790507,-0.153178688236797,0,0,0]])
    # right_distortion = np.array([[-0.192887524,0.706728768, 0.004233541,0.021340116,-1.175486913]])
    right_distortion = np.array([[0.101619261365358,-0.128310783815266, 0, 0,0]])
    # 旋转矩阵
    R = np.array([[1,0.000117007865065619,0.000551464061148515],
                [-0.000116357622553170,1,-0.00117900500140576],
                [-0.000551601626991256,0.00117894064701274,1]])
    # 平移矩阵
    T = np.array([-59.9695960974267,-0.00773947743447085,0.298781994352248])
    size = (1280, 720)

R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(left_camera_matrix, left_distortion,
                                                                  right_camera_matrix, right_distortion, size, R,
                                                                  T)
# 校正查找映射表,将原始图像和校正后的图像上的点一一对应起来
left_map1, left_map2 = cv2.initUndistortRectifyMap(left_camera_matrix, left_distortion, R1, P1, size, cv2.CV_16SC2)
right_map1, right_map2 = cv2.initUndistortRectifyMap(right_camera_matrix, right_distortion, R2, P2, size, cv2.CV_16SC2)

def pcl_show(three, rgbimg):
    color = rgbimg.reshape(720*1280,3)
    three[np.isinf(three)]=0
    three[three >  10000] = 10000
    three[three <  -10000] = -10000
    temp = np.zeros((720, 1280, 4))
    temp[:,:,0:3] = three
    temp=temp.reshape(720*1280,4)
    temp=temp/1000
    temp = temp.astype(np.float32)
    for i in range(color.shape[0]):
        rgb = np.left_shift(color[i,2], 16) + np.left_shift(color[i,1],8) + np.left_shift(color[i,0], 0)
        # c = (color[i,2]<<16)|(color[i,1]<<8)|color[i,0]
        temp[i,3] = rgb
    # temp[:,3] = 1.6777215
    cloud0 = pcl.PointCloud_PointXYZRGBA()
    cloud0.from_array(temp)
    
    # color_cloud = pcl.PointCloud_PointXYZRGB(temp)
    # pcl.savePCD('scen.pcd', color_cloud) 
    visual = pcl.pcl_visualization.CloudViewing()
    visual.ShowColorACloud(cloud0)
    v = True
    while v:
        v = not(visual.WasStopped())

def bm(rimg1, rimg2):
    # rimg1 = cv2.cvtColor(rimg1, cv2.COLOR_RGB2GRAY)
    # rimg2 = cv2.cvtColor(rimg2, cv2.COLOR_RGB2GRAY)

    left_matcher = cv2.StereoBM_create(numDisparities=16, blockSize=9)
    left_matcher.setROI1(validPixROI1)
    left_matcher.setROI2(validPixROI2)
    left_matcher.setPreFilterCap(31)
    left_matcher.setBlockSize(15)
    left_matcher.setMinDisparity(0)
    left_matcher.setNumDisparities(128)
    left_matcher.setTextureThreshold(10)
    left_matcher.setUniquenessRatio(15)
    left_matcher.setSpeckleWindowSize(100)
    left_matcher.setSpeckleRange(1)
    left_matcher.setDisp12MaxDiff(1)
    right_matcher = cv2.ximgproc.createRightMatcher(left_matcher)
    lmbda = 8000
    sigma = 1.5
    wls_filter = cv2.ximgproc.createDisparityWLSFilter(matcher_left=left_matcher)
    wls_filter.setLambda(lmbda)
    wls_filter.setSigmaColor(sigma)
    t0 = time.time()
    displ = left_matcher.compute(rimg1, rimg2)
    dispr = right_matcher.compute(rimg2, rimg1)
    t1 = time.time()
    print('BM计算时间',t1-t0)
    # displ_show = cv2.normalize(displ,displ,alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    # dispr_show = cv2.normalize(dispr,dispr,alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    # cv2.imshow('L', displ_show)
    # cv2.imshow('R', dispr_show)
    displ = np.int16(displ)
    dispr = np.int16(dispr)
    t0 = time.time()
    disparity = wls_filter.filter(displ, rimg1, None, dispr)
    t1 = time.time()
    print('BM融合时间',t1-t0)
    return disparity

def sgbm(rimg1, rimg2):
    # rimg1 = cv2.cvtColor(rimg1, cv2.COLOR_RGB2GRAY)
    # rimg2 = cv2.cvtColor(rimg2, cv2.COLOR_RGB2GRAY)
    maxd = 128
    window_size = 5
    left_matcher = cv2.StereoSGBM_create(
        minDisparity=-maxd,
        numDisparities=maxd * 2,
        blockSize=5,
        P1=8 * 3 * window_size ** 2,
        P2=32 * 3 * window_size ** 2,
        disp12MaxDiff=1,
        uniquenessRatio=15,
        speckleWindowSize=100,
        speckleRange=2,
        preFilterCap=31,
        mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY
    )
    right_matcher = cv2.ximgproc.createRightMatcher(left_matcher)
    lmbda = 8000
    sigma = 1.5
    wls_filter = cv2.ximgproc.createDisparityWLSFilter(matcher_left=left_matcher)
    wls_filter.setLambda(lmbda)
    wls_filter.setSigmaColor(sigma)
    t0 = time.time()
    displ = left_matcher.compute(rimg1, rimg2)
    dispr = right_matcher.compute(rimg2, rimg1)
    t1 = time.time()
    print('SGBM计算时间',t1-t0)
    # displ_show = cv2.normalize(displ,displ,alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    # dispr_show = cv2.normalize(dispr,dispr,alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    # cv2.imshow('L', displ_show)
    # cv2.imshow('R', dispr_show)
    displ = np.int16(displ)
    dispr = np.int16(dispr)
    t0 = time.time()
    disparity = wls_filter.filter(displ, rimg1, None, dispr)
    t1 = time.time()
    print('SGBM融合时间',t1-t0)
    return disparity

def dist(threeD,x,y):
    out=[]
    for i in range(v_size):
        for j in range(v_size):
            if Dst_mode == 'Z':
                temp = threeD[y-v_size//2+i][x-v_size//2+j][2]/1000
                out.append(temp)
            else:
                temp = math.sqrt(threeD[y-v_size//2+i][x-v_size//2+j][0] ** 2 + threeD[y-v_size//2+i][x-v_size//2+j][1] ** 2 + threeD[y-v_size//2+i][x-v_size//2+j][2] ** 2)/1000
                out.append(temp)
    return sum(out)/len(out)

def augment(xyzs):
    axyz = np.ones((len(xyzs), 4))
    axyz[:, :3] = xyzs
    return axyz

def estimate(xyzs):
    axyz = augment(xyzs[:3])
    return np.linalg.svd(axyz)[-1][-1, :]

def is_inlier(coeffs, xyz, threshold):
    return np.abs(coeffs.dot(augment([xyz]).T)) < threshold

def run_ransac(data, estimate, is_inlier, sample_size, goal_inliers, max_iterations, stop_at_goal=True, random_seed=None):
    best_ic = 0
    best_model = None
    random.seed(random_seed)
    # random.sample cannot deal with "data" being a numpy array
    data = list(data)
    for i in range(max_iterations):
        s = random.sample(data, int(sample_size))
        m = estimate(s)
        ic = 0
        for j in range(len(data)):
            if is_inlier(m, data[j]):
                ic += 1

        print(s)
        print('estimate:', m,)
        print('# inliers:', ic)

        if ic > best_ic:
            best_ic = ic
            best_model = m
            if ic > goal_inliers and stop_at_goal:
                break
    print('took iterations:', i+1, 'best model:', best_model, 'explains:', best_ic)
    return best_model, best_ic

# 鼠标回调函数
def onmouse_pick_points(event, x, y, flags, param):
    global flag_mouse_lbutton_down, point_x, point_y
    if leastsq_mode:
        if event == cv2.EVENT_LBUTTONDOWN:
            point_x = x
            point_y = y
        if event == cv2.EVENT_LBUTTONUP:
            threeD = param
            threeD = threeD[point_y:y,point_x:x,:]
            threeD = threeD.reshape(threeD.shape[0]*threeD.shape[1],threeD.shape[2])
            mean = np.mean(threeD[:,2])
            std = np.std(threeD[:,2])
            temp_three = threeD[:,2]
            list = np.where((temp_three>mean+2*std)|(temp_three<mean-2*std))
            threeD_1 = np.delete(threeD, list, axis = 0)
            test_tr = threeD_1[:,2]
            goal_inliers = threeD_1.shape[0] * 0.5
            max_iterations = 20
            m, b = run_ransac(threeD_1, estimate, lambda x, y: is_inlier(x, y, 0.01), 3, goal_inliers, max_iterations)
            a, b, c, d = m
            # z = p[0]*np.mean(threeD[:,0])+p[1]*np.mean(threeD[:,1])+p[2]
            z = (-d - a * threeD_1[:,0] - b * threeD_1[:,1]) / c
            print(point_x,x,point_y,y,-d/c)
            print("中心的距离为：%.3f, %.3f"%(np.median(z),np.mean(z)), "m")
            a= 1    
    else:
        if event == cv2.EVENT_LBUTTONDOWN:
            for i in range(len(param)):
                threeD = param[i]
                # threeD2 = param[1]

                print('\n%d 像素坐标 x = %d, y = %d' % (i, x, y))
                print("世界坐标xyz 是：(%.3f，%.3f，%.3f)"%(threeD[y][x][0] / 1000.0, threeD[y][x][1] / 1000.0, threeD[y][x][2] / 1000.0), "m")
                distance = dist(threeD, x, y)
                print("%dx%d平均距离是：%.3f"%(v_size,v_size,distance),"m")


WIN_NAME = 'RGB'
cv2.namedWindow(WIN_NAME, cv2.WINDOW_AUTOSIZE)


if camera=='kitti':
    imgL = cv2.imread('/media/ye/File/KITTI/data_stereo_flow/testing/image_0/000000_10.png')
    imgR = cv2.imread('/media/ye/File/KITTI/data_stereo_flow/testing/image_1/000000_10.png')
else:
    frame = cv2.imread("/media/ye/File/Stereo/SGBM/data/8-4/test2/8000.jpg")
    img_color = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    w = frame.shape[1]
    frame1 = frame[:, :int(w/2)]
    frame2 = frame[:, int(w/2):]  # 割开双目图像

    imgL = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)  # 将BGR格式转换成灰度图片
    imgR = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
    # cv2.imshow('L', imgL)
    # cv2.imshow('R', imgR)
    # cv2.waitKey(0)

# cv2.remap 重映射，就是把一幅图像中某位置的像素放置到另一个图片指定位置的过程。
# 依据MATLAB测量数据重建无畸变图片
img1_rectified = cv2.remap(imgL, left_map1, left_map2, cv2.INTER_LINEAR)
img2_rectified = cv2.remap(imgR, right_map1, right_map2, cv2.INTER_LINEAR)

# cv2.imshow('L', img1_rectified)
# cv2.imshow('R', img2_rectified)
# cv2.waitKey(0)

imgL_rectified = cv2.remap(frame1, left_map1, left_map2, cv2.INTER_LINEAR)

# cv2.imwrite('data/8-4/imgL.png', imgL_rectified)


# t0 = time.time()
for i in range(1):
    disparity2 = bm(img1_rectified, img2_rectified)
disparity2 = np.where(disparity2 > 0, disparity2, 0)
# t1 = time.time()
# disp = cv2.normalize(disparity, disparity, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)  # 归一化函数算法
disp2 = cv2.normalize(disparity2, disparity2, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)  # 归一化函数算法
# cv2.imwrite('data/8-4/bm.png', disp2)
for i in range(1):
    disparity3 = sgbm(img1_rectified, img2_rectified)
disparity3 = np.where(disparity3 > 0, disparity3, 0)
disp3 = cv2.normalize(disparity3, disparity3, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)  # 归一化函数算法
# cv2.imwrite('sbbm.png', disp3)
    # print('2: ', t1-t0)

disparity_crop = sgbm(img1_rectified[120:600,320:960], img2_rectified[120:600,320:960])
disparity_crop = np.where(disparity_crop > 0, disparity_crop, 0)
disparity_crop = np.pad(disparity_crop, [[120,120], [320,320]])
disp_crop = cv2.normalize(disparity_crop, disparity_crop, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)  # 归一化函数算法

# cv2.imwrite('sbbm_crop.png', disp_crop)


threeD2 = cv2.reprojectImageTo3D(disparity2, Q, handleMissingValues=False)  # 计算三维坐标数据值
threeD2 = threeD2 * 16
np.save('three_bm.npy', threeD2[379:629,435:863,:])
threeD3 = cv2.reprojectImageTo3D(disparity3, Q, handleMissingValues=True)  # 计算三维坐标数据值
threeD3 = threeD3 * 16
np.save('three_sgbm.npy', threeD3[379:629,435:863,:])



if gt_pred:
    if camera=='kitti':
        disparity3 = cv2.imread('/media/ye/File/LEAStereo/predict/kitti2012/images/000000_10.png', cv2.IMREAD_UNCHANGED)

    elif camera=='720p':
        disparity3 = cv2.imread('/media/ye/File/LEAStereo/predict/test_15/4.png', cv2.IMREAD_UNCHANGED)
        disparity3 = np.pad(disparity3, [[72,72], [160,160]])
        
    disparity3 = (disparity3/16).astype(np.int16)
    disp3 = cv2.normalize(disparity3, disparity3, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)  # 归一化函数算法
    threeD3 = cv2.reprojectImageTo3D(disparity3, Q, handleMissingValues=True)  # 计算三维坐标数据值
    threeD3 = threeD3 * 16
    threeD4 = (threeD2,threeD3)
    cv2.imshow('Ref', disp3)
else:
    threeD4 = threeD2
# threeD[y][x] x:0~640; y:0~480;   !!!!!!!!!!
cv2.setMouseCallback(WIN_NAME, onmouse_pick_points, threeD4)  # 鼠标回调事件


dis_color = threeD2[...,2]
dis_color[np.isinf(dis_color)]=0
dis_color[dis_color >  9000] =9000

dis_color = cv2.normalize(dis_color, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
# dis_color = cv2.applyColorMap(dis_color, 2)
# cv2.imshow("depth", dis_color)

# cv2.imshow("left", frame1)
# cv2.imshow('single', disp)
cv2.imshow(WIN_NAME, imgL_rectified)
# cv2.imshow('sgbm', disp3)
cv2.imshow('bm', disp2)# 显示深度图的双目画面
cv2.waitKey()
# 销毁内存
cv2.destroyAllWindows()

pcl_show(threeD3, imgL_rectified)

sgbm_crop.py

'''
Author: YeHanyu
Date: 2022-08-02 11:32:58
LastEditors: YeHanyu
LastEditTime: 2022-08-18 16:42:16
FilePath: /undefined/media/ye/File/Stereo/SGBM/sgbm_crop.py
Description: 

Copyright (c) 2022 by YeHanyu, All Rights Reserved. 
'''

import cv2
import numpy as np
import time
from sklearn.linear_model import RANSACRegressor, LinearRegression
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d

global numDisp
numDisp = 80

# 相机参数读取
class cam_cfg:
    def __init__(self, config_path):
        self.config_path = config_path
        camera_config = cv2.FileStorage(config_path, cv2.FILE_STORAGE_READ)
        self.size = camera_config.getNode("Size").mat().astype(int)[0]
        self.left_matrix = camera_config.getNode("KL").mat()
        self.right_matrix = camera_config.getNode("KR").mat()
        self.left_distortion = camera_config.getNode("DL").mat()
        self.right_distortion = camera_config.getNode("DR").mat()
        self.R = camera_config.getNode("R").mat()
        self.T = camera_config.getNode("T").mat()


def bm(l_img, r_img):
    # l_img = cv2.cvtColor(l_img, cv2.COLOR_RGB2GRAY)
    # r_img = cv2.cvtColor(r_img, cv2.COLOR_RGB2GRAY)
    numDisp = 80
    l_img = np.pad(l_img, [[0, 0], [numDisp, numDisp]])  # pad防止crop溢出
    r_img = np.pad(r_img, [[0, 0], [numDisp, numDisp]])
    left_matcher = cv2.StereoBM_create(numDisparities=16, blockSize=9)
    left_matcher.setROI1(validPixROI1)
    left_matcher.setROI2(validPixROI2)
    left_matcher.setPreFilterCap(31)
    left_matcher.setBlockSize(15)
    left_matcher.setMinDisparity(0)
    left_matcher.setNumDisparities(numDisp - 22)
    left_matcher.setTextureThreshold(10)
    left_matcher.setUniquenessRatio(15)
    left_matcher.setSpeckleWindowSize(100)
    left_matcher.setSpeckleRange(1)
    left_matcher.setDisp12MaxDiff(1)
    right_matcher = cv2.ximgproc.createRightMatcher(left_matcher)
    lmbda = 8000
    sigma = 1.5
    wls_filter = cv2.ximgproc.createDisparityWLSFilter(matcher_left=left_matcher)
    wls_filter.setLambda(lmbda)
    wls_filter.setSigmaColor(sigma)
    t0 = time.time()
    displ = left_matcher.compute(l_img[:, 0:-numDisp], r_img[:, 0:-numDisp])
    dispr = right_matcher.compute(r_img[:, numDisp:], l_img[:, numDisp:])
    t1 = time.time()
    print('BM计算时间: %.2f ms' % ((t1 - t0) * 1000))
    # displ_show = cv2.normalize(displ,displ,alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    # dispr_show = cv2.normalize(dispr,dispr,alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    # cv2.imshow('L', displ_show)
    # cv2.imshow('R', dispr_show)
    displ = np.int16(displ)
    dispr = np.int16(dispr)
    t0 = time.time()
    dispr = np.pad(dispr[:, :-numDisp], [[0, 0], [numDisp, 0]])
    disparity = wls_filter.filter(displ, l_img[:, :-numDisp], None, dispr)
    disparity = np.where(disparity > 0, disparity, 0)
    disp_out = disparity[:, numDisp:]
    t1 = time.time()
    print('BM融合时间: %.2f ms' % ((t1 - t0) * 1000))
    return disp_out


def sgbm(l_img, r_img):

    if l_img.shape[1]==1920:
        l_img = np.pad(l_img, [[0, 0], [numDisp, numDisp]])  # 去除匹配时的黑边
        r_img = np.pad(r_img, [[0, 0], [numDisp, numDisp]])
    left_matcher = cv2.StereoSGBM_create(
        minDisparity=1,
        numDisparities=int(numDisp/1.25),
        blockSize=15,
        P1=8 * 3 * 5 ** 2,
        P2=32 * 3 * 5 ** 2,
        disp12MaxDiff=1,
        uniquenessRatio=10,
        speckleWindowSize=100,
        speckleRange=10,
        preFilterCap=31,
        mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY
    )
    right_matcher = cv2.ximgproc.createRightMatcher(left_matcher)
    lmbda = 8000
    sigma = 1.5
    wls_filter = cv2.ximgproc.createDisparityWLSFilter(matcher_left=left_matcher)
    wls_filter.setLambda(lmbda)
    wls_filter.setSigmaColor(sigma)
    wls_filter2 = cv2.ximgproc.createDisparityWLSFilter(matcher_left=right_matcher)
    wls_filter2.setLambda(lmbda)
    wls_filter2.setSigmaColor(sigma)
    t0 = time.time()
    displ = left_matcher.compute(l_img, r_img)
    dispr = right_matcher.compute(r_img, l_img)
    # displ = left_matcher.compute(l_img, r_img)
    # dispr = right_matcher.compute(r_img, l_img)
    t1 = time.time()
    print('SGBM计算时间: %.2f ms' % ((t1 - t0) * 1000))
    # displ_show = cv2.normalize(displ,displ,alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    # dispr_show = cv2.normalize(dispr,dispr,alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    # cv2.imshow('L', displ_show)
    # cv2.imshow('R', dispr_show)
    displ = np.int16(displ)
    # cv2.imshow('displ', norm_disp(displ))
    dispr = np.int16(dispr)
    # cv2.imshow('dispr', norm_disp(-dispr))
    t0 = time.time()

    disparity = wls_filter.filter(displ, l_img, None, dispr)
    disparity2 = -wls_filter2.filter(dispr, r_img, None, displ)
    disparity = np.where(disparity > 0, disparity, 0)
    disparity2 = np.where(disparity2 > 0, disparity2, 0)
    # cv2.imshow('disparity', norm_disp(disparity))
    t1 = time.time()
    print('SGBM融合时间: %.2f ms' % ((t1 - t0) * 1000))
    disp_out = disparity[:, numDisp:-numDisp]
    disp_out2 = disparity2[:, numDisp:-numDisp]
    return disp_out, disp_out2


'''
description:  计算ROI区域视差，并使用RANSAC进行平面拟合，再计算3D坐标

'''


def calc_dist(x_L, y_L, x_R, y_R):

    l_img = np.pad(img1_rectified, [[0, 0], [numDisp, numDisp]])  # 防止索引为空，放大图像
    r_img = np.pad(img2_rectified, [[0, 0], [numDisp, numDisp]])
    disp_crop, disp_crop2 = sgbm(l_img[y_L:y_R, x_L:x_R + 2 * numDisp], r_img[y_L:y_R, x_L:x_R + 2 * numDisp])  # 使用更大的ROI区域进行计算

    cv2.namedWindow("DISP-CROP", 0)
    cv2.resizeWindow("DISP-CROP", 640, 640)
    cv2.imshow('DISP-CROP', norm_disp(disp_crop))
    cv2.namedWindow("DISP-CROP2", 0)
    cv2.resizeWindow("DISP-CROP2", 640, 640)
    cv2.imshow('DISP-CROP2', norm_disp(disp_crop2))
    cv2.waitKey(0)
    # 提取ROI区域
    # 为disp添加x,y坐标
    h = disp_crop.shape[0]
    w = disp_crop.shape[1]
    x = np.arange(0, w)
    y = np.arange(0, h)
    X, Y = np.meshgrid(x, y)
    disp_3d = np.stack((X, Y, disp_crop), axis=2)
    disp_3d_o = np.stack((X, Y, disp_crop), axis=2)
    # 去除异常点
    disp_3d = disp_3d.reshape(disp_3d.shape[0] * disp_3d.shape[1], disp_3d.shape[2])
    mean = np.mean(disp_3d[:, 2])
    std = np.std(disp_3d[:, 2])
    temp_three = disp_3d[:, 2]
    list = np.where((temp_three > mean + std) | (temp_three < mean - std))
    disp_3d_tmp = np.delete(disp_3d, list, axis=0)
    #  原始的disp点查看
    fig = plt.figure()
    ax = mplot3d.Axes3D(fig, auto_add_to_figure=False)
    fig.add_axes(ax)
    # color_map = plt.get_cmap('jet')
    # ax.scatter3D(disp_3d[:,0], disp_3d[:,1], disp_3d[:,2],  c=disp_3d[:,2],cmap = color_map)
    # plt.show()

    # RANSAC拟合平面
    ransac_disp = RANSACRegressor(LinearRegression(), max_trials=50)
    t0 = time.time()
    ransac_disp.fit(disp_3d_tmp[:, 0:2], disp_3d_tmp[:, -1])
    t1 = time.time()
    print('RANSAC计算时间: %.2f ms' % ((t1 - t0) * 1000))
    z_disp = ransac_disp.predict(disp_3d[:, 0:2])
    z_disp = z_disp.reshape(h, w).astype(np.float32)
    disp_gt_sub = disp_gt[y_L:y_R, x_L:x_R]
    disp_diff = z_disp / 16 - disp_gt_sub
    disp_diff = np.delete(disp_diff, np.where(np.isinf(disp_diff[:, :]))[0])
    disp_diff = np.delete(disp_diff, np.where(np.isneginf(disp_diff[:]))[0])
    print(disp_diff.mean(), disp_diff.std())
    # cv2.imshow('DISP2', norm_disp(disp_diff))
    # cv2.waitKey(0)
    # disp_ref = disp[y_L:y_R,x_L:x_R]  # 参考用，对比裁切后disp和原disp
    # diff = disp_ref - disp_crop
    # ax.plot_surface(disp_3d_o[:,:,0], disp_3d_o[:,:,1], z_disp, rstride=5, cstride=5, color ='lightpink')
    # plt.show()
    # 计算3D坐标
    T = np.eye(4, dtype=float)  # 用于将ROI区域映射回原图所在位置，计算3D坐标
    T[0, 3] = x_L
    T[1, 3] = y_L
    xyz_d = cv2.reprojectImageTo3D(z_disp, Q @ T, handleMissingValues=False)
    xyz_d = xyz_d * 16

    # 输出结果 ref为原始完整disp计算的全图3D坐标
    xyz_ref = xyz[y_L:y_R, x_L:x_R, :]
    print("距离中位数：%.1f  mm, 距离平均数：%.1f  mm" % (np.median(xyz_d[:, :, 2]), np.mean(xyz_d[:, :, 2])))
    print('左上真实坐标: ', xyz_d[0, 0], '右下真实坐标: ', xyz_d[-1, -1])

    print("预测对角线距离: ", np.sqrt(np.sum((xyz_d[0, 0] - xyz_d[-1, -1]) ** 2)))

    # 显示点云与拟合后的点云
    color_map = plt.get_cmap('jet')
    ax.plot_surface(xyz_d[:, :, 0], xyz_d[:, :, 1], xyz_d[:, :, 2], rstride=50, cstride=50, color='lightpink')
    # ax.scatter3D(xyz_d[:,:,0], xyz_d[:,:,1], xyz_d[:,:,2],  c='b')
    ax.scatter3D(xyz_ref[:, :, 0], xyz_ref[:, :, 1], xyz_ref[:, :, 2], c=xyz_ref[:, :, 2], cmap=color_map)
    plt.show()


# 鼠标回调函数
def onmouse_pick_points(event, x, y, flags, param):
    global x_L, y_L
    if event == cv2.EVENT_LBUTTONDOWN:
        # t0 = time.time()
        # temp= np.argwhere((left_map1[:,:,0]==x)&(left_map1[:,:,1]==y))  # 提取remap后检测框的坐标
        # i = 1
        # while temp.size == 0:  # remap后坑可能存在空映射，向右下方寻找寻找最近的映射点
        #     temp = np.argwhere((left_map1[:,:,0]<=x+i)&(left_map1[:,:,0]>=x)&(left_map1[:,:,1]<=y+i)&(left_map1[:,:,1]>=y))
        #     i+=1
        # y_L,x_L = temp[0]
        # t1 = time.time()
        # print('Remap坐标时间: %.2f ms'%((t1-t0)*1000))
        x_L, y_L = x, y                        # 直到找到一个最近的有效点。操作在原图上，因此需要借助map_rev得到在校正之后的坐标

        # x_L, y_L = map_rev[y, x]                         # 直到找到一个最近的有效点。操作在原图上，因此需要借助map_rev得到在校正之后的坐标
        # num = 1
        # while x_L == -1:
        #     x_L, y_L = map_rev[y + num, x + num]
        #     num += 1

    if event == cv2.EVENT_LBUTTONUP:
        # temp= np.argwhere((left_map1[:,:,0]==x)&(left_map1[:,:,1]==y))
        # i = 1
        # while temp.size == 0: # remap后坑可能存在空映射，向左上方寻找寻找最近的映射点
        #     temp = np.argwhere((left_map1[:,:,0]<=x)&(left_map1[:,:,0]>=x-i)&(left_map1[:,:,1]<=y)&(left_map1[:,:,1]>=y-i))
        #     i+=1
        # y_R,x_R = temp[0]
        x_R, y_R = x, y
        # x_R, y_R = map_rev[y, x]
        # num = 1
        # while x_R == -1:
        #     x_R, y_R = map_rev[y - num, x - num]
        #     num += 1
        # print('左上像素坐标: ', x_L, y_L, ' 右下像素坐标: ', x_R, y_R)
        calc_dist(x_L, y_L, x_R, y_R)  # 计算视差等数据


# RANSAC 3D点云代码
def calc_3d_dist(x_L, y_L, x_R, y_R):
    disp_ref = disp
    T = np.eye(4, dtype=float)
    T[0, 3] = x_L
    T[1, 3] = y_L
    t0 = time.time()
    xyz = cv2.reprojectImageTo3D(disp_ref, Q @ T, handleMissingValues=False)  # 计算三维坐标数据值
    t1 = time.time()
    print('3D坐标计算时间: %.2f ms' % ((t1 - t0) * 1000))
    xyz = xyz * 16
    xyz = xyz.reshape(xyz.shape[0] * xyz.shape[1], xyz.shape[2])
    tmp = xyz[:, 2]
    inf_list = np.where(tmp == np.inf)
    xyz = np.delete(xyz, inf_list, axis=0)
    mean = np.mean(xyz[:, 2])
    std = np.std(xyz[:, 2])
    temp_three = xyz[:, 2]
    list = np.where((temp_three > mean + std) | (temp_three < mean - std))
    xyz = np.delete(xyz, list, axis=0)
    test_tr = xyz[:, 2]
    # goal_inliers = xyz.shape[0] * 0.4
    # max_iterations = 5
    # t0 = time.time()
    # m, b = run_ransac(xyz, estimate, lambda x, y: is_inlier(x, y, 0.01), 3, goal_inliers, max_iterations)
    # t1 = time.time()
    t0 = time.time()
    lr = LinearRegression().fit(xyz[:, 0:2], xyz[:, -1])
    t1 = time.time()
    print('最小二乘计算时间: %.2f ms' % ((t1 - t0) * 1000))
    z = lr.predict(xyz[:, 0:2])
    print("距离中位数：%.1f  mm, 距离平均数：%.1f  mm, 原平均数：%.1f  mm" % (np.median(z), np.mean(z), mean))

    stride = xyz.shape[0] // 10000 + 1
    xyz_s = xyz[::stride, :]
    ransac = RANSACRegressor(LinearRegression(), max_trials=100, residual_threshold=mean * 0.01,
                             stop_n_inliers=xyz_s.shape[0] // 1.1, stop_probability=1)
    t0 = time.time()
    ransac.fit(xyz_s[:, 0:2], xyz_s[:, -1])
    t1 = time.time()
    print('RANSAC计算时间: %.2f ms' % ((t1 - t0) * 1000))
    # 获取内点和异常点集合
    inlier_mask = ransac.inlier_mask_

    # a, b, c, d = m
    # z = p[0]*np.mean(xyz[:,0])+p[1]*np.mean(xyz[:,1])+p[2]
    # z = (-d - a * xyz[:,0] - b * xyz[:,1]) / c
    list = np.where(inlier_mask == 0)
    xyz_ss = np.delete(xyz_s, list, axis=0)
    z = ransac.predict(xyz_ss[:, 0:2])
    dif = xyz_ss[:, -1] - z
    # xyz_ss[:,-1] = z
    x_median = np.median(xyz_ss[:, 0])
    y_median = np.median(xyz_ss[:, 1])
    z_median = ransac.predict(((x_median, y_median), (0, 0)))

    print("距离中位数：%.1f mm, 距离平均数：%.1f  mm, 原平均数：%.1f  mm" % (np.median(z), np.mean(z), z_median[0]))


def norm_disp(disp):
    # 归一化显示视差
    disp[np.isinf(disp)] = 0
    disp = np.where(disp < 0, 0, disp)
    disp = np.where(disp > 200, 200, disp)
    mean = np.mean(disp)
    std = np.std(disp)
    disp_norm = np.where(disp < mean + 2 * std, disp, mean + 2 * std)
    disp_norm = cv2.normalize(disp_norm, disp_norm, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX,
                              dtype=cv2.CV_8U)  # 归一化函数算法
    disp_norm = cv2.applyColorMap(disp_norm, 14)
    return disp_norm


if __name__ == '__main__':

    # 读取相机参数
    camera_cfg = cam_cfg('/mnt/sda1/Stereo/data/config/cam25/camera24.yml')
    # 双目极线校正
    R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(camera_cfg.left_matrix,
                                                                      camera_cfg.left_distortion,
                                                                      camera_cfg.right_matrix,
                                                                      camera_cfg.right_distortion, camera_cfg.size,
                                                                      camera_cfg.R,
                                                                      camera_cfg.T,
                                                                      flags=cv2.CALIB_ZERO_DISPARITY,
                                                                      alpha=-1)

    # 校正查找映射表,将原始图像和校正后的图像上的点一一对应起来
    left_map1, left_map2 = cv2.initUndistortRectifyMap(camera_cfg.left_matrix, camera_cfg.left_distortion, R1, P1,
                                                       camera_cfg.size, cv2.CV_16SC2)
    right_map1, right_map2 = cv2.initUndistortRectifyMap(camera_cfg.right_matrix, camera_cfg.right_distortion, R2, P2,
                                                         camera_cfg.size, cv2.CV_16SC2)

    map_rev = -np.ones(left_map1.shape, dtype=np.int16)
    t0 = time.time()

    for y, list in enumerate(left_map1):   # 每一行
        for x, point in enumerate(list):    # 每个点
            if point[1] < left_map1.shape[0] and point[0] < left_map1.shape[1]: # 有效区域映射回去
                map_rev[point[1], point[0]] = (x, y)
            # print('1')
    t1 = time.time()
    print('remap_rev', (t1 - t0) * 1000)



    WIN_NAME = 'RGB'
    cv2.namedWindow(WIN_NAME, cv2.WINDOW_AUTOSIZE)

    img = cv2.imread('/mnt/sda1/Stereo/data/imgs/00.jpg')
    w = img.shape[1]
    img_R = img[:, :int(w / 2)]
    img_L = img[:, int(w / 2):]  # 割开双目图像

    # cv2.imshow('RGB_L', img_L)
    imgL = cv2.cvtColor(img_L, cv2.COLOR_BGR2GRAY)  # 将BGR格式转换成灰度图片
    imgR = cv2.cvtColor(img_R, cv2.COLOR_BGR2GRAY)

    # cv2.imshow('L', imgL)
    # cv2.imshow('R', imgR)
    # cv2.waitKey(0)
    # l_tmp = left_map1[400:800, 400:800, :]-(400, 400)
    # r_tmp = right_map1[400:800, 400:800, :]-(400, 400)
    # l_tmp2 = left_map2[400:800, 400:800]
    # r_tmp2 = right_map2[400:800, 400:800]
    # l_tmpimg = imgL[400:800, 400:800]
    # r_tmpimg = imgR[400:800, 400:800]
    # img1_re = cv2.remap(l_tmpimg, l_tmp.astype(np.int16), l_tmp2, cv2.INTER_LINEAR)
    # img2_re = cv2.remap(r_tmpimg, r_tmp.astype(np.int16), r_tmp2, cv2.INTER_LINEAR)


    # cv2.remap 重映射，就是把一幅图像中某位置的像素放置到另一个图片指定位置的过程。
    # 依据MATLAB测量数据重建无畸变图片
    t0 = time.time()
    img1_rectified = cv2.remap(imgL, left_map1, left_map2, cv2.INTER_LINEAR)
    img2_rectified = cv2.remap(imgR, right_map1, right_map2, cv2.INTER_LINEAR)
    # cv2.imwrite('/mnt/sda1/CREStereo/img/left.png', img1_rectified)
    # cv2.imwrite('/mnt/sda1/CREStereo/img/right.png', img2_rectified)
    t1 = time.time()
    print('remap', (t1 - t0) * 1000)
    # cv2.imshow('L', img1_rectified)
    # cv2.imshow('R', img1_re)
    # cv2.waitKey(0)

    # mlimg1 = cv2.imread('/mnt/sda1/Stereo/data/cam-test/cam1-1920-1080-2.1mm/image0/04.png')
    # mlimg2 = cv2.imread('/mnt/sda1/Stereo/data/cam-test/cam1-1920-1080-2.1mm/image1/04.png')
    # disp2 = sgbm(mlimg1, mlimg2)
    # 原始图计算视差，参考用
    # img1_rectified = cv2.imread("/media/hx/datav/data/octogons2/im0.png", 0)
    #
    # img2_rectified = cv2.imread("/media/hx/datav/data/octogons2/im1.png", 0)
    # disp = (sgbm(img1_rectified, img2_rectified) / 16).astype(np.float32)
    # import readpfm
    # disp_gt, scale = readpfm.readPFM("/media/hx/datav/data/octogons2/disp0.pfm")
    # cv2.imwrite("disp_gt.png", norm_disp(disp_gt))
    # cv2.imwrite("disp_sgbm.png", norm_disp(disp))

    # disp_diff = disp - disp_gt
    # disp_diff = np.delete(disp_diff, np.where(np.isinf(disp_diff[:, :]))[0])
    # disp_diff = np.delete(disp_diff, np.where(np.isneginf(disp_diff[:]))[0])

    # np.save('/mnt/sda1/CREStereo/ref.npy', disp)
    # np.save('ref.npy', disp)

# 3D坐标计算
#     xyz = cv2.reprojectImageTo3D(disp, Q, handleMissingValues=True)  # 计算三维坐标数据值
#     xyz = xyz * 16

    cv2.setMouseCallback(WIN_NAME, onmouse_pick_points, 0)  # 鼠标回调事件

    cv2.imshow(WIN_NAME, img1_rectified)
    # cv2.imshow('DISP', norm_disp(disp))  # 显示深度图的双目画面
    # cv2.imshow('DISP2', norm_disp(disp2))  # 显示深度图的双目画面
    cv2.waitKey()
    # 销毁内存
    cv2.destroyAllWindows()