Casting NN subtraction

1#!/usr/bin/env python3
2
3import cv2
4import depthai as dai
5from pathlib import Path
6
7SHAPE = 720
8
9p = dai.Pipeline()
10
11camRgb = p.create(dai.node.ColorCamera)
12nn = p.create(dai.node.NeuralNetwork)
13script = p.create(dai.node.Script)
14rgbXout = p.create(dai.node.XLinkOut)
15cast = p.create(dai.node.Cast)
16castXout = p.create(dai.node.XLinkOut)
17
18camRgb.setVideoSize(SHAPE, SHAPE)
19camRgb.setPreviewSize(SHAPE, SHAPE)
20camRgb.setInterleaved(False)
21
22nnBlobPath = (Path(__file__).parent / Path('../models/diff_openvino_2022.1_6shave.blob')).resolve().absolute()
23nn.setBlobPath(nnBlobPath)
24
25script.setScript("""
26old = node.io['in'].get()
27while True:
28    frame = node.io['in'].get()
29    node.io['img1'].send(old)
30    node.io['img2'].send(frame)
31    old = frame
32""")
33
34rgbXout.setStreamName("rgb")
35castXout.setStreamName("cast")
36cast.setOutputFrameType(dai.RawImgFrame.Type.GRAY8)
37
38# Linking
39camRgb.preview.link(script.inputs['in'])
40script.outputs['img1'].link(nn.inputs['img1'])
41script.outputs['img2'].link(nn.inputs['img2'])
42camRgb.video.link(rgbXout.input)
43nn.out.link(cast.input)
44cast.output.link(castXout.input)
45
46# Pipeline is defined, now we can connect to the device
47with dai.Device(p) as device:
48    qCam = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
49    qCast = device.getOutputQueue(name="cast", maxSize=4, blocking=False)
50
51
52    while True:
53        colorFrame = qCam.get()
54        assert isinstance(colorFrame, dai.ImgFrame)
55        cv2.imshow("Color", colorFrame.getCvFrame())
56
57        inCast = qCast.get()
58        assert isinstance(inCast, dai.ImgFrame)
59        cv2.imshow("Diff", inCast.getCvFrame())
60
61        if cv2.waitKey(1) == ord('q'):
62            break