Skip to content

Commit 0a2781e

Browse files
Vishwa2684Vishwa2684
committed
Write a custom calculator to process yolox tiny's ov tensors
1 parent 155f27e commit 0a2781e

7 files changed

Lines changed: 327 additions & 9 deletions

File tree

‎demos/mediapipe/bytetrack/bytetrack_ovms.pbtxt‎

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -47,21 +47,34 @@ node {
4747
output_stream: "OVTENSORS2:detection_tensors"
4848
node_options: {
4949
[type.googleapis.com/mediapipe.OpenVINOInferenceCalculatorOptions]: {
50-
input_order_list :["images"]
51-
output_order_list :["output"]
50+
input_order_list :["image"]
51+
output_order_list :["boxes","labels"]
5252
}
5353
}
5454
}
5555

56-
### WRITE YOLO SPECIFIC CALCULATORS
56+
### OpenVINOYoloXTensorsToDetectionsCalculator was developed for TFLite specific model, but we can use OVYoloXTensorsToDetectionsCalculator for OpenVINO model.
5757

58-
node{
59-
calculator: "OpenVINOYoloXTensorsToDetectionsCalculator"
60-
input_stream: "TENSORS:detection_tensors"
61-
output_stream: "DETECTIONS:detections"
62-
node_options: {
63-
[type.googleapis.com/mediapipe.OpenVINOYoloXTensorsToDetectionsCalculatorOptions] {
58+
#node{
59+
# calculator: "OpenVINOYoloXTensorsToDetectionsCalculator"
60+
# input_stream: "TENSORS:detection_tensors"
61+
# output_stream: "DETECTIONS:detections"
62+
# node_options: {
63+
# [type.googleapis.com/mediapipe.OpenVINOYoloXTensorsToDetectionsCalculatorOptions] {
64+
# conf_thresh: 0.1
65+
# }
66+
# }
67+
# }
68+
69+
node {
70+
calculator: "OVYoloXTensorsToDetectionsCalculator"
71+
input_stream: "TENSORS:detection_tensors"
72+
output_stream: "DETECTIONS:detections"
73+
74+
node_options: {
75+
[type.googleapis.com/mediapipe.OVYoloXTensorsToDetectionsCalculatorOptions] {
6476
conf_thresh: 0.1
77+
input_size: 416.0
6578
}
6679
}
6780
}

‎demos/mediapipe/bytetrack/download_models.py‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
import glob
2+
import json
3+
import os
4+
import requests
5+
import openvino as ov
6+
from huggingface_hub import snapshot_download
7+
8+
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
9+
10+
COCO_CLASSES_URL = "https://raw.githubusercontent.com/openvinotoolkit/open_model_zoo/master/data/dataset_classes/coco_80cl.txt"
11+
12+
MODEL_REPO = "OpenVINO/yolox_tiny-fp16-ov"
13+
# ---------------------------------------------------------
14+
# 1. Download model
15+
# ---------------------------------------------------------
16+
17+
model_dir = snapshot_download(repo_id=MODEL_REPO)
18+
19+
xml_path = glob.glob(os.path.join(model_dir, "*.xml"))[0]
20+
config_path = os.path.join(model_dir, "config.json")
21+
22+
print("Found IR :", xml_path)
23+
print("Found config:", config_path)
24+
25+
26+
# ---------------------------------------------------------
27+
# 2. Read config.json
28+
# ---------------------------------------------------------
29+
30+
with open(config_path, "r") as f:
31+
config = json.load(f)
32+
33+
print("\nModel config:")
34+
print("model_name :", config.get("model_name"))
35+
print("model_type :", config.get("model_type"))
36+
print("input_type :", config.get("input_dtype"))
37+
print("mean_values:", config.get("mean_values"))
38+
print("scale_values:", config.get("scale_values"))
39+
40+
41+
# ---------------------------------------------------------
42+
# 3. Parse mean and scale values
43+
# ---------------------------------------------------------
44+
mean_values = [float(x) for x in config["mean_values"].split()]
45+
scale_values = [float(x) for x in config["scale_values"].split()]
46+
47+
print("\nParsed preprocessing:")
48+
print("mean :", mean_values)
49+
print("scale:", scale_values)
50+
51+
# ---------------------------------------------------------
52+
# 4. Load OpenVINO model
53+
# ---------------------------------------------------------
54+
core = ov.Core()
55+
model = core.read_model(xml_path)
56+
# ---------------------------------------------------------
57+
# 5. Configure preprocessing
58+
# ---------------------------------------------------------
59+
ppp = ov.preprocess.PrePostProcessor(model)
60+
inp = ppp.input(0)
61+
62+
# Input coming from user/image:
63+
# f32 NHWC
64+
inp.tensor().set_element_type(ov.Type.f16) .set_layout(ov.Layout("NHWC"))
65+
66+
# Model expects:
67+
# float32 NCHW
68+
inp.model().set_layout(ov.Layout("NCHW"))
69+
70+
# Preprocessing:
71+
72+
inp.preprocess().convert_element_type(ov.Type.f16).convert_layout(ov.Layout("NCHW")).scale(255.0).mean(mean_values).scale(scale_values)
73+
# ---------------------------------------------------------
74+
# 6. Build and save
75+
# ---------------------------------------------------------
76+
77+
model = ppp.build()
78+
79+
output_path = "yolox_tiny_float32/1/yolox_tiny_full_preprocess.xml"
80+
81+
ov.save_model(model, output_path)
82+
83+
print("\nSaved:", os.path.abspath(output_path))
84+
85+
with open("coco_80cl.txt", "wb") as f:
86+
f.write(requests.get(COCO_CLASSES_URL).content)
87+
88+
print("Downloaded successfully")

‎demos/mediapipe/bytetrack/requirements.txt‎

316 Bytes
Binary file not shown.

‎src/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -757,6 +757,7 @@ ovms_cc_library(
757757
"//src/image_gen:image_gen_calculator",
758758
"//src/audio/speech_to_text:s2t_calculator",
759759
"//src/audio/text_to_speech:t2s_calculator",
760+
"//src/yolox:ov_yolox_tensors_to_detections_calculator",
760761
"//src/audio:audio_utils",
761762
"//src/image_gen:imagegen_init",
762763
"//src/llm:openai_responses_handler",

‎src/yolox/BUILD‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
load("@mediapipe//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library", "mediapipe_proto_library")
2+
3+
mediapipe_proto_library(
4+
name = "ov_yolox_tensors_to_detections_calculator_proto",
5+
srcs = [
6+
"ov_yolox_tensors_to_detections_calculator.proto",
7+
],
8+
deps = [
9+
"@mediapipe//mediapipe/framework:calculator_options_proto",
10+
"@mediapipe//mediapipe/framework:calculator_proto"
11+
],
12+
visibility = [
13+
"//src:__pkg__",
14+
],
15+
)
16+
17+
18+
cc_library(
19+
name = "ov_yolox_tensors_to_detections_calculator",
20+
srcs = [
21+
"ov_yolox_tensors_to_detections_calculator.cc",
22+
],
23+
deps = [
24+
"//third_party:openvino",
25+
"@mediapipe//mediapipe/framework:calculator_framework",
26+
"@mediapipe//mediapipe/framework/formats:detection_cc_proto",
27+
"@mediapipe//mediapipe/framework/formats:location_data_cc_proto",
28+
"@mediapipe//mediapipe/framework/port:ret_check",
29+
"@mediapipe//mediapipe/framework/port:status",
30+
"@mediapipe//mediapipe/framework:calculator_cc_proto",
31+
":ov_yolox_tensors_to_detections_calculator_cc_proto",
32+
],
33+
visibility = [
34+
"//src:__pkg__",
35+
],
36+
alwayslink = True
37+
)

‎src/yolox/ov_yolox_tensors_to_detections_calculator.cc‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
//*****************************************************************************
2+
// Copyright (c) 2026 Intel Corporation
3+
//
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
//
8+
// http://www.apache.org/licenses/LICENSE-2.0
9+
//
10+
// Unless required by applicable law or agreed to in writing, software
11+
// distributed under the License is distributed on an "AS IS" BASIS,
12+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// See the License for the specific language governing permissions and
14+
// limitations under the License.
15+
//*****************************************************************************
16+
17+
#include <cstdint>
18+
#include <vector>
19+
20+
#include <openvino/openvino.hpp>
21+
22+
#include "src/yolox/ov_yolox_tensors_to_detections_calculator.pb.h"
23+
#include "mediapipe/framework/calculator_framework.h"
24+
#include "mediapipe/framework/formats/detection.pb.h"
25+
#include "mediapipe/framework/formats/location_data.pb.h"
26+
#include "mediapipe/framework/port/ret_check.h"
27+
#include "mediapipe/framework/port/status.h"
28+
29+
namespace mediapipe {
30+
31+
class OVYoloXTensorsToDetectionsCalculator : public CalculatorBase {
32+
static const std::string INPUT_TAG_NAME;
33+
static const std::string OUTPUT_TAG_NAME;
34+
35+
float confidenceThreshold_ = 0.1f;
36+
float inputSize_ = 416.0f;
37+
38+
public:
39+
static absl::Status GetContract(CalculatorContract* cc) {
40+
RET_CHECK(!cc->Inputs().GetTags().empty());
41+
RET_CHECK(!cc->Outputs().GetTags().empty());
42+
43+
cc->Inputs()
44+
.Tag(INPUT_TAG_NAME)
45+
.Set<std::vector<ov::Tensor>>();
46+
47+
cc->Outputs()
48+
.Tag(OUTPUT_TAG_NAME)
49+
.Set<std::vector<Detection>>();
50+
51+
return absl::OkStatus();
52+
}
53+
54+
absl::Status Open(CalculatorContext* cc) final {
55+
const auto& options =
56+
cc->Options<
57+
mediapipe::OVYoloXTensorsToDetectionsCalculatorOptions>();
58+
59+
confidenceThreshold_ =
60+
options.has_conf_thresh() ? options.conf_thresh() : 0.1f;
61+
62+
inputSize_ =
63+
options.has_input_size() ? options.input_size() : 416.0f;
64+
65+
return absl::OkStatus();
66+
}
67+
68+
absl::Status Process(CalculatorContext* cc) final {
69+
if (cc->Inputs().Tag(INPUT_TAG_NAME).IsEmpty()) {
70+
return absl::OkStatus();
71+
}
72+
73+
const auto& tensors =
74+
cc->Inputs()
75+
.Tag(INPUT_TAG_NAME)
76+
.Get<std::vector<ov::Tensor>>();
77+
78+
RET_CHECK_EQ(tensors.size(), 2u);
79+
80+
const auto& boxesTensor = tensors[0];
81+
const auto& labelsTensor = tensors[1];
82+
83+
RET_CHECK_EQ(boxesTensor.get_element_type(), ov::element::f32);
84+
RET_CHECK_EQ(labelsTensor.get_element_type(), ov::element::i64);
85+
86+
const auto boxesShape = boxesTensor.get_shape();
87+
const auto labelsShape = labelsTensor.get_shape();
88+
89+
RET_CHECK_EQ(boxesShape.size(), 3u);
90+
RET_CHECK_EQ(boxesShape[0], 1u);
91+
RET_CHECK_EQ(boxesShape[2], 5u);
92+
93+
RET_CHECK_EQ(labelsShape.size(), 2u);
94+
RET_CHECK_EQ(labelsShape[0], 1u);
95+
RET_CHECK_EQ(labelsShape[1], boxesShape[1]);
96+
97+
const size_t numBoxes = boxesShape[1];
98+
99+
const float* boxes = boxesTensor.data<const float>();
100+
const int64_t* labels = labelsTensor.data<const int64_t>();
101+
102+
RET_CHECK(boxes != nullptr);
103+
RET_CHECK(labels != nullptr);
104+
105+
auto detections =
106+
absl::make_unique<std::vector<Detection>>();
107+
108+
for (size_t i = 0; i < numBoxes; ++i) {
109+
const float x1 = boxes[i * 5 + 0];
110+
const float y1 = boxes[i * 5 + 1];
111+
const float x2 = boxes[i * 5 + 2];
112+
const float y2 = boxes[i * 5 + 3];
113+
const float confidence = boxes[i * 5 + 4];
114+
115+
if (confidence < confidenceThreshold_) {
116+
continue;
117+
}
118+
119+
if (x2 <= x1 || y2 <= y1) {
120+
continue;
121+
}
122+
123+
Detection detection;
124+
125+
auto* locationData =
126+
detection.mutable_location_data();
127+
128+
locationData->set_format(
129+
LocationData::RELATIVE_BOUNDING_BOX);
130+
131+
auto* boundingBox =
132+
locationData->mutable_relative_bounding_box();
133+
134+
boundingBox->set_xmin(x1 / inputSize_);
135+
boundingBox->set_ymin(y1 / inputSize_);
136+
boundingBox->set_width((x2 - x1) / inputSize_);
137+
boundingBox->set_height((y2 - y1) / inputSize_);
138+
139+
detection.add_score(confidence);
140+
detection.add_label_id(
141+
static_cast<int>(labels[i]));
142+
143+
detections->emplace_back(std::move(detection));
144+
}
145+
146+
cc->Outputs()
147+
.Tag(OUTPUT_TAG_NAME)
148+
.Add(detections.release(), cc->InputTimestamp());
149+
150+
return absl::OkStatus();
151+
}
152+
};
153+
154+
const std::string OVYoloXTensorsToDetectionsCalculator::INPUT_TAG_NAME{
155+
"TENSORS"};
156+
157+
const std::string OVYoloXTensorsToDetectionsCalculator::OUTPUT_TAG_NAME{
158+
"DETECTIONS"};
159+
160+
REGISTER_CALCULATOR(OVYoloXTensorsToDetectionsCalculator);
161+
162+
} // namespace mediapipe

‎src/yolox/ov_yolox_tensors_to_detections_calculator.proto‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
syntax = "proto2";
2+
3+
package mediapipe;
4+
5+
import "mediapipe/framework/calculator.proto";
6+
7+
message OVYoloXTensorsToDetectionsCalculatorOptions {
8+
extend .mediapipe.CalculatorOptions {
9+
optional OVYoloXTensorsToDetectionsCalculatorOptions ext = 211376658;
10+
}
11+
12+
// Minimum confidence required for a detection.
13+
optional float conf_thresh = 1 [default = 0.10];
14+
15+
// Input resolution of the YOLOX model.
16+
optional float input_size = 3 [default = 416.0];
17+
}

0 commit comments

Comments
 (0)