oak-examples/tutorials/custom-models/concat.py at main · luxonis/oak-examples · 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
import depthai as dai
from pathlib import Path
from depthai_nodes.node import ParsingNeuralNetwork
from utils.arguments import initialize_argparser

NN_CONCAT_SIZE = (300, 300)

_, args = initialize_argparser()

device = dai.Device(dai.DeviceInfo(args.device)) if args.device else dai.Device()
visualizer = dai.RemoteConnection(httpPort=8082)


with dai.Pipeline(device) as pipeline:
    platform = device.getPlatformAsString()

    cam_rgb = pipeline.create(dai.node.Camera).build(
        boardSocket=dai.CameraBoardSocket.CAM_A
    )
    cam_left = pipeline.create(dai.node.Camera).build(
        boardSocket=dai.CameraBoardSocket.CAM_B
    )
    cam_right = pipeline.create(dai.node.Camera).build(
        boardSocket=dai.CameraBoardSocket.CAM_C
    )

    cam_left_out = cam_left.requestOutput(
        size=NN_CONCAT_SIZE,
        type=dai.ImgFrame.Type.BGR888i
        if platform == "RVC4"
        else dai.ImgFrame.Type.BGR888p,
        fps=args.fps_limit,
    )
    cam_right_out = cam_right.requestOutput(
        size=NN_CONCAT_SIZE,
        type=dai.ImgFrame.Type.BGR888i
        if platform == "RVC4"
        else dai.ImgFrame.Type.BGR888p,
        fps=args.fps_limit,
    )
    cam_rgb_out = cam_rgb.requestOutput(
        size=NN_CONCAT_SIZE,
        type=dai.ImgFrame.Type.BGR888i
        if platform == "RVC4"
        else dai.ImgFrame.Type.BGR888p,
        fps=args.fps_limit,
    )

    # CONCAT
    concat_nn_archive_path = Path(__file__).parent / Path(
        f"models/concat.{platform.lower()}.tar.xz"
    )
    concat_nn_archive = dai.NNArchive(archivePath=str(concat_nn_archive_path))
    nn_concat = pipeline.create(ParsingNeuralNetwork)
    nn_concat.setNNArchive(concat_nn_archive)
    cam_rgb_out.link(nn_concat.inputs["img1"])
    cam_left_out.link(nn_concat.inputs["img2"])
    cam_right_out.link(nn_concat.inputs["img3"])

    visualizer.addTopic("Concat", nn_concat.out, "images")

    print("Pipeline created.")

    pipeline.start()
    visualizer.registerPipeline(pipeline)

    while pipeline.isRunning():
        pipeline.processTasks()
        key = visualizer.waitKey(1)
        if key == ord("q"):
            print("Got q key from the remote connection!")
            break