RGB-Thermal Align

1#!/usr/bin/env python3
2
3"""
4Due to an issue with our calibration, you might receive the following error when running this script on early release OAK Thermal devices:
5```bash
6[ImageAlign(4)] [error] Failed to get calibration data: Extrinsic connection between the requested cameraId's doesn't exist. Please recalibrate or modify your calibration data
7```
8If this happens, please download the calibration data + script from https://drive.google.com/drive/folders/1Q_MZMqWMKDC1eOqVHGPeDO-NJgFmnY5U,
9place them into the same folder, connect the camera to the computer and run the script. This will update the
10calibration and add required extrinsics between the camera sensors.
11"""
12
13import cv2
14import depthai as dai
15import numpy as np
16import time
17from datetime import timedelta
18
19FPS = 25.0
20
21RGB_SOCKET = dai.CameraBoardSocket.CAM_A
22COLOR_RESOLUTION = dai.ColorCameraProperties.SensorResolution.THE_1080_P
23
24class FPSCounter:
25    def __init__(self):
26        self.frameTimes = []
27
28    def tick(self):
29        now = time.time()
30        self.frameTimes.append(now)
31        self.frameTimes = self.frameTimes[-100:]
32
33    def getFps(self):
34        if len(self.frameTimes) <= 1:
35            return 0
36        # Calculate the FPS
37        return (len(self.frameTimes) - 1) / (self.frameTimes[-1] - self.frameTimes[0])
38
39device = dai.Device()
40
41thermalWidth, thermalHeight = -1, -1
42thermalFound = False
43for features in device.getConnectedCameraFeatures():
44    if dai.CameraSensorType.THERMAL in features.supportedTypes:
45        thermalFound = True
46        thermalSocket = features.socket
47        thermalWidth, thermalHeight = features.width, features.height
48        break
49if not thermalFound:
50    raise RuntimeError("No thermal camera found!")
51
52
53ISP_SCALE = 3
54
55calibrationHandler = device.readCalibration()
56rgbDistortion = calibrationHandler.getDistortionCoefficients(RGB_SOCKET)
57distortionModel = calibrationHandler.getDistortionModel(RGB_SOCKET)
58if distortionModel != dai.CameraModel.Perspective:
59    raise RuntimeError("Unsupported distortion model for RGB camera. This example supports only Perspective model.")
60
61pipeline = dai.Pipeline()
62
63# Define sources and outputs
64camRgb = pipeline.create(dai.node.ColorCamera)
65thermalCam = pipeline.create(dai.node.Camera)
66thermalCam.setBoardSocket(thermalSocket)
67thermalCam.setFps(FPS)
68
69sync = pipeline.create(dai.node.Sync)
70out = pipeline.create(dai.node.XLinkOut)
71align = pipeline.create(dai.node.ImageAlign)
72cfgIn = pipeline.create(dai.node.XLinkIn)
73
74
75camRgb.setBoardSocket(RGB_SOCKET)
76camRgb.setResolution(COLOR_RESOLUTION)
77camRgb.setFps(FPS)
78camRgb.setIspScale(1,ISP_SCALE)
79
80out.setStreamName("out")
81
82sync.setSyncThreshold(timedelta(seconds=0.5 / FPS))
83
84cfgIn.setStreamName("config")
85
86cfg = align.initialConfig.get()
87staticDepthPlane = cfg.staticDepthPlane
88
89# Linking
90align.outputAligned.link(sync.inputs["aligned"])
91camRgb.isp.link(sync.inputs["rgb"])
92camRgb.isp.link(align.inputAlignTo)
93thermalCam.raw.link(align.input)
94sync.out.link(out.input)
95cfgIn.out.link(align.inputConfig)
96
97
98rgbWeight = 0.4
99thermalWeight = 0.6
100
101
102def updateBlendWeights(percentRgb):
103    """
104    Update the rgb and depth weights used to blend depth/rgb image
105    @param[in] percent_rgb The rgb weight expressed as a percentage (0..100)
106    """
107    global thermalWeight
108    global rgbWeight
109    rgbWeight = float(percentRgb) / 100.0
110    thermalWeight = 1.0 - rgbWeight
111
112def updateDepthPlane(depth):
113    global staticDepthPlane
114    staticDepthPlane = depth
115
116# Connect to device and start pipeline
117with device:
118    device.startPipeline(pipeline)
119    queue = device.getOutputQueue("out", 8, False)
120    cfgQ = device.getInputQueue("config")
121
122    # Configure windows; trackbar adjusts blending ratio of rgb/depth
123    windowName = "rgb-thermal"
124
125    # Set the window to be resizable and the initial size
126    cv2.namedWindow(windowName, cv2.WINDOW_NORMAL)
127    cv2.resizeWindow(windowName, 1280, 720)
128    cv2.createTrackbar(
129        "RGB Weight %",
130        windowName,
131        int(rgbWeight * 100),
132        100,
133        updateBlendWeights,
134    )
135    cv2.createTrackbar(
136        "Static Depth Plane [mm]",
137        windowName,
138        0,
139        2000,
140        updateDepthPlane,
141    )
142    fpsCounter = FPSCounter()
143    while True:
144        messageGroup = queue.get()
145        assert isinstance(messageGroup, dai.MessageGroup)
146        frameRgb = messageGroup["rgb"]
147        assert isinstance(frameRgb, dai.ImgFrame)
148        thermalAligned = messageGroup["aligned"]
149        assert isinstance(thermalAligned, dai.ImgFrame)
150        frameRgbCv = frameRgb.getCvFrame()
151        fpsCounter.tick()
152
153        rgbIntrinsics = calibrationHandler.getCameraIntrinsics(RGB_SOCKET, int(frameRgbCv.shape[1]), int(frameRgbCv.shape[0]))
154
155        cvFrameUndistorted = cv2.undistort(
156            frameRgbCv,
157            np.array(rgbIntrinsics),
158            np.array(rgbDistortion),
159        )
160
161        # Colorize the aligned depth
162        thermalFrame = thermalAligned.getCvFrame().astype(np.float32)
163        # Create a mask for nan values
164        mask = np.isnan(thermalFrame)
165        # Replace nan values with a mean for visualization
166        thermalFrame[mask] = np.nanmean(thermalFrame)
167        thermalFrame = cv2.normalize(thermalFrame, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
168        colormappedFrame = cv2.applyColorMap(thermalFrame, cv2.COLORMAP_MAGMA)
169        # Apply the mask back with black pixels (0)
170        colormappedFrame[mask] = 0
171
172        blended = cv2.addWeighted(cvFrameUndistorted, rgbWeight, colormappedFrame, thermalWeight, 0)
173
174        cv2.putText(
175            blended,
176            f"FPS: {fpsCounter.getFps():.2f}",
177            (10, 30),
178            cv2.FONT_HERSHEY_SIMPLEX,
179            1,
180            (255, 255, 255),
181            2,
182        )
183
184        cv2.imshow(windowName, blended)
185
186        key = cv2.waitKey(1)
187        if key == ord("q"):
188            break
189
190        cfg.staticDepthPlane = staticDepthPlane
191        cfgQ.send(cfg)