RGB-Thermal Align

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