Neural Network Multi-input Combined

View as Markdown

1#!/usr/bin/env python3
2import cv2
3import depthai as dai
4import numpy as np
5from pathlib import Path
6
7# Get the absolute path of the current script's directory
8script_dir = Path(__file__).resolve().parent
9examplesRoot = (script_dir / Path('../')).resolve()  # This resolves the parent directory correctly
10models = examplesRoot / 'models'
11tagImage = models / 'lenna.png'
12
13# Decode the image using OpenCV
14lenaImage = cv2.imread(str(tagImage.resolve()))
15lenaImage = cv2.resize(lenaImage, (256, 256))
16lenaImage = cv2.cvtColor(lenaImage, cv2.COLOR_BGR2RGB)
17lenaImage = np.array(lenaImage)
18
19device = dai.Device()
20platform = device.getPlatform()
21if platform == dai.Platform.RVC2:
22    lenaImage = np.transpose(lenaImage, (2, 0, 1))
23    nnTensorType = dai.TensorInfo.DataType.U8F
24elif platform == dai.Platform.RVC4:
25    # Add an empty dimension to the beginning
26    lenaImage = np.expand_dims(lenaImage, axis=0)
27    nnTensorType = dai.TensorInfo.DataType.FP16
28
29inputNNData = dai.NNData()
30inputNNData.addTensor("image1", lenaImage, dataType=nnTensorType)
31inputNNData.addTensor("image2", lenaImage, dataType=nnTensorType)
32
33
34with dai.Pipeline(device) as pipeline:
35    model = dai.NNModelDescription("depthai-test-models/simple-concatenate-model")
36    model.platform = platform.name
37
38    nnArchive = dai.NNArchive(dai.getModelFromZoo(model))
39
40    neuralNetwork = pipeline.create(dai.node.NeuralNetwork)
41    neuralNetwork.setNNArchive(nnArchive)
42    nnDataInputQueue = neuralNetwork.input.createInputQueue()
43    qNNData = neuralNetwork.out.createOutputQueue()
44    pipeline.start()
45    while pipeline.isRunning():
46        nnDataInputQueue.send(inputNNData)
47        inNNData: dai.NNData = qNNData.get()
48        tensor : np.ndarray = inNNData.getFirstTensor()
49        # Drop the first dimension
50        tensor = tensor.squeeze().astype(np.uint8)
51        # Check the shape - in case 3 is not the last dimension, permute it to the last
52        if tensor.shape[0] == 3:
53            tensor = np.transpose(tensor, (1, 2, 0))
54        cv2.imshow("Combined image", tensor)
55        key = cv2.waitKey(1)
56        if key == ord('q'):
57            break

1#include <atomic>
2#include <csignal>
3#include <opencv2/opencv.hpp>
4#include <xtensor/containers/xadapt.hpp>
5#include <xtensor/containers/xarray.hpp>
6
7#include "depthai/depthai.hpp"
8#include "depthai/modelzoo/Zoo.hpp"
9
10std::atomic<bool> quitEvent(false);
11
12void signalHandler(int) {
13    quitEvent = true;
14}
15
16int main() {
17    signal(SIGTERM, signalHandler);
18    signal(SIGINT, signalHandler);
19
20    // Decode the image using OpenCV
21    cv::Mat lenaImageCv = cv::imread(LENNA_PATH);
22    cv::resize(lenaImageCv, lenaImageCv, cv::Size(256, 256));
23    cv::cvtColor(lenaImageCv, lenaImageCv, cv::COLOR_BGR2RGB);
24
25    // Create xt::xarray from cv::Mat
26    std::vector<uint8_t> lenaImageData(lenaImageCv.data, lenaImageCv.data + lenaImageCv.total() * lenaImageCv.channels());
27    xt::xarray<uint8_t> lenaImage = xt::adapt(lenaImageData);
28
29    // Create pipeline
30    dai::Pipeline pipeline;
31
32    // Create model description
33    dai::NNModelDescription model;
34    model.model = "depthai-test-models/simple-concatenate-model";
35    model.platform = pipeline.getDefaultDevice()->getPlatformAsString();
36    dai::NNArchive archive(dai::getModelFromZoo(model));
37
38    // Create and set up nodes
39    auto neuralNetwork = pipeline.create<dai::node::NeuralNetwork>();
40    neuralNetwork->setNNArchive(archive);
41    auto nnDataInputQueue = neuralNetwork->input.createInputQueue();
42    auto qNNData = neuralNetwork->out.createOutputQueue();
43
44    // Prepare input data
45    auto inputNNData = std::make_shared<dai::NNData>();
46    auto platform = pipeline.getDefaultDevice()->getPlatform();
47
48    if(platform == dai::Platform::RVC2) {
49        // Transpose to CHW format
50        lenaImage = xt::transpose(lenaImage, {2, 0, 1});
51        inputNNData->addTensor("image1", lenaImage, dai::TensorInfo::DataType::U8F);
52        inputNNData->addTensor("image2", lenaImage, dai::TensorInfo::DataType::U8F);
53    } else {
54        // Add empty dimension at front (NHWC format)
55        lenaImage = xt::expand_dims(lenaImage, 0);
56        inputNNData->addTensor("image1", lenaImage, dai::TensorInfo::DataType::FP16);
57        inputNNData->addTensor("image2", lenaImage, dai::TensorInfo::DataType::FP16);
58    }
59
60    // Start pipeline
61    pipeline.start();
62
63    // Main loop
64    while(pipeline.isRunning() && !quitEvent) {
65        nnDataInputQueue->send(inputNNData);
66        auto inNNData = qNNData->get<dai::NNData>();
67        auto tensor = inNNData->getFirstTensor<float>();
68        auto tensor_uint8 = xt::eval(xt::squeeze(xt::cast<uint8_t>(tensor), 0));
69
70        cv::Mat output;
71        if(tensor_uint8.shape()[0] == 3) {
72            tensor_uint8 = xt::transpose(tensor_uint8, {1, 2, 0});
73        }
74        output = cv::Mat(tensor_uint8.shape()[0], tensor_uint8.shape()[1], CV_8UC3);
75        std::memcpy(output.data, tensor_uint8.data(), tensor_uint8.size());
76
77        cv::imshow("Combined image", output);
78
79        char key = cv::waitKey(1);
80        if(key == 'q') {
81            break;
82        }
83    }
84
85    pipeline.stop();
86    pipeline.wait();
87
88    return 0;
89}