RGB-Thermal Align

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