04_visual.py

import cv2
import numpy as np

# === Load Q Matrix ===
Q = np.load("q_matrix.npy")

# === Load Rectification Maps ===
left_map1 = np.load("left_map1.npy")
left_map2 = np.load("left_map2.npy")
right_map1 = np.load("right_map1.npy")
right_map2 = np.load("right_map2.npy")

# === Stereo Camera IDs ===
LEFT_ID = 2
RIGHT_ID = 4

# === StereoSGBM Setup ===
stereo = cv2.StereoSGBM_create(
    minDisparity=0,
    numDisparities=16 * 6,  # Must be divisible by 16
    blockSize=5,
    P1=8 * 3 * 5 ** 2,
    P2=32 * 3 * 5 ** 2,
    disp12MaxDiff=1,
    uniquenessRatio=10,
    speckleWindowSize=100,
    speckleRange=32
)

# === Video Writer Setup ===
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter("depth_demo.avi", fourcc, 10.0, (1280, 480))

clicked_point = None

# === Mouse Callback ===
def mouse_callback(event, x, y, flags, param):
    global clicked_point
    if event == cv2.EVENT_LBUTTONDOWN:
        if x >= 640:  # Only if click is on disparity map
            clicked_point = (x - 640, y)

cv2.namedWindow("RGB (Left) + Depth Map")
cv2.setMouseCallback("RGB (Left) + Depth Map", mouse_callback)

print("[*] Press 'q' to stop recording.")

camL = cv2.VideoCapture(LEFT_ID)
camR = cv2.VideoCapture(RIGHT_ID)

while True:
    retL, frameL = camL.read()
    retR, frameR = camR.read()

    if not retL or not retR:
        print("Frame grab failed.")
        break

    rectL = cv2.remap(frameL, left_map1, left_map2, cv2.INTER_LINEAR)
    rectR = cv2.remap(frameR, right_map1, right_map2, cv2.INTER_LINEAR)

    grayL = cv2.cvtColor(rectL, cv2.COLOR_BGR2GRAY)
    grayR = cv2.cvtColor(rectR, cv2.COLOR_BGR2GRAY)

    disparity = stereo.compute(grayL, grayR).astype(np.float32) / 16.0

    # Reproject to 3D space
    points_3D = cv2.reprojectImageTo3D(disparity, Q)

    # Visualize disparity
    disp_vis = cv2.normalize(disparity, None, 0, 255, cv2.NORM_MINMAX)
    disp_vis = np.uint8(disp_vis)
    disp_color = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)

    # === Handle click and show distance ===
    if clicked_point:
        x, y = clicked_point
        if 0 <= x < points_3D.shape[1] and 0 <= y < points_3D.shape[0]:
            X, Y, Z = points_3D[y, x]
            if not np.isinf(Z) and not np.isnan(Z):
                distance = np.linalg.norm([X, Y, Z])
                cv2.circle(disp_color, (x, y), 5, (255, 255, 255), -1)
                cv2.putText(disp_color, f"{distance:.2f} m", (x + 10, y - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

    # Combine views
    combined = np.hstack((rectL, disp_color))
    out.write(combined)
    cv2.imshow("RGB (Left) + Depth Map", combined)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

camL.release()
camR.release()
out.release()
cv2.destroyAllWindows()