Enables use of multiple TPUs in OPI5

This commit is contained in:
Tuomas Järvinen
2024-07-27 11:28:47 +03:00
parent 147213cec6
commit 7052a05d55
17 changed files with 336 additions and 65 deletions
@@ -15,6 +15,7 @@ AiEngineInferencevOnnxRuntime::AiEngineInferencevOnnxRuntime(QString modelPath,
{
qDebug() << "TUOMAS AiEngineInferencevOnnxRuntime() mModelPath=" << mModelPath;
/*
mClassNames = {
"Armoured vehicle",
"Truck",
@@ -27,6 +28,90 @@ AiEngineInferencevOnnxRuntime::AiEngineInferencevOnnxRuntime(QString modelPath,
"Additional protection tank",
"Smoke"
};
*/
mClassNames = {
"person",
"bicycle",
"car",
"motorcycle",
"airplane",
"bus",
"train",
"truck",
"boat",
"traffic light",
"fire hydrant",
"stop sign",
"parking meter",
"bench",
"bird",
"cat",
"dog",
"horse",
"sheep",
"cow",
"elephant",
"bear",
"zebra",
"giraffe",
"backpack",
"umbrella",
"handbag",
"tie",
"suitcase",
"frisbee",
"skis",
"snowboard",
"sports ball",
"kite",
"baseball bat",
"baseball glove",
"skateboard",
"surfboard",
"tennis racket",
"bottle",
"wine glass",
"cup",
"fork",
"knife",
"spoon",
"bowl",
"banana",
"apple",
"sandwich",
"orange",
"broccoli",
"carrot",
"hot dog",
"pizza",
"donut",
"cake",
"chair",
"couch",
"potted plant",
"bed",
"dining table",
"toilet",
"tv",
"laptop",
"mouse",
"remote",
"keyboard",
"cell phone",
"microwave",
"oven",
"toaster",
"sink",
"refrigerator",
"book",
"clock",
"vase",
"scissors",
"teddy bear",
"hair drier",
"toothbrush"
};
}
@@ -37,7 +122,7 @@ cv::Mat AiEngineInferencevOnnxRuntime::drawLabels(const cv::Mat &image, const st
for (const auto &detection : detections)
{
cv::rectangle(result, detection.box, cv::Scalar(0, 255, 0), 2);
std::string label = mClassNames[detection.classId] + ": " + std::to_string(detection.conf);
std::string label = mClassNames[detection.classId].toStdString() + ": " + std::to_string(detection.conf);
int baseLine;
cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
@@ -67,31 +152,66 @@ cv::Mat AiEngineInferencevOnnxRuntime::drawLabels(const cv::Mat &image, const st
void AiEngineInferencevOnnxRuntime::performInferenceSlot(cv::Mat frame)
{
mActive = true;
//qDebug() << __PRETTY_FUNCTION__;
cv::Mat scaledImage = resizeAndPad(frame);
std::vector<Yolov8Result> detections = mPredictor.predict(scaledImage);
AiEngineInferenceResult result;
try {
mActive = true;
cv::Mat scaledImage = resizeAndPad(frame);
//cv::imwrite("/tmp/frame.png", scaledImage);
for (uint i = 0; i < detections.size(); i++) {
const Yolov8Result &detection = detections[i];
std::vector<Yolov8Result> detections = mPredictor.predict(scaledImage);
// Add detected objects to the results
AiEngineObject object;
object.classId = detection.classId;
object.propability = detection.conf;
object.rectangle.top = detection.box.y;
object.rectangle.left = detection.box.x;
object.rectangle.bottom = detection.box.y + detection.box.height;
object.rectangle.right = detection.box.x + detection.box.width;
result.objects.append(object);
// Only keep following detected objects.
// car = 2
// train = 6
// cup = 41
// banana = 46
auto it = std::remove_if(detections.begin(), detections.end(),
[](const Yolov8Result& result) {
return result.classId != 2 &&
result.classId != 6 &&
result.classId != 41 &&
result.classId != 46;
});
detections.erase(it, detections.end());
AiEngineInferenceResult result;
for (uint i = 0; i < detections.size(); i++) {
const Yolov8Result &detection = detections[i];
// Add detected objects to the results
AiEngineObject object;
object.classId = detection.classId;
object.classStr = mClassNames[detection.classId];
object.propability = detection.conf;
object.rectangle.top = detection.box.y;
object.rectangle.left = detection.box.x;
object.rectangle.bottom = detection.box.y + detection.box.height;
object.rectangle.right = detection.box.x + detection.box.width;
result.objects.append(object);
}
if (result.objects.empty() == false) {
result.frame = drawLabels(scaledImage, detections);
emit resultsReady(result);
}
mActive = false;
}
if (result.objects.empty() == false) {
qDebug() << __PRETTY_FUNCTION__ << "detections:" << detections.size();
result.frame = drawLabels(scaledImage, detections);
emit resultsReady(result);
catch (const cv::Exception& e) {
std::cout << "OpenCV exception caught: " << e.what() << std::endl;
}
catch (const std::exception& e) {
std::cout << "Standard exception caught: " << e.what() << std::endl;
}
catch (...) {
std::cout << "Unknown exception caught" << std::endl;
}
mActive = false;
}
void AiEngineInferencevOnnxRuntime::initialize(int number)
{
(void)number;
}
@@ -9,6 +9,7 @@ class AiEngineInferencevOnnxRuntime : public AiEngineInference
Q_OBJECT
public:
explicit AiEngineInferencevOnnxRuntime(QString modelPath, QObject *parent = nullptr);
void initialize(int number);
public slots:
void performInferenceSlot(cv::Mat frame) override;
@@ -16,5 +17,5 @@ public slots:
private:
cv::Mat drawLabels(const cv::Mat &image, const std::vector<Yolov8Result> &detections);
YOLOPredictor mPredictor;
std::vector<std::string> mClassNames;
QVector<QString> mClassNames;
};