trashbot/rectification.py at main · fyang7705-sys/trashbot · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
import av
import cv2
import numpy as np
import threading

# =========================
# 配置
# =========================
DEVICE = r"video=USB2.0 Camera RGB"
W, H, FPS = 1280, 480, 60

# =========================
# 1. 加载标定参数
# =========================
data = np.load("stereo_calibration.npz")

mtx_l = data["mtx_l"]
dist_l = data["dist_l"]
mtx_r = data["mtx_r"]
dist_r = data["dist_r"]
R = data["R"]
T = data["T"]

print("标定参数加载成功")

# =========================
# 2. 计算 Rectification
# =========================
image_size = (640, 480)

R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(
    mtx_l, dist_l,
    mtx_r, dist_r,
    image_size,
    R, T,
    flags=cv2.CALIB_ZERO_DISPARITY,
    alpha=0
)

# =========================
# 3. 生成映射表（核心）
# =========================
map1_l, map2_l = cv2.initUndistortRectifyMap(
    mtx_l, dist_l, R1, P1, image_size, cv2.CV_16SC2)

map1_r, map2_r = cv2.initUndistortRectifyMap(
    mtx_r, dist_r, R2, P2, image_size, cv2.CV_16SC2)

print("Rectification 映射表生成完成")

# =========================
# 采集线程（PyAV）
# =========================
latest_frame = None
lock = threading.Lock()

def capture_thread():
    global latest_frame

    options = {
        "video_size": f"{W}x{H}",
        "framerate": str(FPS),
        "vcodec": "mjpeg",
        "rtbufsize": "64M",
    }

    container = av.open(DEVICE, format="dshow", options=options)
    stream = next(s for s in container.streams if s.type == "video")

    for packet in container.demux(stream):
        for frame in packet.decode():
            img = frame.to_ndarray(format="bgr24")
            with lock:
                latest_frame = img

# =========================
# 4. 主循环：实时校正
# =========================
def main():
    global latest_frame

    threading.Thread(target=capture_thread, daemon=True).start()

    print("双目校正预览启动（按 q 退出）")

    while True:
        with lock:
            if latest_frame is None:
                continue
            frame = latest_frame.copy()

        # 分割左右
        left = frame[:, :640]
        right = frame[:, 640:]

        # =========================
        # 关键：执行校正
        # =========================
        rect_l = cv2.remap(left, map1_l, map2_l, cv2.INTER_LINEAR)
        rect_r = cv2.remap(right, map1_r, map2_r, cv2.INTER_LINEAR)

        # =========================
        # 可视化验证（画水平线）
        # =========================
        for y in range(0, 480, 40):
            cv2.line(rect_l, (0, y), (640, y), (0, 255, 0), 1)
            cv2.line(rect_r, (0, y), (640, y), (0, 255, 0), 1)

        combined = np.hstack((rect_l, rect_r))

        cv2.imshow("Rectified Stereo", combined)

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

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()