diff --git a/ci/lib_search.py b/ci/lib_search.py index 29bb3c38e9..86901e416a 100644 --- a/ci/lib_search.py +++ b/ci/lib_search.py @@ -113,6 +113,7 @@ def check_dir(start_dir): 'ovms-c/dist', 'requirements.txt', 'requirements_win.txt', + 'bytetrack_ovms.pbtxt', 'resnet_images.txt', "resnet_labels.txt", 'rest_sdk_v2.10.16.patch', diff --git a/demos/mediapipe/bytetrack/README.md b/demos/mediapipe/bytetrack/README.md new file mode 100644 index 0000000000..bcf038874b --- /dev/null +++ b/demos/mediapipe/bytetrack/README.md @@ -0,0 +1,70 @@ +# ByteTrack Demo Setup + +End-to-end demo: video source (webcam / file) → OpenVINO Model Server (YOLOX Tiny + ByteTrack) → output (screen / file). + +--- + +## Steps + +### 1. Clone the repository + +Clone the repository, switch to the `gsoc_bytetrack` branch, and move into the demo directory: + +```bash +git clone https://github.com/Vishwa2684/model_server +cd model_server +git checkout gsoc_bytetrack +cd demos/mediapipe/bytetrack +``` + +### 2. Install requirements + +Install all the Python dependencies needed by the client and the model download script. Run this from inside the `demos/mediapipe/bytetrack` directory: + +```bash +pip install -r requirements.txt +``` + +### 3. Download a model + +Download the detector model that the OpenVINO Model Server will use. This same command also fetches the COCO class list used for labeling detections: + +```bash +python download_models.py --model OpenVINO/yolox_tiny-fp16-ov +``` + +> Swap `--model-repo` for any of the repo IDs listed below to use a different YOLOX size. + +| Model | HuggingFace Repo | +|---|---| +| YOLOX-Tiny (fp16 precision)| `OpenVINO/yolox_tiny-fp16-ov` | +| YOLOX-Tiny (int8 precision)| `OpenVINO/yolox_tiny-int8-ov` | + +`yolox_tiny-fp16-ov` is the default used in this demo. + +This step populates the local model directory that `config.json` (used by the OpenVINO Model Server in step 4) points to, and that ByteTrack consumes downstream for tracking. + +### 4. Start the OpenVINO Model Server + +Bring up the OpenVINO Model Server as a Docker container. This mounts your current directory into the container so it can read `config.json`, and exposes port 9000 for the client to connect to: + +```bash +docker run -d -v $PWD:/demo -p 9000:9000 openvino/model_server:latest --config_path /demo/config.json --port 9000 +``` + +Leave this container running in the background — the client in the next step connects to it over gRPC. + +### 5. Run the demo — local webcam → screen + +With the model server running, run the client script. This reads directly from your local webcam, runs it through detection + ByteTrack tracking, and renders the annotated output live in a window on your screen: + +```bash +cd ../../real_time_stream_analysis/python +python client.py --grpc_address localhost:9000 --input_stream 0 --output_stream screen --model_name ByteTrack --input_name input_video +``` + +- `--grpc_address localhost:9000` — address of the OpenVINO Model Server started in step 4. +- `--input_stream 0` — camera device ID `0` (use `1`, `2`, etc. if you have multiple cameras and want a different one). +- `--output_stream screen` — opens a live preview window instead of writing to a file or stream. + +A window should open showing your webcam feed with tracked bounding boxes drawn on it in real time. To use different input and output streams for real time. Read the documentation on [real time stream analysis](../../real_time_stream_analysis/python/README.md) \ No newline at end of file diff --git a/demos/mediapipe/bytetrack/bytetrack_ovms.pbtxt b/demos/mediapipe/bytetrack/bytetrack_ovms.pbtxt new file mode 100644 index 0000000000..9969e38350 --- /dev/null +++ b/demos/mediapipe/bytetrack/bytetrack_ovms.pbtxt @@ -0,0 +1,138 @@ +input_stream: "IMAGE:input_video" +output_stream: "IMAGE:output" + +node: { + calculator: "ImageTransformationCalculator" + input_stream: "IMAGE:input_video" + output_stream: "IMAGE:transformed_input_video" + node_options: { + [type.googleapis.com/mediapipe.ImageTransformationCalculatorOptions] { + output_width: 416 + output_height: 416 + } + } +} + +node { + calculator: "OpenVINOConverterCalculator" + input_stream: "IMAGE:transformed_input_video" + output_stream: "TENSORS:image_tensor" + node_options: { + [type.googleapis.com/mediapipe.OpenVINOConverterCalculatorOptions] { + enable_normalization: true + use_custom_normalization: true + custom_div: 1.0 + custom_sub: 0.0 + } + } +} + +node { + calculator: "OpenVINOModelServerSessionCalculator" + output_side_packet: "SESSION:session" + node_options: { + [type.googleapis.com/mediapipe.OpenVINOModelServerSessionCalculatorOptions]: { + servable_name: "yolox_tiny-fp16-ov" # servable name inside OVMS + servable_version: "1" + } + } +} + +node { + calculator: "OpenVINOInferenceCalculator" + input_side_packet: "SESSION:session" + input_stream: "OVTENSORS:image_tensor" + output_stream: "OVTENSORS2:detection_tensors" + node_options: { + [type.googleapis.com/mediapipe.OpenVINOInferenceCalculatorOptions]: { + input_order_list :["image"] + output_order_list :["boxes","labels"] + } + } +} + +### OpenVINOYoloXTensorsToDetectionsCalculator was developed for TFLite specific model, but we can use OVYoloXTensorsToDetectionsCalculator for OpenVINO model. + +#node{ +# calculator: "OpenVINOYoloXTensorsToDetectionsCalculator" +# input_stream: "TENSORS:detection_tensors" +# output_stream: "DETECTIONS:detections" +# node_options: { +# [type.googleapis.com/mediapipe.OpenVINOYoloXTensorsToDetectionsCalculatorOptions] { +# conf_thresh: 0.1 +# } +# } +# } + +node { + calculator: "OVYoloXTensorsToDetectionsCalculator" + input_stream: "TENSORS:detection_tensors" + output_stream: "DETECTIONS:detections" + + node_options: { + [type.googleapis.com/mediapipe.OVYoloXTensorsToDetectionsCalculatorOptions] { + conf_thresh: 0.1 + input_size: 416.0 + } + } +} + +# Performs non-max suppression to remove excessive detections. +node { + calculator: "NonMaxSuppressionCalculator" + input_stream: "detections" + output_stream: "filtered_detections" + node_options: { + [type.googleapis.com/mediapipe.NonMaxSuppressionCalculatorOptions] { + min_suppression_threshold: 0.45 + max_num_detections: 100 + overlap_type: INTERSECTION_OVER_UNION + return_empty_detections: true + } + } +} + + +# Maps detection label IDs to the corresponding label text. The label map is +# provided in the label_map_path option. +node { + calculator: "DetectionLabelIdToTextCalculator" + input_stream: "filtered_detections" + output_stream: "output_detections" + node_options: { + [type.googleapis.com/mediapipe.DetectionLabelIdToTextCalculatorOptions] { + label_map_path: "/demo/coco_80cl.txt" + } + } +} + +node { + calculator: "ByteTrackCalculator" + input_stream: "DETECTIONS:output_detections" + output_stream: "DETECTIONS:tracked_detections" + options: { + [mediapipe.ByteTrackCalculatorOptions.ext] { + track_high_threshold:0.7 + track_low_threshold:0.55 + new_track_threshold:0.35 + matching_threshold: 0.8 + track_buffer: 60 + fuse_score: false + } + } +} + +# Converts the detections to drawing primitives for annotation overlay. +node { + calculator: "DetectionColorByIdCalculator" + input_stream: "DETECTIONS:tracked_detections" + output_stream: "RENDER_DATA:detections_render_data" +} + +# Draws annotations and overlays them on top of the input images. +node { + calculator: "AnnotationOverlayCalculator" + input_stream: "IMAGE:input_video" + input_stream: "detections_render_data" + output_stream: "IMAGE:output" +} \ No newline at end of file diff --git a/demos/mediapipe/bytetrack/config.json b/demos/mediapipe/bytetrack/config.json new file mode 100644 index 0000000000..f907d2d92b --- /dev/null +++ b/demos/mediapipe/bytetrack/config.json @@ -0,0 +1,16 @@ +{ + "model_config_list": [ + {"config": { + "name": "yolox_tiny-fp16-ov", + "base_path": "yolox_tiny-fp16-ov" + } + } + ], + "mediapipe_config_list": [ + { + "name":"ByteTrack", + "base_path":"./", + "graph_path":"bytetrack_ovms.pbtxt" + } + ] +} \ No newline at end of file diff --git a/demos/mediapipe/bytetrack/download_models.py b/demos/mediapipe/bytetrack/download_models.py new file mode 100644 index 0000000000..4e68753818 --- /dev/null +++ b/demos/mediapipe/bytetrack/download_models.py @@ -0,0 +1,120 @@ +# +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import glob +import json +import argparse +import os +import openvino as ov +from huggingface_hub import snapshot_download + +os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1" + +parser = argparse.ArgumentParser() +parser.add_argument( + "--model-repo", + default="OpenVINO/yolox_tiny-fp16-ov", + help="Hugging Face model repository", +) +args = parser.parse_args() + +MODEL_REPO = args.model_repo +MODEL_NAME = MODEL_REPO.split("/")[-1] +# --------------------------------------------------------- +# 1. Download model +# --------------------------------------------------------- + +model_dir = snapshot_download(repo_id=MODEL_REPO) + +xml_path = glob.glob(os.path.join(model_dir, "*.xml"))[0] +config_path = os.path.join(model_dir, "config.json") + +print("Found IR :", xml_path) +print("Found config:", config_path) + + +# --------------------------------------------------------- +# 2. Read config.json +# --------------------------------------------------------- + +with open(config_path, "r") as f: + config = json.load(f) + +print("\nModel config:") +print("model_name :", config.get("model_name")) +print("model_type :", config.get("model_type")) +print("input_type :", config.get("input_dtype")) +print("mean_values:", config.get("mean_values")) +print("scale_values:", config.get("scale_values")) +print("classes:", config.get("labels")) + +# --------------------------------------------------------- +# 3. Prepare classes list +# --------------------------------------------------------- +classes = config.get("labels").split(" ") + +# --------------------------------------------------------- +# 4. Parse mean and scale values +# --------------------------------------------------------- +mean_values = [float(x) for x in config["mean_values"].split()] +scale_values = [float(x) for x in config["scale_values"].split()] + +print("\nParsed preprocessing:") +print("mean :", mean_values) +print("scale:", scale_values) + +# --------------------------------------------------------- +# 5. Load OpenVINO model +# --------------------------------------------------------- +core = ov.Core() +model = core.read_model(xml_path) +# --------------------------------------------------------- +# 6. Configure preprocessing +# --------------------------------------------------------- +ppp = ov.preprocess.PrePostProcessor(model) +inp = ppp.input(0) + +# Input coming from user/image: +# f32 NHWC +inp.tensor().set_element_type(ov.Type.f32).set_layout(ov.Layout("NHWC")) + +# Model expects: +# float32 NCHW +inp.model().set_layout(ov.Layout("NCHW")) + +# Preprocessing: + +inp.preprocess().convert_element_type(ov.Type.f32).convert_layout( + ov.Layout("NCHW") +).scale(255.0).mean(mean_values).scale(scale_values) +# --------------------------------------------------------- +# 7. Build and save +# --------------------------------------------------------- + +model = ppp.build() + +output_path = f"{MODEL_NAME}/1/{MODEL_NAME}.xml" +os.makedirs(os.path.dirname(output_path), exist_ok=True) +ov.save_model(model, output_path) + +print("\nSaved:", os.path.abspath(output_path)) + +with open("coco_80cl.txt", "w") as f: + n = len(classes) + for i, c in enumerate(classes): + f.write(c + ("\n" if i < n - 1 else "")) + +print("Downloaded successfully") diff --git a/demos/mediapipe/bytetrack/requirements.txt b/demos/mediapipe/bytetrack/requirements.txt new file mode 100644 index 0000000000..66e7fd02ec Binary files /dev/null and b/demos/mediapipe/bytetrack/requirements.txt differ diff --git a/demos/real_time_stream_analysis/python/client.py b/demos/real_time_stream_analysis/python/client.py index 00e76e6561..7da26ab66a 100755 --- a/demos/real_time_stream_analysis/python/client.py +++ b/demos/real_time_stream_analysis/python/client.py @@ -23,6 +23,8 @@ parser = argparse.ArgumentParser() parser.add_argument('--grpc_address', required=False, default='localhost:9022', help='Specify url to grpc service') +parser.add_argument('--ffmpeg_output_width', required=False, default=None, type=int, help='Width of the output video') +parser.add_argument('--ffmpeg_output_height', required=False, default=None, type=int, help='Height of the output video') parser.add_argument('--input_stream', required=False, default="rtsp://localhost:8080/channel1", type=str, help='Url of input rtsp stream') parser.add_argument('--output_stream', required=False, default="rtsp://localhost:8080/channel2", type=str, help='Url of output rtsp stream') parser.add_argument('--model_name', required=False, default="holisticTracking", type=str, help='Name of the model') @@ -54,6 +56,6 @@ def postprocess(frame, result): backend = StreamClient.OutputBackends.cv2 exact = True -client = StreamClient(postprocess_callback = postprocess, preprocess_callback=preprocess, output_backend=backend, source=args.input_stream, sink=args.output_stream, exact=exact, benchmark=args.benchmark, verbose=args.verbose) +client = StreamClient(postprocess_callback = postprocess, preprocess_callback=preprocess, output_backend=backend, source=args.input_stream, sink=args.output_stream, exact=exact, benchmark=args.benchmark, verbose=args.verbose, ffmpeg_output_width=args.ffmpeg_output_width, ffmpeg_output_height=args.ffmpeg_output_height) client.start(ovms_address=args.grpc_address, input_name=args.input_name, model_name=args.model_name, datatype = StreamClient.Datatypes.uint8, batch = False, limit_stream_duration = args.limit_stream_duration, limit_frames = args.limit_frames, streaming_api=True) diff --git a/spelling-whitelist.txt b/spelling-whitelist.txt index 8d2ea40cfe..cbf12d136c 100644 --- a/spelling-whitelist.txt +++ b/spelling-whitelist.txt @@ -39,6 +39,10 @@ windows_parse_tests.bat:136: SEH ==> SHE windows_parse_tests.bat:141: SEH ==> SHE windows_parse_tests.bat:144: SEH ==> SHE src/test/llm/output_parsers/gemma4_output_parser_test.cpp +src/test/llm/output_parsers/qwen3_output_parser_test.cpp:719: thi ==> the, this +STrack +strack +nd src/test/llm/output_parsers/qwen3_output_parser_test.cpp:697: thi ==> the, this extras/chat_template_examples/chat_template_onyx.jinja src/test/llm/chat_templates/chat_template_onyx.jinja diff --git a/src/BUILD b/src/BUILD index 1e2368327b..6f8e32e4c4 100644 --- a/src/BUILD +++ b/src/BUILD @@ -762,6 +762,8 @@ ovms_cc_library( "//src/image_gen:image_gen_calculator", "//src/audio/speech_to_text:s2t_calculator", "//src/audio/text_to_speech:t2s_calculator", + "//src/yolox:ov_yolox_tensors_to_detections_calculator", + "//src/bytetrack/calculators:bytetrack_calculators", "//src/audio:audio_utils", "//src/image_gen:imagegen_init", "//src/llm:openai_responses_handler", diff --git a/src/bytetrack/calculators/BUILD b/src/bytetrack/calculators/BUILD new file mode 100644 index 0000000000..daec9152b8 --- /dev/null +++ b/src/bytetrack/calculators/BUILD @@ -0,0 +1,127 @@ +# +# Copyright (c) 2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +load("@mediapipe//mediapipe/framework/port:build_config.bzl", "mediapipe_proto_library") + +licenses(["notice"]) + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "bytetrack_calculators", + deps = [ + "//src/bytetrack/utils:detection_color_by_id_calculator", + ":kalman_filter", + ":kalman_matrices", + ":strack", + ":bytetrack_calculator", + ":render_data_passthrough", + ], + alwayslink = 1 +) + +mediapipe_proto_library( + name = "bytetrack_calculator_proto", + visibility = ["//visibility:public"], + srcs = ["bytetrack_calculator.proto"], + deps=[ + "@mediapipe//mediapipe/framework:calculator_options_proto", + "@mediapipe//mediapipe/framework:calculator_proto", + ], + alwayslink = 1 +) + +cc_library( + name = "kalman_matrices", + hdrs = ["kalman_matrices.h"], + deps = ["@eigen_archive//:eigen3"], + alwayslink = 1 +) + + +cc_library( + name = "matching_utils", + hdrs = ["matching_utils.h"], + deps = [ + ":strack", + "@eigen_archive//:eigen3", + ], + alwayslink = 1 +) + +cc_library( + name = "kalman_filter", + srcs = ["kalman_filter.cc"], + hdrs = ["kalman_filter.h"], + deps = [ + ":kalman_matrices", + "@mediapipe//mediapipe/framework/formats:detection_cc_proto", + "@eigen_archive//:eigen3", + ], + alwayslink = 1 +) + +cc_library( + name = "basetrack", + srcs = ["basetrack.cc"], + hdrs = ["basetrack.h"], + alwayslink = 1 +) + +cc_library( + name = "strack", + srcs = ["strack.cc"], + hdrs = ["strack.h"], + deps = [ + ":kalman_filter", + ":basetrack", + "@mediapipe//mediapipe/framework/formats:detection_cc_proto", + "@eigen_archive//:eigen3", + ], + alwayslink = 1 +) + +cc_library( + name = "bytetrack_calculator", + srcs = ["bytetrack_calculator.cc"], + hdrs = ["bytetrack_calculator.h"], + deps = [ + "@mediapipe//mediapipe/framework:calculator_framework", + "@mediapipe//mediapipe/framework/formats:detection_cc_proto", + ":bytetrack_calculator_cc_proto", + "@mediapipe//mediapipe/framework/port:status", + ":matching_utils", + ":strack", + ":kalman_filter", + ":kalman_matrices", + "@eigen_archive//:eigen3", + ], + alwayslink = 1, +) + +# This calculator is for debugging to check whether render data is correctly passed through the graph. +cc_library( + name = "render_data_passthrough", + srcs = ["render_data_passthrough_calculator.cc"], + visibility = ["//visibility:public"], + deps = [ + "@mediapipe//mediapipe/framework:calculator_framework", + "@mediapipe//mediapipe/framework/formats:detection_cc_proto", + "@mediapipe//mediapipe/framework/port:status", + "@mediapipe//mediapipe/util:render_data_cc_proto", + ], + alwayslink = 1 +) diff --git a/src/bytetrack/calculators/basetrack.cc b/src/bytetrack/calculators/basetrack.cc new file mode 100644 index 0000000000..11aa702fba --- /dev/null +++ b/src/bytetrack/calculators/basetrack.cc @@ -0,0 +1,27 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** +// basetrack.cpp +#include "src/bytetrack/calculators/basetrack.h" + +#include + +namespace mediapipe { +namespace bytetrack { + +std::atomic BaseTrack::count_{0}; + +} // namespace bytetrack +} // namespace mediapipe diff --git a/src/bytetrack/calculators/basetrack.h b/src/bytetrack/calculators/basetrack.h new file mode 100644 index 0000000000..d1f42dc012 --- /dev/null +++ b/src/bytetrack/calculators/basetrack.h @@ -0,0 +1,60 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +#ifndef SRC_BYTETRACK_CALCULATORS_BASETRACK_H_ +#define SRC_BYTETRACK_CALCULATORS_BASETRACK_H_ + +#include + +namespace mediapipe { +namespace bytetrack { + +class BaseTrack { +public: + enum class TrackState { NEW, + TRACKED, + LOST, + REMOVED }; + + // mirrors Python's next_id() — static counter owned here + static int next_id() { return count_.fetch_add(1, std::memory_order_relaxed) + 1; } + static void reset_id() { count_.store(0, std::memory_order_relaxed); } + + // Accessors + int track_id() const { return track_id_; } + int frame_id() const { return frame_id_; } + int start_frame() const { return start_frame_; } + float score() const { return score_; } + TrackState state() const { return state_; } + bool is_activated() const { return is_activated_; } + + void MarkLost() { state_ = TrackState::LOST; } + void MarkRemoved() { state_ = TrackState::REMOVED; } + +protected: + int track_id_ = 0; + int frame_id_ = 0; + int start_frame_ = 0; + float score_ = 0.f; + bool is_activated_ = false; + TrackState state_ = TrackState::NEW; + static std::atomic count_; +}; + +} // namespace bytetrack +} // namespace mediapipe + +#endif // SRC_BYTETRACK_CALCULATORS_BASETRACK_H_ diff --git a/src/bytetrack/calculators/bytetrack_calculator.cc b/src/bytetrack/calculators/bytetrack_calculator.cc new file mode 100644 index 0000000000..9feb32eb13 --- /dev/null +++ b/src/bytetrack/calculators/bytetrack_calculator.cc @@ -0,0 +1,440 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +#include "src/bytetrack/calculators/bytetrack_calculator.h" +#include "src/bytetrack/calculators/bytetrack_calculator.pb.h" + +#include +#include +#include +#include +#include + +#include + +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/detection.pb.h" +#include "src/bytetrack/calculators/strack.h" +#include "src/bytetrack/calculators/kalman_filter.h" +#include "src/bytetrack/calculators/matching_utils.h" + +// node { +// calculator: "ByteTrackCalculator" +// input_stream: "DETECTIONS:detections_with_id" +// output_stream: "DETECTIONS:tracked_detections" +// options: { +// [mediapipe.ByteTrackCalculatorOptions.ext] { +// track_high_threshold:0.5 +// track_low_threshold:0.1 +// new_track_threshold:0.25 +// matching_threshold: 0.75 +// track_buffer: 30 +// fuse_score: true +// } +// } +// } + +namespace mediapipe { +REGISTER_CALCULATOR(ByteTrackCalculator); + +absl::Status ByteTrackCalculator::GetContract(CalculatorContract* cc) { + cc->Inputs().Get("DETECTIONS", 0).Set>(); + cc->Outputs().Get("DETECTIONS", 0).Set>(); + return absl::OkStatus(); +} + +absl::Status ByteTrackCalculator::Open(CalculatorContext* cc) { + options_ = cc->Options(); + match_thresh_ = options_.matching_threshold(); + track_buffer_ = options_.track_buffer(); + track_high_thresh_ = options_.track_high_threshold(); + track_low_thresh_ = options_.track_low_threshold(); + new_track_thresh_ = options_.new_track_threshold(); + fuse_score_ = options_.fuse_score(); + return absl::OkStatus(); +} + +absl::Status ByteTrackCalculator::Process(CalculatorContext* cc) { + if (cc->Inputs().Get("DETECTIONS", 0).IsEmpty()) { + return absl::OkStatus(); // nothing to do + } + int64_t current_ts = cc->InputTimestamp().Microseconds(); + if (last_timestamp_ > 0 && frame_id_ > 1) { + float dt_sec = (current_ts - last_timestamp_) / 1e6f; + float instant_fps = 1.0f / dt_sec; + // Smooth it with a running average to avoid jitter + estimated_fps_ = 0.9f * estimated_fps_ + 0.1f * instant_fps; + max_time_lost_ = static_cast(estimated_fps_ / 30.0f * track_buffer_); + LOG(INFO) << "MAX TIME LOST: " << max_time_lost_; + } + last_timestamp_ = current_ts; + frame_id_++; + LOG(INFO) << "Frame " << frame_id_ + << " tracked_=" << tracked_stracks_.size() + << " lost_=" << lost_stracks_.size(); + + // LOG(INFO) << "ByteTrackCalculator::Process called, frame " << frame_id_; + auto high_dets = std::make_unique>(); + auto low_dets = std::make_unique>(); + + if (!cc->Inputs().Tag("DETECTIONS").IsEmpty()) { + const auto& input = cc->Inputs().Tag("DETECTIONS").Get>(); + for (const auto& detection : input) { + float score = detection.score(0); + if (score >= track_high_thresh_) + high_dets->push_back(detection); + else if (score >= track_low_thresh_) + low_dets->push_back(detection); + } + } + + std::vector activated_stracks; + std::vector refind_stracks; + std::vector lost_stracks; + std::vector removed_stracks; + + LOG(INFO) << "High dets: " << high_dets->size() + << " Low dets: " << low_dets->size(); + + // create tracks from high score detections + std::vector detections; + if (high_dets->size() > 0) { + for (auto& d : *high_dets) { + detections.emplace_back(d); + } + } + + // add newly detected tracks + std::vector unconfirmed; + std::vector tracked_stracks; + for (auto& track : tracked_stracks_) { + if (!track.is_activated()) { + unconfirmed.push_back(&track); + } else { + tracked_stracks.push_back(&track); + } + } + + /////////////////////// First association ////////////////////////////////////// + std::vector lost_ptrs; + for (auto& t : lost_stracks_) + lost_ptrs.push_back(&t); + + auto strack_pool = JointStracks(tracked_stracks, lost_ptrs); + bytetrack::STrack::MultiPredict(strack_pool); + auto dists = BuildIoUCostMatrix(strack_pool, detections); + if (fuse_score_) { + dists = FuseScore(dists, detections); + } + // printDists(dists); + // gives 3 vectors + auto [matches, u_track, u_detection] = bytetrack::LinearAssignment(dists, match_thresh_); + + LOG(INFO) << "First step association, " + << " matches size: " << matches.size() + << ", U_track size: " << u_track.size() + << ", U_detection size: " << u_detection.size(); + + for (int k = 0; k < matches.rows(); ++k) { + int itracked = matches(k, 0); + int idet = matches(k, 1); + auto* track = strack_pool[itracked]; + auto det = detections[idet]; + if (track->state() == bytetrack::BaseTrack::TrackState::TRACKED) { + track->Update(detections[idet], frame_id_); + activated_stracks.push_back(*track); + } else { + track->ReActivate(det, frame_id_, false); + refind_stracks.push_back(*track); + } + } + LOG(INFO) << " After 1st assoc: activated=" << activated_stracks.size() + << " refind=" << refind_stracks.size(); + /////////////////////// Second association /////////////////////////// + std::vector detections_second; + if (low_dets->size() > 0) { + for (const auto& d : *low_dets) { + detections_second.emplace_back(d); + } + } + + std::vector r_tracked_stracks; + for (int i : u_track) { + if (strack_pool[i]->state() == bytetrack::BaseTrack::TrackState::TRACKED) { + r_tracked_stracks.push_back(strack_pool[i]); + } + } + + dists = BuildIoUCostMatrix(r_tracked_stracks, detections_second); + if (fuse_score_) { + dists = FuseScore(dists, detections_second); + } + + auto [matches2, u_track2, u_detection_second] = bytetrack::LinearAssignment(dists, 0.5f); + + LOG(INFO) << "Second step association, " + << " matches2 size: " << matches2.size() + << ", U_track2 size: " << u_track2.size() + << ", U_detection_second size: " << u_detection_second.size(); + + for (int k = 0; k < matches2.rows(); ++k) { + int itracked = matches2(k, 0); + int idet = matches2(k, 1); + auto* track = r_tracked_stracks[itracked]; + auto det = detections_second[idet]; + if (track->state() == bytetrack::BaseTrack::TrackState::TRACKED) { + track->Update(det, frame_id_); + activated_stracks.push_back(*track); + } else { + track->ReActivate(det, frame_id_, false); + refind_stracks.push_back(*track); + } + } + + // mark lost tracks + for (int it : u_track2) { + auto* track = r_tracked_stracks[it]; + if (track->state() != bytetrack::BaseTrack::TrackState::LOST) { + track->MarkLost(); + lost_stracks.push_back(*track); + } + } + LOG(INFO) << " After 2nd assoc: lost=" << lost_stracks.size(); + /////////////////////// DEAL W UNCONFIRMED TRACKS /////////////////////////// + std::vector detections_uc; + for (int i : u_detection) { + detections_uc.push_back(detections[i]); + } + + dists = BuildIoUCostMatrix(unconfirmed, detections_uc); + if (fuse_score_) { + dists = FuseScore(dists, detections_uc); + } + + auto [matches3, u_unconfirmed, u_detection_3] = bytetrack::LinearAssignment(dists, 0.7f); + + for (int k = 0; k < matches3.rows(); ++k) { + int itracked = matches3(k, 0); + int idet = matches3(k, 1); + unconfirmed[itracked]->Update(detections_uc[idet], frame_id_); + activated_stracks.push_back(*unconfirmed[itracked]); + } + + for (int it : u_unconfirmed) { + auto* track = unconfirmed[it]; + track->MarkRemoved(); + } + /////////////////////// INITIALIZE NEW TRACKS ///////////////////////////// + for (int inew : u_detection_3) { + auto track = detections_uc[inew]; + if (track.score() < new_track_thresh_) { + continue; + } + track.Activate(&kalman_filter_, frame_id_); + activated_stracks.push_back(track); + } + LOG(INFO) << " After unconfirmed+new: activated=" << activated_stracks.size(); + /////////////////////// UPDATE STATE ///////////////////////////// + for (auto& track : lost_stracks_) { + LOG(INFO) << "Time diff update state " << frame_id_ - track.frame_id() << ","; + if (frame_id_ - track.frame_id() > max_time_lost_) { + track.MarkRemoved(); + removed_stracks.push_back(track); + } + } + + LOG(INFO) << "actiavted_stracks size: " << activated_stracks.size() + << ",refind_stracks size: " << refind_stracks.size() + << ",lost_stracks size: " << lost_stracks.size() + << ",removed_stracks size: " << removed_stracks.size(); + + // Filter tracked_stracks_ to only TRACKED state, then join activated + refind + // (mirrors: self.tracked_stracks = [t for t in self.tracked_stracks if t.state == Tracked]) + { + std::vector only_tracked; + for (auto& t : tracked_stracks_) { + if (t.state() == bytetrack::BaseTrack::TrackState::TRACKED) + only_tracked.push_back(t); + } + tracked_stracks_ = only_tracked; + } + + // joint_stracks(tracked_stracks_, activated_stracks) + { + std::vector cur_ptrs, act_ptrs, ref_ptrs; + for (auto& t : tracked_stracks_) + cur_ptrs.push_back(&t); + for (auto& t : activated_stracks) + act_ptrs.push_back(&t); + for (auto& t : refind_stracks) + ref_ptrs.push_back(&t); + + auto joined = JointStracks(cur_ptrs, act_ptrs); + // joint_stracks(tracked_stracks_, refind_stracks) + joined = JointStracks(joined, ref_ptrs); + + std::vector joined_tracks; + joined_tracks.reserve(joined.size()); + for (auto* t : joined) + joined_tracks.push_back(*t); + tracked_stracks_ = std::move(joined_tracks); + } + + // sub_stracks(lost_stracks_, tracked_stracks_) then extend with lost_stracks (local) + { + std::vector lost_ptrs2, new_tracked_ptrs, removed_ptrs, local_lost_ptrs; + for (auto& t : lost_stracks_) + lost_ptrs2.push_back(&t); + for (auto& t : tracked_stracks_) + new_tracked_ptrs.push_back(&t); // fresh pointers! + for (auto& t : removed_stracks) + removed_ptrs.push_back(&t); + for (auto& t : lost_stracks) + local_lost_ptrs.push_back(&t); + + auto new_lost = SubStracks(lost_ptrs2, new_tracked_ptrs); + new_lost = JointStracks(new_lost, local_lost_ptrs); // extend + new_lost = SubStracks(new_lost, removed_ptrs); + + // lost_stracks_.clear(); + // for (auto* t : new_lost) lost_stracks_.push_back(*t); + std::vector l_stracks; + l_stracks.reserve(new_lost.size()); + for (auto* t : new_lost) { + if (t->state() != bytetrack::BaseTrack::TrackState::REMOVED) + l_stracks.push_back(*t); + } + lost_stracks_ = std::move(l_stracks); + } + + // remove duplicates + { + std::vector tracked_ptrs2, lost_ptrs3; + for (auto& t : tracked_stracks_) + tracked_ptrs2.push_back(&t); + for (auto& t : lost_stracks_) + lost_ptrs3.push_back(&t); + auto [dedup_tracked, dedup_lost] = RemoveDuplicateStracks(tracked_ptrs2, lost_ptrs3); + std::vector new_tracked, new_lost; + new_tracked.reserve(dedup_tracked.size()); + new_lost.reserve(dedup_lost.size()); + for (auto* t : dedup_tracked) + new_tracked.push_back(*t); + for (auto* t : dedup_lost) + new_lost.push_back(*t); + tracked_stracks_ = std::move(new_tracked); + lost_stracks_ = std::move(new_lost); + } + + LOG(INFO) << "After update state"; + LOG(INFO) << " End of frame: tracked=" << tracked_stracks_.size() + << " lost=" << lost_stracks_.size(); + + ////////////////////////////// OUTPUT /////////////////////////////////// + auto output = std::make_unique>(); + for (const auto& t : tracked_stracks_) { + if (!t.is_activated()) + continue; + Detection d; + d.set_detection_id(t.track_id()); + d.add_label(t.label()); + d.add_score(t.score()); + auto* loc = d.mutable_location_data(); + loc->set_format(LocationData::RELATIVE_BOUNDING_BOX); + auto* rb = loc->mutable_relative_bounding_box(); + Eigen::Vector4f box = t.tlwh(); + rb->set_xmin(box(0)); + rb->set_ymin(box(1)); + rb->set_width(box(2)); + rb->set_height(box(3)); + output->push_back(d); + } + LOG(INFO) << "After building detections"; + LOG(INFO) << " Output size: " << output->size(); + // cc->Outputs().Get("DETECTIONS",0).Add( + // output.release(), cc->InputTimestamp()); + cc->Outputs().Get("DETECTIONS", 0).Add(output.release(), cc->Inputs().Get("DETECTIONS", 0).Value().Timestamp()); + return absl::OkStatus(); +} + +std::vector ByteTrackCalculator::JointStracks(std::vector& a, std::vector& b) { + std::unordered_map exists; + std::vector res; + + for (auto* t : a) { + exists[t->track_id()] = true; + res.push_back(t); + } + + for (auto* t : b) { + int tid = t->track_id(); + if (exists.find(tid) == exists.end()) { + exists[tid] = true; + res.push_back(t); + } + } + return res; +} + +std::vector ByteTrackCalculator::SubStracks(std::vector& a, std::vector& b) { + std::unordered_map exists; + std::vector res; + for (auto* t : a) { + exists[t->track_id()] = t; + } + for (auto* t : b) { + int tid = t->track_id(); + if (exists.find(tid) != exists.end()) { + exists.erase(tid); + } + } + for (auto& i : exists) { + res.push_back(i.second); + } + return res; +} + +//// WIP +std::pair, std::vector> +ByteTrackCalculator::RemoveDuplicateStracks(std::vector& a, std::vector& b) { + auto pdist = BuildIoUCostMatrix(a, b); + + std::vector dupa, dupb; + for (int p = 0; p < (int)a.size(); ++p) { + for (int q = 0; q < (int)b.size(); ++q) { + if (pdist(p, q) < 0.15f) { // high overlap — duplicate + int timep = a[p]->frame_id() - a[p]->start_frame(); + int timeq = b[q]->frame_id() - b[q]->start_frame(); + if (timep > timeq) + dupb.push_back(q); + else + dupa.push_back(p); + } + } + } + + std::vector resa, resb; + for (int i = 0; i < (int)a.size(); ++i) + if (std::find(dupa.begin(), dupa.end(), i) == dupa.end()) + resa.push_back(a[i]); + for (int i = 0; i < (int)b.size(); ++i) + if (std::find(dupb.begin(), dupb.end(), i) == dupb.end()) + resb.push_back(b[i]); + + return {resa, resb}; +} + +} // namespace mediapipe diff --git a/src/bytetrack/calculators/bytetrack_calculator.h b/src/bytetrack/calculators/bytetrack_calculator.h new file mode 100644 index 0000000000..dd56e3ab9b --- /dev/null +++ b/src/bytetrack/calculators/bytetrack_calculator.h @@ -0,0 +1,71 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +#ifndef SRC_BYTETRACK_CALCULATORS_BYTETRACK_CALCULATOR_H_ +#define SRC_BYTETRACK_CALCULATORS_BYTETRACK_CALCULATOR_H_ + +#include +#include + +#include "src/bytetrack/calculators/bytetrack_calculator.pb.h" + +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/detection.pb.h" +#include "src/bytetrack/calculators/strack.h" +#include "src/bytetrack/calculators/kalman_filter.h" + +namespace mediapipe { +class ByteTrackCalculator : public CalculatorBase { +public: + static absl::Status GetContract(CalculatorContract* cc); + absl::Status Open(CalculatorContext* cc) override; + absl::Status Process(CalculatorContext* cc) override; + +private: + ::mediapipe::ByteTrackCalculatorOptions options_; + std::vector tracked_stracks_; + std::vector lost_stracks_; + + int track_buffer_; + float det_thresh_; + float match_thresh_; + float track_high_thresh_; + float track_low_thresh_; + float new_track_thresh_; + bool fuse_score_; + int frame_id_ = 0; + int max_time_lost_ = 30; + int64_t last_timestamp_ = -1; + float estimated_fps_ = 30.0f; + + bytetrack::KalmanFilter kalman_filter_; + + static std::vector JointStracks( + std::vector& a, + std::vector& b); + + static std::vector SubStracks( + std::vector& a, + std::vector& b); + + static std::pair, std::vector> + RemoveDuplicateStracks( + std::vector& a, + std::vector& b); +}; +} // namespace mediapipe + +#endif // SRC_BYTETRACK_CALCULATORS_BYTETRACK_CALCULATOR_H_ diff --git a/src/bytetrack/calculators/bytetrack_calculator.proto b/src/bytetrack/calculators/bytetrack_calculator.proto new file mode 100644 index 0000000000..d3e0af5603 --- /dev/null +++ b/src/bytetrack/calculators/bytetrack_calculator.proto @@ -0,0 +1,33 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +syntax = "proto2"; + +package mediapipe; + +import "mediapipe/framework/calculator.proto"; + +message ByteTrackCalculatorOptions { + extend mediapipe.CalculatorOptions { + optional ByteTrackCalculatorOptions ext = 247258241; + } + optional float matching_threshold = 2 [default = 0.8]; + optional int32 track_buffer = 3 [default = 30]; + optional float track_high_threshold = 4 [default = 0.25]; + optional float track_low_threshold = 5 [default = 0.1]; + optional float new_track_threshold = 6 [default = 0.25]; + optional bool fuse_score = 7 [default = true]; +} diff --git a/src/bytetrack/calculators/kalman_filter.cc b/src/bytetrack/calculators/kalman_filter.cc new file mode 100644 index 0000000000..bf83fac362 --- /dev/null +++ b/src/bytetrack/calculators/kalman_filter.cc @@ -0,0 +1,123 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +#include "src/bytetrack/calculators/kalman_matrices.h" + +#include + +#include + +#include "mediapipe/framework/formats/detection.pb.h" +#include "src/bytetrack/calculators/kalman_filter.h" + +namespace mediapipe { +namespace bytetrack { + +KalmanFilter::KalmanFilter() : + motion_mat_(MakeMotionMatrix()), + update_mat_(MakeUpdateMatrix()), + kStdWeightPos_(kStdWeightPos), + kStdWeightVel_(kStdWeightVel) {} + +std::pair KalmanFilter::Initiate(const Eigen::Vector4f detection) { + Mean mean; + mean << detection(0), detection(1), detection(2), detection(3), 0.f, 0.f, 0.f, 0.f; + + Eigen::Matrix std_vals; + std_vals << 2.f * kStdWeightPos_ * detection(3), + 2.f * kStdWeightPos_ * detection(3), + 1e-2f, + 2.f * kStdWeightPos_ * detection(3), + 10.f * kStdWeightVel_ * detection(3), + 10.f * kStdWeightVel_ * detection(3), + 1e-5f, + 10.f * kStdWeightVel_ * detection(3); + + Cov cov = std_vals.array().square().matrix().asDiagonal(); + return {mean, cov}; +} + +std::pair KalmanFilter::Predict(const Mean& mean, const Cov& cov) const { + const float h = mean(3); + Eigen::Matrix std_pv; + std_pv << kStdWeightPos_ * h, kStdWeightPos_ * h, 1e-2f, kStdWeightPos_ * h, + kStdWeightVel_ * h, kStdWeightVel_ * h, 1e-5f, kStdWeightVel_ * h; + const Cov motion_cov = std_pv.array().square().matrix().asDiagonal(); + const Mean est_mean = mean * motion_mat_.transpose(); + const Cov est_cov = motion_mat_ * cov * motion_mat_.transpose() + motion_cov; + + return {est_mean, est_cov}; +} + +std::pair KalmanFilter::Update(const Mean& mean, const Cov& cov, const Eigen::Vector4f xyah) const { + const float x_c = xyah(0); + const float y_c = xyah(1); + const float ar = xyah(2); + const float h = xyah(3); + + Eigen::Matrix measurement; + measurement << x_c, y_c, ar, h; + // Innovation covariance in measurement space + Eigen::Matrix noise_std; + noise_std << kStdWeightPos_ * mean(3), + kStdWeightPos_ * mean(3), + 1e-1f, + kStdWeightPos_ * mean(3); + const Eigen::Matrix innov_cov = noise_std.array().square().matrix().asDiagonal(); + + const Eigen::Matrix proj_mean = mean * update_mat_.transpose(); + const Eigen::Matrix proj_cov = update_mat_ * cov * update_mat_.transpose() + innov_cov; + + // Kalman gain via Cholesky solve: K = (P H^T) (H P H^T + R)^{-1} + const Eigen::Matrix PHt = cov * update_mat_.transpose(); + const Eigen::Matrix K = proj_cov.llt().solve(PHt.transpose()).transpose(); + + const Mean updated_mean = mean + (measurement - proj_mean) * K.transpose(); + const Cov updated_cov = cov - K * proj_cov * K.transpose(); + + return {updated_mean, updated_cov}; +} + +std::pair KalmanFilter::MultiPredict(const MeanMatrix& means, const CovMatrix& covs) const { + const int N = means.rows(); + + // Build Nx8 std matrix — each row is the std devs for one track + MeanMatrix std_mat(N, 8); + for (int i = 0; i < N; ++i) { + const float h = means(i, 3); + std_mat.row(i) << kStdWeightPos_ * h, kStdWeightPos_ * h, 1e-2f, kStdWeightPos_ * h, + kStdWeightVel_ * h, kStdWeightVel_ * h, 1e-5f, kStdWeightVel_ * h; + } + + // Predicted means: (N,8) @ F^T + MeanMatrix pred_means = means * motion_mat_.transpose(); + + // Predicted covariances per track + CovMatrix pred_covs(N); + for (int i = 0; i < N; ++i) { + const Cov motion_cov = std_mat.row(i) + .array() + .square() + .matrix() + .asDiagonal(); + pred_covs[i] = motion_mat_ * covs[i] * motion_mat_.transpose() + motion_cov; + } + + return {pred_means, pred_covs}; +} + +} // namespace bytetrack +} // namespace mediapipe diff --git a/src/bytetrack/calculators/kalman_filter.h b/src/bytetrack/calculators/kalman_filter.h new file mode 100644 index 0000000000..5fc6ca9575 --- /dev/null +++ b/src/bytetrack/calculators/kalman_filter.h @@ -0,0 +1,52 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +#ifndef SRC_BYTETRACK_CALCULATORS_KALMAN_FILTER_H_ +#define SRC_BYTETRACK_CALCULATORS_KALMAN_FILTER_H_ + +#include +#include + +#include + +#include "mediapipe/framework/formats/detection.pb.h" + +namespace mediapipe { +namespace bytetrack { + +using Mean = Eigen::Matrix; +using Cov = Eigen::Matrix; +using MeanMatrix = Eigen::Matrix; +using CovMatrix = std::vector; +class KalmanFilter { +public: + KalmanFilter(); + std::pair Initiate(const Eigen::Vector4f detection); + std::pair Predict(const Mean& mean, const Cov& cov) const; + std::pair Update(const Mean& mean, const Cov& cov, const Eigen::Vector4f xyah) const; + std::pair MultiPredict(const MeanMatrix& means, const CovMatrix& covs) const; + +private: + Eigen::Matrix motion_mat_; + Eigen::Matrix update_mat_; + float kStdWeightPos_; + float kStdWeightVel_; +}; + +} // namespace bytetrack +} // namespace mediapipe + +#endif // SRC_BYTETRACK_CALCULATORS_KALMAN_FILTER_H_ diff --git a/src/bytetrack/calculators/kalman_matrices.h b/src/bytetrack/calculators/kalman_matrices.h new file mode 100644 index 0000000000..c53c97721a --- /dev/null +++ b/src/bytetrack/calculators/kalman_matrices.h @@ -0,0 +1,43 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** +// kalman_matrices.h +#ifndef SRC_BYTETRACK_CALCULATORS_KALMAN_MATRICES_H_ +#define SRC_BYTETRACK_CALCULATORS_KALMAN_MATRICES_H_ + +#include + +namespace mediapipe { +namespace bytetrack { + +static constexpr float kStdWeightPos = 1.0f / 20.0f; +static constexpr float kStdWeightVel = 1.0f / 160.0f; + +inline Eigen::Matrix MakeMotionMatrix() { + Eigen::Matrix F = Eigen::Matrix::Identity(); + F.block<4, 4>(0, 4) = Eigen::Matrix4f::Identity(); + return F; +} + +inline Eigen::Matrix MakeUpdateMatrix() { + Eigen::Matrix H = Eigen::Matrix::Zero(); + H.block<4, 4>(0, 0) = Eigen::Matrix4f::Identity(); + return H; +} + +} // namespace bytetrack +} // namespace mediapipe + +#endif // SRC_BYTETRACK_CALCULATORS_KALMAN_MATRICES_H_ diff --git a/src/bytetrack/calculators/matching_utils.h b/src/bytetrack/calculators/matching_utils.h new file mode 100644 index 0000000000..5283608931 --- /dev/null +++ b/src/bytetrack/calculators/matching_utils.h @@ -0,0 +1,354 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +#ifndef SRC_BYTETRACK_CALCULATORS_MATCHING_UTILS_H_ +#define SRC_BYTETRACK_CALCULATORS_MATCHING_UTILS_H_ + +#include +#include +#include +#include + +#include + +#include "src/bytetrack/calculators/strack.h" + +namespace mediapipe { +namespace bytetrack { + +struct AssignmentResult { + Eigen::MatrixXi matches; // shape (K, 2) — col 0 = track idx, col 1 = box idx + Eigen::VectorXi unmatched_tracks; // shape (P,) + Eigen::VectorXi unmatched_boxes; // shape (Q,) +}; + +// IoU between two [top, left, bottom, right] boxes +inline float ComputeIoU(const Eigen::Vector4f& a, + const Eigen::Vector4f& b) { + float inter_x1 = std::max(a[0], b[0]); // left + float inter_y1 = std::max(a[1], b[1]); // top + float inter_x2 = std::min(a[2], b[2]); // right + float inter_y2 = std::min(a[3], b[3]); // bottom + + float inter_w = std::max(0.f, inter_x2 - inter_x1); + float inter_h = std::max(0.f, inter_y2 - inter_y1); + float inter_area = inter_w * inter_h; + if (inter_area == 0.f) + return 0.f; + + float area_a = (a[2] - a[0]) * (a[3] - a[1]); // w * h + float area_b = (b[2] - b[0]) * (b[3] - b[1]); + return inter_area / (area_a + area_b - inter_area); +} + +inline Eigen::MatrixXf BuildIoUCostMatrix( + const std::vector& tracks, + const std::vector& detections) { + int N = tracks.size(); + int M = detections.size(); + + Eigen::MatrixXf cost(N, M); + + for (int i = 0; i < N; ++i) { + auto tb = tracks[i]->tlbr(); + for (int j = 0; j < M; ++j) { + auto bb = detections[j].tlbr(); + cost(i, j) = 1.f - ComputeIoU(tb, bb); + } + } + + return cost; +} + +inline Eigen::MatrixXf FuseScore( + const Eigen::MatrixXf& cost_matrix, + const std::vector& detections) { + if (cost_matrix.size() == 0) { + return cost_matrix; + } + Eigen::MatrixXf iou_sim = 1.0f - cost_matrix.array(); + + Eigen::RowVectorXf det_scores(detections.size()); + for (int i = 0; i < (int)detections.size(); i++) { + det_scores(i) = detections[i].score(); + } + + Eigen::MatrixXf det_scores_mat = det_scores.replicate(cost_matrix.rows(), 1); + Eigen::MatrixXf fuse_sim = iou_sim.array() * det_scores_mat.array(); + Eigen::MatrixXf fuse_cost = 1.0f - fuse_sim.array(); + + return fuse_cost; +} + +inline Eigen::MatrixXf BuildIoUCostMatrix( + const std::vector& a, + const std::vector& b) { + int N = a.size(); + int M = b.size(); + + Eigen::MatrixXf cost(N, M); + + for (int i = 0; i < N; ++i) { + auto ta = a[i]->tlbr(); + for (int j = 0; j < M; ++j) { + auto tb = b[j]->tlbr(); + cost(i, j) = 1.f - ComputeIoU(ta, tb); + } + } + + return cost; +} + +// inline AssignmentResult LinearAssignment(const Eigen::MatrixXf& cost, float thresh){ +// int N = (int)cost.rows(); +// int M = (int)cost.cols(); + +// // Collect candidates below threshold +// std::vector> entries; +// entries.reserve(N * M); +// for (int i = 0; i < N; ++i) +// for (int j = 0; j < M; ++j) +// if (cost(i,j) <= thresh) +// entries.emplace_back(cost(i,j), i, j); + +// std::sort(entries.begin(), entries.end()); + +// // Greedy assignment +// std::vector track_used(N, false); +// std::vector box_used(M, false); +// std::vector> matched; +// matched.reserve(std::min(N, M)); + +// for (auto& [c, i, j] : entries) { +// if (!track_used[i] && !box_used[j]) { +// matched.emplace_back(i, j); +// track_used[i] = true; +// box_used[j] = true; +// } +// } + +// // Pack into Eigen outputs +// int K = (int)matched.size(); +// Eigen::MatrixXi matches(K, 2); // (K,2) — mirrors np.empty((0,2)) when K=0 +// for (int k = 0; k < K; ++k) { +// matches(k, 0) = matched[k].first; +// matches(k, 1) = matched[k].second; +// } + +// // Count unmatched first, then fill — avoids push_back on Eigen vectors +// int n_ut = (int)std::count(track_used.begin(), track_used.end(), false); +// int n_ub = (int)std::count(box_used.begin(), box_used.end(), false); + +// Eigen::VectorXi unmatched_tracks(n_ut); +// Eigen::VectorXi unmatched_boxes(n_ub); + +// for (int i = 0, k = 0; i < N; ++i) +// if (!track_used[i]) unmatched_tracks(k++) = i; +// for (int j = 0, k = 0; j < M; ++j) +// if (!box_used[j]) unmatched_boxes(k++) = j; + +// return {matches, unmatched_tracks, unmatched_boxes}; +// } + +inline AssignmentResult LinearAssignment(const Eigen::MatrixXf& cost, float thresh) { + int N = (int)cost.rows(); + int M = (int)cost.cols(); + + // Empty matrix early exit — mirrors Python: if cost_matrix.size == 0 + if (N == 0 || M == 0) { + Eigen::MatrixXi matches(0, 2); + Eigen::VectorXi u_tracks(N), u_boxes(M); + for (int i = 0; i < N; ++i) + u_tracks(i) = i; + for (int j = 0; j < M; ++j) + u_boxes(j) = j; + return {matches, u_tracks, u_boxes}; + } + + // Pad to square S x S — mirrors lap.lapjv extend_cost=True + int S = std::max(N, M); + const float INF = 1e9f; + + Eigen::MatrixXf cost_sq = Eigen::MatrixXf::Constant(S, S, INF); + for (int i = 0; i < N; ++i) + for (int j = 0; j < M; ++j) + cost_sq(i, j) = cost(i, j); + + // Dual variables and assignment vectors + std::vector u(S, 0.f), v(S, 0.f); + std::vector row2col(S, -1), col2row(S, -1); + + // Phase 1: Column reduction — init v[j] to column minimum + for (int j = 0; j < S; ++j) { + int best_i = 0; + float best_v = cost_sq(0, j); + for (int i = 1; i < S; ++i) { + if (cost_sq(i, j) < best_v) { + best_v = cost_sq(i, j); + best_i = i; + } + } + v[j] = best_v; + if (row2col[best_i] == -1) { + row2col[best_i] = j; + col2row[j] = best_i; + } + } + + // Phase 2: Augmenting row reduction (2 passes) + for (int pass = 0; pass < 2; ++pass) { + for (int i = 0; i < S; ++i) { + if (row2col[i] != -1) + continue; + int j1 = -1, j2 = -1; + float u1 = INF, u2 = INF; + for (int j = 0; j < S; ++j) { + float h = cost_sq(i, j) - v[j]; + if (h < u2) { + if (h < u1) { + u2 = u1; + j2 = j1; + u1 = h; + j1 = j; + } else { + u2 = h; + j2 = j; + } + } + } + u[i] = u1; + if (j1 != -1) { + if (col2row[j1] == -1) { + row2col[i] = j1; + col2row[j1] = i; + } else { + v[j1] -= (u2 - u1); + } + } + } + } + + // Phase 3: Augmentation via shortest path (Dijkstra with potentials) + std::vector dist(S); + std::vector pred(S, -1); + std::vector visited(S, false); + + for (int i_start = 0; i_start < S; ++i_start) { + if (row2col[i_start] != -1) + continue; + + std::fill(dist.begin(), dist.end(), INF); + std::fill(pred.begin(), pred.end(), -1); + std::fill(visited.begin(), visited.end(), false); + + for (int j = 0; j < S; ++j) + dist[j] = cost_sq(i_start, j) - u[i_start] - v[j]; + + int j_end = -1; + float d_min = INF; + + for (int iter = 0; iter < S; ++iter) { + // Pick unvisited col with smallest dist + int j_min = -1; + d_min = INF; + for (int j = 0; j < S; ++j) + if (!visited[j] && dist[j] < d_min) { + d_min = dist[j]; + j_min = j; + } + + if (j_min == -1) + break; + visited[j_min] = true; + + if (col2row[j_min] == -1) { + j_end = j_min; + break; + } + + // Relax edges through the row that owns j_min + int i_next = col2row[j_min]; + u[i_next] = cost_sq(i_next, j_min) - v[j_min] - d_min; // update dual + for (int j = 0; j < S; ++j) { + if (visited[j]) + continue; + float nd = d_min + cost_sq(i_next, j) - u[i_next] - v[j]; + if (nd < dist[j]) { + dist[j] = nd; + pred[j] = j_min; + } + } + } + + // Update col duals along the path + for (int j = 0; j < S; ++j) + if (visited[j]) + v[j] += dist[j] - d_min; + u[i_start] += d_min; + + // Augment: flip assignments along path back to i_start + int j_cur = j_end; + while (j_cur != -1) { + int i_cur = (pred[j_cur] == -1) ? i_start : col2row[pred[j_cur]]; + col2row[j_cur] = i_cur; + row2col[i_cur] = j_cur; + j_cur = pred[j_cur]; + } + } + + // Extract matches — mirrors: for ix, mx in enumerate(x): if mx >= 0 + // Apply cost_limit=thresh filter here + std::vector track_used(N, false); + std::vector box_used(M, false); + std::vector> matched; + + for (int i = 0; i < N; ++i) { + int j = row2col[i]; + if (j < M && cost(i, j) <= thresh) { + matched.emplace_back(i, j); + track_used[i] = true; + box_used[j] = true; + } + } + + // Pack into AssignmentResult + int K = (int)matched.size(); + Eigen::MatrixXi matches(K, 2); + for (int k = 0; k < K; ++k) { + matches(k, 0) = matched[k].first; + matches(k, 1) = matched[k].second; + } + + int n_ut = (int)std::count(track_used.begin(), track_used.end(), false); + int n_ub = (int)std::count(box_used.begin(), box_used.end(), false); + + Eigen::VectorXi unmatched_tracks(n_ut); + Eigen::VectorXi unmatched_boxes(n_ub); + + for (int i = 0, k = 0; i < N; ++i) + if (!track_used[i]) + unmatched_tracks(k++) = i; + for (int j = 0, k = 0; j < M; ++j) + if (!box_used[j]) + unmatched_boxes(k++) = j; + + return {matches, unmatched_tracks, unmatched_boxes}; +} + +} // namespace bytetrack +} // namespace mediapipe + +#endif // SRC_BYTETRACK_CALCULATORS_MATCHING_UTILS_H_ diff --git a/src/bytetrack/calculators/render_data_passthrough_calculator.cc b/src/bytetrack/calculators/render_data_passthrough_calculator.cc new file mode 100644 index 0000000000..cb9c43d40c --- /dev/null +++ b/src/bytetrack/calculators/render_data_passthrough_calculator.cc @@ -0,0 +1,47 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +#include "mediapipe/framework/calculator_framework.h" + +#include "mediapipe/util/render_data.pb.h" + +namespace mediapipe { + +class PassThroughRenderDataCalculator : public CalculatorBase { +public: + static absl::Status GetContract(CalculatorContract* cc) { + cc->Inputs().Tag("RENDER_DATA").Set(); + cc->Outputs().Tag("RENDER_DATA").Set(); + return absl::OkStatus(); + } + + absl::Status Process(CalculatorContext* cc) override { + const auto& render_data = + cc->Inputs().Tag("RENDER_DATA").Get(); + + LOG(INFO) << "RenderData: num objects = " + << render_data.render_annotations_size(); + + // Forward unchanged + cc->Outputs().Tag("RENDER_DATA").AddPacket(MakePacket(render_data).At(cc->InputTimestamp())); + + return absl::OkStatus(); + } +}; + +REGISTER_CALCULATOR(PassThroughRenderDataCalculator); + +} // namespace mediapipe diff --git a/src/bytetrack/calculators/strack.cc b/src/bytetrack/calculators/strack.cc new file mode 100644 index 0000000000..40000cf9f2 --- /dev/null +++ b/src/bytetrack/calculators/strack.cc @@ -0,0 +1,167 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +#include "src/bytetrack/calculators/basetrack.h" + +#include +#include + +#include + +#include "mediapipe/framework/formats/detection.pb.h" +#include "src/bytetrack/calculators/kalman_filter.h" +#include "src/bytetrack/calculators/kalman_matrices.h" +#include "src/bytetrack/calculators/strack.h" + +namespace mediapipe { +namespace bytetrack { + +KalmanFilter STrack::shared_kalman; + +STrack::STrack(const Detection& det) { + const auto& loc = det.location_data(); + score_ = det.score_size() > 0 ? det.score(0) : 0.0f; + label_ = det.label_size() > 0 ? det.label(0) : ""; + if (loc.format() == LocationData::RELATIVE_BOUNDING_BOX) { + const auto& rb = loc.relative_bounding_box(); + tlwh_ << rb.xmin(), rb.ymin(), rb.width(), rb.height(); + } else if (loc.format() == LocationData::BOUNDING_BOX) { + const auto& b = loc.bounding_box(); + tlwh_ << b.xmin(), b.ymin(), b.width(), b.height(); + } else { + tlwh_.setZero(); + } +} + +void STrack::Predict() { + Mean mean_state = mean(); + if (state() != TrackState::TRACKED) { + mean_state(7) = 0.f; + } + auto [new_mean, new_cov] = kf_->Predict(mean_state, cov()); + mean_ = new_mean; + cov_ = new_cov; +} + +void STrack::Activate(KalmanFilter* kalman_filter, int frame_id) { + kf_ = kalman_filter; + track_id_ = next_id(); + auto [new_mean, new_cov] = kf_->Initiate(TlwhToXyah(tlwh_)); + mean_ = new_mean; + cov_ = new_cov; + + tracklet_len_ = 0; + state_ = TrackState::TRACKED; + if (frame_id == 1) + is_activated_ = true; + frame_id_ = frame_id; + start_frame_ = frame_id; +} + +void STrack::ReActivate(const STrack& new_track, int frame_id, bool new_id) { + Eigen::Vector4f new_tlwh = new_track.tlwh_; + auto [new_mean, new_cov] = kf_->Update(mean_, cov_, TlwhToXyah(new_tlwh)); + mean_ = new_mean; + cov_ = new_cov; + tracklet_len_ = 0; + state_ = TrackState::TRACKED; + is_activated_ = true; + + frame_id_ = frame_id; + if (new_id) + track_id_ = next_id(); + score_ = new_track.score(); +} + +void STrack::Update(const STrack& new_track, int frame_id) { + frame_id_ = frame_id; + tracklet_len_ += 1; + + Eigen::Vector4f new_tlwh = new_track.tlwh_; + auto [new_mean, new_cov] = kf_->Update(mean_, cov_, TlwhToXyah(new_tlwh)); + mean_ = new_mean; + cov_ = new_cov; + + state_ = TrackState::TRACKED; + is_activated_ = true; + score_ = new_track.score(); +} + +void STrack::MultiPredict(std::vector& tracks) { + std::size_t n = tracks.size(); + if (n > 0) { + MeanMatrix multi_mean(n, 8); + for (int i = 0; i < n; i++) { + multi_mean.row(i) = tracks[i]->mean_; + } + CovMatrix multi_cov(n); + for (int i = 0; i < n; i++) { + multi_cov[i] = tracks[i]->cov_; + } + for (int i = 0; i < n; i++) { + if (tracks[i]->state_ != TrackState::TRACKED) { + multi_mean(i, 7) = 0.f; + } + } + auto [updated_means, updated_covs] = shared_kalman.MultiPredict(multi_mean, multi_cov); + for (int i = 0; i < n; i++) { + tracks[i]->mean_ = updated_means.row(i); + tracks[i]->cov_ = updated_covs[i]; + } + } +} + +Eigen::Vector4f STrack::TlwhToXyah(const Eigen::Vector4f& tlwh) { + Eigen::Vector4f xyah; + float x = tlwh(0); + float y = tlwh(1); + float w = tlwh(2); + float h = tlwh(3); + + xyah(0) = x + w / 2.0f; + xyah(1) = y + h / 2.0f; + xyah(2) = w / h; + xyah(3) = h; + + return xyah; +} + +Eigen::Vector4f STrack::tlwh() const { + // Before activation — return the raw detection box + if (kf_ == nullptr) + return tlwh_; + + // After activation — reconstruct from Kalman mean + // mean_ = [cx, cy, ar, h, vx, vy, var, vh] + Eigen::Vector4f ret; + ret(0) = mean_(0); + ret(1) = mean_(1); + ret(2) = mean_(2) * mean_(3); + ret(3) = mean_(3); + ret(0) -= ret(2) / 2.f; + ret(1) -= ret(3) / 2.f; + return ret; +} + +Eigen::Vector4f STrack::tlbr() const { + Eigen::Vector4f ret = tlwh(); + ret(2) += ret(0); + ret(3) += ret(1); + return ret; +} + +} // namespace bytetrack +} // namespace mediapipe diff --git a/src/bytetrack/calculators/strack.h b/src/bytetrack/calculators/strack.h new file mode 100644 index 0000000000..45cefb6848 --- /dev/null +++ b/src/bytetrack/calculators/strack.h @@ -0,0 +1,76 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +#ifndef SRC_BYTETRACK_CALCULATORS_STRACK_H_ +#define SRC_BYTETRACK_CALCULATORS_STRACK_H_ + +#include +#include + +#include + +#include "mediapipe/framework/formats/detection.pb.h" +#include "src/bytetrack/calculators/kalman_filter.h" +#include "src/bytetrack/calculators/basetrack.h" + +namespace mediapipe { +namespace bytetrack { + +using Mean = Eigen::Matrix; +using Cov = Eigen::Matrix; +using MeanMatrix = Eigen::Matrix; +using CovMatrix = std::vector; + +class STrack : public BaseTrack { +public: + static KalmanFilter shared_kalman; + + explicit STrack(const Detection& det); + + void Activate(KalmanFilter* kalman_filter, int frame_id); + void ReActivate(const STrack& new_track, int frame_id, bool new_id = false); + void Update(const STrack& new_track, int frame_id); + + void Predict(); + static void MultiPredict(std::vector& tracks); + + Eigen::Vector4f tlwh() const; + Eigen::Vector4f tlbr() const; + // static Eigen::Vector4f TlbrToTlwh(const Eigen::Vector4f& tlbr); + // static Eigen::Vector4f TlwhToTlbr(const Eigen::Vector4f& tlwh); + static Eigen::Vector4f TlwhToXyah(const Eigen::Vector4f& tlwh); + + // KalmanState ToProto() const; + int tracklet_len() const { return tracklet_len_; } + const std::string& label() const { return label_; } + const Mean& mean() const { return mean_; } + const Cov& cov() const { return cov_; } + +private: + // STrack-only — Kalman state and detection origin + // score_, is_activated_, tracklet_len_ are inherited from BaseTrack — do NOT redeclare + Eigen::Vector4f tlwh_; + std::string label_; + KalmanFilter* kf_ = nullptr; + Mean mean_ = Mean::Zero(); + Cov cov_ = Cov::Zero(); + int tracklet_len_ = 0; +}; + +} // namespace bytetrack +} // namespace mediapipe + +#endif // SRC_BYTETRACK_CALCULATORS_STRACK_H_ diff --git a/src/bytetrack/utils/BUILD b/src/bytetrack/utils/BUILD new file mode 100644 index 0000000000..b3d3a88928 --- /dev/null +++ b/src/bytetrack/utils/BUILD @@ -0,0 +1,45 @@ +# +# Copyright (c) 2020 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +load("@mediapipe//mediapipe/framework/port:build_config.bzl", "mediapipe_proto_library") + +licenses(["notice"]) + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "detection_color_by_id_calculator", + srcs = ["detection_color_by_id_calculator.cc"], + deps=[ + ":detection_color_by_id_calculator_cc_proto", + "@mediapipe//mediapipe/util:render_data_cc_proto", + "@mediapipe//mediapipe/framework:calculator_framework", + "@mediapipe//mediapipe/framework/formats:detection_cc_proto", + "@mediapipe//mediapipe/framework/formats:location_data_cc_proto", + "@mediapipe//mediapipe/framework/port:ret_check", + "@mediapipe//mediapipe/util:color_cc_proto", + ], + alwayslink = 1, +) + +mediapipe_proto_library( + name = "detection_color_by_id_calculator_proto", + srcs = ["detection_color_by_id_calculator.proto"], + deps = [ + "@mediapipe//mediapipe/framework:calculator_options_proto", + "@mediapipe//mediapipe/framework:calculator_proto", + ], +) \ No newline at end of file diff --git a/src/bytetrack/utils/detection_color_by_id_calculator.cc b/src/bytetrack/utils/detection_color_by_id_calculator.cc new file mode 100644 index 0000000000..9b8d85b4d3 --- /dev/null +++ b/src/bytetrack/utils/detection_color_by_id_calculator.cc @@ -0,0 +1,143 @@ +//***************************************************************************** +// Copyright (c) 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** +#include "src/bytetrack/utils/detection_color_by_id_calculator.pb.h" + +#include +#include +#include +#include +#include + +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/detection.pb.h" +#include "mediapipe/util/render_data.pb.h" +#include "mediapipe/util/color.pb.h" + +namespace mediapipe { + +class DetectionColorByIdCalculator : public CalculatorBase { +public: + static absl::Status GetContract(CalculatorContract* cc) { + cc->Inputs().Tag("DETECTIONS").Set>(); + cc->Outputs().Tag("RENDER_DATA").Set(); + return absl::OkStatus(); + } + absl::Status Open(CalculatorContext* cc) override { + const auto& options = + cc->Options(); + thickness_ = options.has_thickness() ? options.thickness() : 5.0f; + saturation_ = options.has_saturation() ? options.saturation() : 0.85f; + value_ = options.has_value() ? options.value() : 0.95f; + return absl::OkStatus(); + } + absl::Status Process(CalculatorContext* cc) override { + const auto& detections = + cc->Inputs().Tag("DETECTIONS").Get>(); + + auto render_data = std::make_unique(); + + for (const auto& det : detections) { + int id = det.detection_id(); + mediapipe::Color color = IdToColor(id); + + const auto& bbox = det.location_data().relative_bounding_box(); + + // ── 1. Bounding box ────────────────────────────────────────────────── + { + auto* a = render_data->add_render_annotations(); + *a->mutable_color() = color; + a->set_thickness(thickness_); + + auto* rect = a->mutable_rectangle(); + rect->set_left(bbox.xmin()); + rect->set_top(bbox.ymin()); + rect->set_right(bbox.xmin() + bbox.width()); + rect->set_bottom(bbox.ymin() + bbox.height()); + rect->set_normalized(true); + } + + // ── 2. Label ─────────────────────────────────────────────────────── + { + auto* a = render_data->add_render_annotations(); + *a->mutable_color() = color; // same color as box + a->set_thickness(thickness_ - 1.0f); + + auto* text = a->mutable_text(); + std::string label = "ID:" + std::to_string(id); + if (!det.label().empty()) + label += " " + det.label(0); + if (det.score_size() > 0) { + char buf[8]; + std::snprintf(buf, sizeof(buf), " %.2f", det.score(0)); + label += buf; + } + + text->set_display_text(label); + text->set_normalized(true); + text->set_left(bbox.xmin() + 0.005f); + text->set_baseline(bbox.ymin() + 0.04f); + text->set_font_height(0.035f); + } + } + + cc->Outputs().Tag("RENDER_DATA").Add(render_data.release(), cc->InputTimestamp()); + return absl::OkStatus(); + } + +private: + mediapipe::Color IdToColor(int id) { + // Golden angle ensures max visual distance between consecutive IDs + const float kGoldenAngle = 137.508f; + float hue = std::fmod(id * kGoldenAngle, 360.0f); + float chroma = value_ * saturation_; + float x = chroma * (1.0f - std::fabs(std::fmod(hue / 60.0f, 2.0f) - 1.0f)); + float m = value_ - chroma; + + float r = 0, g = 0, b = 0; + if (hue < 60) { + r = chroma; + g = x; + } else if (hue < 120) { + r = x; + g = chroma; + } else if (hue < 180) { + g = chroma; + b = x; + } else if (hue < 240) { + g = x; + b = chroma; + } else if (hue < 300) { + r = x; + b = chroma; + } else { + r = chroma; + b = x; + } + + mediapipe::Color color; + color.set_r(static_cast((r + m) * 255)); + color.set_g(static_cast((g + m) * 255)); + color.set_b(static_cast((b + m) * 255)); + return color; + } + float thickness_ = 4.0f; + float saturation_ = 0.85f; + float value_ = 0.95f; +}; + +REGISTER_CALCULATOR(DetectionColorByIdCalculator); + +} // namespace mediapipe diff --git a/src/bytetrack/utils/detection_color_by_id_calculator.proto b/src/bytetrack/utils/detection_color_by_id_calculator.proto new file mode 100644 index 0000000000..d7bcdd05b5 --- /dev/null +++ b/src/bytetrack/utils/detection_color_by_id_calculator.proto @@ -0,0 +1,31 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +syntax = "proto2"; + +package mediapipe; + +import "mediapipe/framework/calculator.proto"; + +message DetectionColorByIdCalculatorOptions{ + // hue is calculated using track_id* 137.508f % 360, In order to get unique color for bbox based on id + extend CalculatorOptions { + optional DetectionColorByIdCalculatorOptions ext = 259397841; + } + optional float saturation = 1 [default = 0.85]; + optional float value = 2 [default = 0.95]; + optional float thickness = 3 [default = 5.0]; +} \ No newline at end of file diff --git a/src/test/mediapipeflow_test.cpp b/src/test/mediapipeflow_test.cpp index c9fa5c22f2..6daa122d0e 100644 --- a/src/test/mediapipeflow_test.cpp +++ b/src/test/mediapipeflow_test.cpp @@ -3811,6 +3811,7 @@ TEST(WhitelistRegistered, MediapipeCalculatorsList) { "BeginLoopUint64tCalculator", "BoxDetectorCalculator", "BoxTrackerCalculator", + "ByteTrackCalculator", "CallbackCalculator", "CallbackPacketCalculator", "CallbackWithHeaderCalculator", @@ -3849,6 +3850,9 @@ TEST(WhitelistRegistered, MediapipeCalculatorsList) { "DetectionClassificationCombinerCalculator", "DetectionClassificationResultCalculator", "DetectionClassificationSerializationCalculator", + "DetectionColorByIdCalculator", + "PassThroughRenderDataCalculator", + "OVYoloXTensorsToDetectionsCalculator", "DetectionExtractionCalculator", "DetectionLabelIdToTextCalculator", "DetectionLetterboxRemovalCalculator", diff --git a/src/yolox/BUILD b/src/yolox/BUILD new file mode 100644 index 0000000000..027fac018e --- /dev/null +++ b/src/yolox/BUILD @@ -0,0 +1,53 @@ +# +# Copyright (c) 2020 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +load("@mediapipe//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library", "mediapipe_proto_library") + +mediapipe_proto_library( + name = "ov_yolox_tensors_to_detections_calculator_proto", + srcs = [ + "ov_yolox_tensors_to_detections_calculator.proto", + ], + deps = [ + "@mediapipe//mediapipe/framework:calculator_options_proto", + "@mediapipe//mediapipe/framework:calculator_proto" + ], + visibility = [ + "//src:__pkg__", + ], +) + + +cc_library( + name = "ov_yolox_tensors_to_detections_calculator", + srcs = [ + "ov_yolox_tensors_to_detections_calculator.cc", + ], + deps = [ + "//third_party:openvino", + "@mediapipe//mediapipe/framework:calculator_framework", + "@mediapipe//mediapipe/framework/formats:detection_cc_proto", + "@mediapipe//mediapipe/framework/formats:location_data_cc_proto", + "@mediapipe//mediapipe/framework/port:ret_check", + "@mediapipe//mediapipe/framework/port:status", + "@mediapipe//mediapipe/framework:calculator_cc_proto", + ":ov_yolox_tensors_to_detections_calculator_cc_proto", + ], + visibility = [ + "//src:__pkg__", + ], + alwayslink = True +) \ No newline at end of file diff --git a/src/yolox/ov_yolox_tensors_to_detections_calculator.cc b/src/yolox/ov_yolox_tensors_to_detections_calculator.cc new file mode 100644 index 0000000000..6ad749a454 --- /dev/null +++ b/src/yolox/ov_yolox_tensors_to_detections_calculator.cc @@ -0,0 +1,162 @@ +//***************************************************************************** +// Copyright (c) 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +#include +#include + +#include + +#include "src/yolox/ov_yolox_tensors_to_detections_calculator.pb.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/detection.pb.h" +#include "mediapipe/framework/formats/location_data.pb.h" +#include "mediapipe/framework/port/ret_check.h" +#include "mediapipe/framework/port/status.h" + +namespace mediapipe { + +class OVYoloXTensorsToDetectionsCalculator : public CalculatorBase { + static const std::string INPUT_TAG_NAME; + static const std::string OUTPUT_TAG_NAME; + + float confidenceThreshold_ = 0.1f; + float inputSize_ = 416.0f; + +public: + static absl::Status GetContract(CalculatorContract* cc) { + RET_CHECK(!cc->Inputs().GetTags().empty()); + RET_CHECK(!cc->Outputs().GetTags().empty()); + + cc->Inputs() + .Tag(INPUT_TAG_NAME) + .Set>(); + + cc->Outputs() + .Tag(OUTPUT_TAG_NAME) + .Set>(); + + return absl::OkStatus(); + } + + absl::Status Open(CalculatorContext* cc) final { + const auto& options = + cc->Options< + mediapipe::OVYoloXTensorsToDetectionsCalculatorOptions>(); + + confidenceThreshold_ = + options.has_conf_thresh() ? options.conf_thresh() : 0.1f; + + inputSize_ = + options.has_input_size() ? options.input_size() : 416.0f; + + return absl::OkStatus(); + } + + absl::Status Process(CalculatorContext* cc) final { + if (cc->Inputs().Tag(INPUT_TAG_NAME).IsEmpty()) { + return absl::OkStatus(); + } + + const auto& tensors = + cc->Inputs() + .Tag(INPUT_TAG_NAME) + .Get>(); + + RET_CHECK_EQ(tensors.size(), 2u); + + const auto& boxesTensor = tensors[0]; + const auto& labelsTensor = tensors[1]; + + RET_CHECK_EQ(boxesTensor.get_element_type(), ov::element::f32); + RET_CHECK_EQ(labelsTensor.get_element_type(), ov::element::i64); + + const auto boxesShape = boxesTensor.get_shape(); + const auto labelsShape = labelsTensor.get_shape(); + + RET_CHECK_EQ(boxesShape.size(), 3u); + RET_CHECK_EQ(boxesShape[0], 1u); + RET_CHECK_EQ(boxesShape[2], 5u); + + RET_CHECK_EQ(labelsShape.size(), 2u); + RET_CHECK_EQ(labelsShape[0], 1u); + RET_CHECK_EQ(labelsShape[1], boxesShape[1]); + + const size_t numBoxes = boxesShape[1]; + + const float* boxes = boxesTensor.data(); + const int64_t* labels = labelsTensor.data(); + + RET_CHECK(boxes != nullptr); + RET_CHECK(labels != nullptr); + + auto detections = + absl::make_unique>(); + + for (size_t i = 0; i < numBoxes; ++i) { + const float x1 = boxes[i * 5 + 0]; + const float y1 = boxes[i * 5 + 1]; + const float x2 = boxes[i * 5 + 2]; + const float y2 = boxes[i * 5 + 3]; + const float confidence = boxes[i * 5 + 4]; + + if (confidence < confidenceThreshold_) { + continue; + } + + if (x2 <= x1 || y2 <= y1) { + continue; + } + + Detection detection; + + auto* locationData = + detection.mutable_location_data(); + + locationData->set_format( + LocationData::RELATIVE_BOUNDING_BOX); + + auto* boundingBox = + locationData->mutable_relative_bounding_box(); + + boundingBox->set_xmin(x1 / inputSize_); + boundingBox->set_ymin(y1 / inputSize_); + boundingBox->set_width((x2 - x1) / inputSize_); + boundingBox->set_height((y2 - y1) / inputSize_); + + detection.add_score(confidence); + detection.add_label_id( + static_cast(labels[i])); + + detections->emplace_back(std::move(detection)); + } + + cc->Outputs() + .Tag(OUTPUT_TAG_NAME) + .Add(detections.release(), cc->InputTimestamp()); + + return absl::OkStatus(); + } +}; + +const std::string OVYoloXTensorsToDetectionsCalculator::INPUT_TAG_NAME{ + "TENSORS"}; + +const std::string OVYoloXTensorsToDetectionsCalculator::OUTPUT_TAG_NAME{ + "DETECTIONS"}; + +REGISTER_CALCULATOR(OVYoloXTensorsToDetectionsCalculator); + +} // namespace mediapipe diff --git a/src/yolox/ov_yolox_tensors_to_detections_calculator.proto b/src/yolox/ov_yolox_tensors_to_detections_calculator.proto new file mode 100644 index 0000000000..81cacceeb4 --- /dev/null +++ b/src/yolox/ov_yolox_tensors_to_detections_calculator.proto @@ -0,0 +1,33 @@ +//***************************************************************************** +// Copyright 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//***************************************************************************** + +syntax = "proto2"; + +package mediapipe; + +import "mediapipe/framework/calculator.proto"; + +message OVYoloXTensorsToDetectionsCalculatorOptions { + extend .mediapipe.CalculatorOptions { + optional OVYoloXTensorsToDetectionsCalculatorOptions ext = 211376658; + } + + // Minimum confidence required for a detection. + optional float conf_thresh = 1 [default = 0.10]; + + // Input resolution of the YOLOX model. + optional float input_size = 3 [default = 416.0]; +} \ No newline at end of file