diff --git a/.gitignore b/.gitignore index f2d1c34..84441a4 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,12 @@ +# C/C++ build artifacts +build/ +cmake-build-*/ +CMakeCache.txt +CMakeFiles/ +*.o +*.a +*.obj + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..3c9808f --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,53 @@ +cmake_minimum_required(VERSION 3.16) +project(videoanalytics LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_EXTENSIONS OFF) + +if(NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE Release) +endif() + +find_package(OpenCV REQUIRED) +find_package(Threads REQUIRED) +find_package(nlohmann_json QUIET) + +# ZeroMQ (libzmq). Prefer pkg-config, fall back to plain library search. +find_package(PkgConfig QUIET) +if(PkgConfig_FOUND) + pkg_check_modules(ZMQ QUIET libzmq) +endif() +if(NOT ZMQ_FOUND) + find_path(ZMQ_INCLUDE_DIRS zmq.h) + find_library(ZMQ_LIBRARIES NAMES zmq libzmq) +endif() +if(NOT ZMQ_LIBRARIES) + message(FATAL_ERROR "libzmq not found. Install libzmq3-dev (or equivalent).") +endif() + +add_library(va_core + src/config.cpp + src/camera.cpp + src/broker.cpp + src/worker.cpp + src/gui.cpp + src/psnr.cpp + src/detector.cpp +) +target_include_directories(va_core PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${OpenCV_INCLUDE_DIRS} + ${ZMQ_INCLUDE_DIRS} +) +target_link_libraries(va_core PUBLIC + ${OpenCV_LIBS} + ${ZMQ_LIBRARIES} + Threads::Threads +) +if(nlohmann_json_FOUND) + target_link_libraries(va_core PUBLIC nlohmann_json::nlohmann_json) +endif() + +add_executable(videoanalytics src/main.cpp) +target_link_libraries(videoanalytics PRIVATE va_core) diff --git a/README.md b/README.md index 9a657f4..49c0973 100644 --- a/README.md +++ b/README.md @@ -1,52 +1,117 @@ # videoanalytics -This project is about videosources emergency detection made with ZeroMQ infrastructure +This project is about videosources emergency detection made with ZeroMQ infrastructure. + +This is the C++ port of the project (originally written in Python). The +architecture is unchanged: `cameras -> broker -> workers -> gui`, all connected +over ZeroMQ. ## Used technologies -- ZeroMQ -- YOLOv8 +- ZeroMQ (libzmq / cppzmq) +- YOLOv8 (via OpenCV DNN, ONNX models) - PSNR - OpenCV -- Multiprocessing & Threading +- C++17 threads + +## Architecture + +``` +cameras (PUB) --> broker (SUB -> PUSH) --> workers (PULL -> PUSH) --> gui (PULL) +``` + +- **cameras** – one thread per source reads a video and publishes JPEG frames + tagged with a 2-byte source id. +- **broker** – subscribes to every source and routes each frame to the worker + whose index matches the source id. +- **workers** – decode the frame, run the fire and forklift detectors plus PSNR + concurrently, and forward `[address, jpeg, detections_json, psnr]`. +- **gui** – renders the latest annotated frame of every source on a grid with + bounding boxes and `WARNING` / `FIRE` overlays. ## Prerequisites -- Torch with CUDA -- Python3.12 +- A C++17 compiler +- CMake >= 3.16 +- OpenCV 4 (with the `dnn` module) +- ZeroMQ (`libzmq`) and the `cppzmq` header (`zmq.hpp`) +- nlohmann/json + +On Debian/Ubuntu: + +```bash +sudo apt-get install -y build-essential cmake libzmq3-dev libopencv-dev nlohmann-json3-dev pkg-config +``` + +## Models + +The original project shipped Ultralytics PyTorch checkpoints (`.pt`). The C++ +port runs the exported ONNX versions through OpenCV's DNN module. Export them +once with Ultralytics: + +```bash +yolo export model=forklift_8s.pt format=onnx +yolo export model=fire.pt format=onnx +``` + +Place the resulting `.onnx` files under `./models` and point `config.json` at +them. Class names default to `forklift` and `fire`; adjust them in +`src/worker.cpp` if your models use different labels. ## Preparations -- `pip install -r requirements.txt` -- Download all necessary files from [Google Drive](https://drive.google.com/drive/u/0/folders/1OI_XtRNcwbm-JvojeKGR_x1SE1GuonQq) : - - (optional) place test videos & meanframes to `./videos` folder - - place YOLO models to `./models` folder -- set up `config.json` with actual files locations, ip addresses & ports, psnr threshold +- Download the necessary media/models and place test videos & meanframes in + `./videos` and YOLO ONNX models in `./models`. +- Set up `config.json` with actual file locations, ip addresses & ports, and the + psnr threshold. Notes: -> meanframes must have same resolution as original videos, '.jpg' hardcoded
-> meanframes were created on "normal" parts of videos with `psnr.py` +> meanframes must have the same resolution as the original videos, '.jpg' hardcoded
+> meanframes are created on "normal" parts of videos with the `psnr` subcommand -## Meanframes for PSNR on custom videos -Create meanframes for PSNR function: -- Prepare "normal" videofragment (not containing emergency situations) in video editor -- Use `psnr.py`: - - `python psnr.py ` -- Meanframe will be created in same directory with same filename in '.jpg' format +## Build + +```bash +mkdir -p build && cd build +cmake .. +make -j +``` + +This produces a single `videoanalytics` executable. ## Usage -After you've set up config.json, just: + +Run the full pipeline (broker + workers + cameras + gui): + ```bash -python run.py +./build/videoanalytics ``` -`Ctrl+C` to stop + +`ESC` (in the GUI window) to stop. + +Individual components can also be launched separately: + +```bash +./build/videoanalytics broker +./build/videoanalytics workers +./build/videoanalytics cameras +./build/videoanalytics gui +``` + +## Meanframes for PSNR on custom videos +Create meanframes for the PSNR function: +- Prepare a "normal" video fragment (not containing emergency situations). +- Run the meanframe tool: + - `./build/videoanalytics psnr ` +- The meanframe is created in the same directory with the same filename in + '.jpg' format. ## TODOs - create meanframes automatically - train models on another datasets - add resolution to config - Ctrl+C handler -- Auto scale processes +- Auto scale processes/threads - Test on realtime videosources ## Example -![example](./media/example_videoanalytics.gif) \ No newline at end of file +![example](./media/example_videoanalytics.gif) diff --git a/broker.py b/broker.py deleted file mode 100644 index c323832..0000000 --- a/broker.py +++ /dev/null @@ -1,26 +0,0 @@ -import zmq -from config import config - - -def main(): - # Prepare our context and router - context = zmq.Context() - - subscriber = context.socket(zmq.SUB) - for s in config['sources']: - subscriber.connect(s['port']) - subscriber.setsockopt(zmq.SUBSCRIBE, b"") - - workers = [] - for w in config['workers']: - worker = context.socket(zmq.PUSH) - worker.connect(w) - workers.append(worker) - - while True: - [address, contents] = subscriber.recv_multipart() - workers[int.from_bytes(address)].send_multipart([address, contents]) - - # We never get here but clean up anyhow - subscriber.close() - context.term() \ No newline at end of file diff --git a/cameras/cam.py b/cameras/cam.py deleted file mode 100644 index b2310cd..0000000 --- a/cameras/cam.py +++ /dev/null @@ -1,45 +0,0 @@ -import zmq -import cv2 -from multiprocessing import Process -from time import sleep -from config import config - -def start_source(id, source: dict): - """Starts videosource process - - Args: - id (int): defines source id - source (dict): videosource dict having 'port' and 'path' fields - """ - context = zmq.Context() - publisher = context.socket(zmq.PUB) - publisher.bind(source['port']) - - cap = cv2.VideoCapture(source['path']) - count = 0 - while True: - success, img = cap.read() - - # skip few frames to be more realtime-like - count += 5 - cap.set(cv2.CAP_PROP_POS_FRAMES, count) - if not success: - break - - img = cv2.imencode('.jpg', img)[1].tobytes() - publisher.send_multipart([id.to_bytes(2, 'big'), img]) - sleep(0.1) # delay to avoid frames skipping - - publisher.close() - context.term() - - -def run(): - """generate videosource processes""" - - processes = [] - for id, source in enumerate(config['sources']): - processes.append(Process(target=start_source, args=(id, source))) - print('inited process') - for process in processes: - process.start() diff --git a/colors.py b/colors.py deleted file mode 100644 index 90ad4be..0000000 --- a/colors.py +++ /dev/null @@ -1,57 +0,0 @@ -import cv2 -import numpy as np -import matplotlib.pyplot as plt - -def calculate_color_distribution(frame): - # Reshape the frame to be a 2D array of pixels - pixels = frame.reshape(-1, 3) - - # Use numpy to count the occurrences of each color - unique_colors, counts = np.unique(pixels, axis=0, return_counts=True) - - # Normalize the counts to get a distribution - distribution = counts / np.sum(counts) - - return unique_colors, distribution - -def plot_color_distribution(unique_colors, distribution): - # Plot the color distribution - plt.figure(figsize=(10, 5)) - plt.bar(unique_colors[:,0], distribution, width=0.01) - plt.xlabel('Color') - plt.ylabel('Distribution') - plt.title('Color Distribution of Video Frame') - plt.show() - -def main(): - # Open the video file - video_path = './videos/alumin.mp4' - cap = cv2.VideoCapture(video_path) - - # Check if video opened successfully - if not cap.isOpened(): - print("Error opening video file") - return - - # Read and process each frame - while True: - ret, frame = cap.read() - if not ret: - break - - # Calculate color distribution for the current frame - unique_colors, distribution = calculate_color_distribution(frame) - - # Plot the color distribution - plot_color_distribution(unique_colors, distribution) - - # Press 'q' to quit - if cv2.waitKey(1) & 0xFF == ord('q'): - break - - # Release the video capture object - cap.release() - cv2.destroyAllWindows() - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/config.json b/config.json index c2c827f..36819e4 100644 --- a/config.json +++ b/config.json @@ -28,8 +28,8 @@ "tcp://127.0.0.1:5565" ], "models": { - "fire": "./models/fire.pt", - "forklift": "./models/forklift_8s.pt" + "fire": "./models/fire.onnx", + "forklift": "./models/forklift_8s.onnx" }, "gui": "tcp://127.0.0.1:5562", "psnr_threshold": 25 diff --git a/config.py b/config.py deleted file mode 100644 index a74e9bf..0000000 --- a/config.py +++ /dev/null @@ -1,5 +0,0 @@ -import json - -global config -with open('./config.json', 'r') as f: - config = json.load(f) \ No newline at end of file diff --git a/gui.py b/gui.py deleted file mode 100644 index 97c8f1a..0000000 --- a/gui.py +++ /dev/null @@ -1,93 +0,0 @@ -import zmq -import numpy as np -import cv2 -import json -from config import config - - -class GUI: - def __init__(self): - self.width = 1920 - self.height = 1080 - self._canvas = np.ones((self.height, self.width, 3), dtype=np.uint8) * 40 # gray canvas - self.grid = (2, 2) # (WIDTH, HEIGHT) - self.cells = np.zeros((self.grid[0]*self.grid[1], self.height, self.width, 3)) - - @property - def canvas(self): - return self._canvas - - def update_cell(self, n: int, img): - """Updated n-th videosource - - Args: - n (int): source id - img (_type_): image - """ - idx = (n // self.grid[0], n % self.grid[0]) - w_pix = int(self.width / self.grid[0]) - h_pix = int(self.height / self.grid[1]) - - w_start = idx[0] * w_pix - h_start = idx[1] * h_pix - - img = cv2.resize(img, (w_pix, h_pix), interpolation=cv2.INTER_CUBIC) - self._canvas[h_start:h_start+h_pix, w_start:w_start+w_pix] = img - - -def main(): - - # Prepare our context and publisher - context = zmq.Context() - puller = context.socket(zmq.PULL) - puller.bind(config['gui']) - print('GUI') - gui = GUI() - - while True: - addr, img, dets, psnr = puller.recv_multipart() - print(addr, dets, psnr) - addr = int.from_bytes(addr) - psnr = psnr.decode() - - img = cv2.imdecode(np.frombuffer(img, dtype=np.uint8), cv2.IMREAD_COLOR) - dets = json.loads(dets.decode()) - - isFire = False - for d in dets: - # if d['name'] == 'forklift' and addr == 2: - # continue - print(f'CONFIDENCE {d['confidence']} CLASS {d['name']}') - x1 = int(d['box']['x1']) - y1 = int(d['box']['y1']) - x2 = int(d['box']['x2']) - y2 = int(d['box']['y2']) - img = cv2.rectangle(img, (x1, y1), (x2, y2), (255,0,0), 2) - name = d['name'] - if name == 'fire': - isFire = True - img = cv2.putText(img, name, (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2, cv2.LINE_AA) - - img = cv2.putText(img, f'PSNR: {psnr}', (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA) - - if int(psnr) <= config['psnr_threshold']: - img = cv2.putText(img, 'WARNING!', (250, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 4, cv2.LINE_AA) - - if isFire: - img = cv2.putText(img, 'FIRE!', (450, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 4, cv2.LINE_AA) - - gui.update_cell(addr, img) - - cv2.imshow('Image', gui.canvas) - if cv2.waitKey(10) & 0xFF == 27: - cv2.destroyAllWindows() - break - - - # We never get here but clean up anyhow - puller.close() - context.term() - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/psnr.py b/psnr.py deleted file mode 100644 index 60ad8dc..0000000 --- a/psnr.py +++ /dev/null @@ -1,71 +0,0 @@ -import cv2 -import numpy as np -from skimage.metrics import peak_signal_noise_ratio -import matplotlib.pyplot as plt -import argparse -import sys - - -def get_meanframe(normal_videopath, meanframe_path): - """Creates meanframe of video - - Args: - normal_videopath (str): path of "normal" video - meanframe_path (str): path to save meanframe to - """ - before = cv2.VideoCapture(normal_videopath) - success, img = before.read() - frames = 0 - avg = np.zeros((720,1280,3), dtype=np.float64) - - while True: - success, img = before.read() - if not success or img is None: - break - avg = np.add(avg, img) - frames += 1 - - before.release() - avg = np.divide(avg, frames).astype(np.uint8) - - cv2.imwrite(meanframe_path, avg) - -def plot_psnr(meanframe_path, video_path): - """Plot PSNR plot (X axis - frames, Y axis - PSNR in dB) - - Args: - meanframe_path (str): destination of meanframe in .jpg format - video_path (str): destination of video with emergency situation - """ - meanframe = cv2.imread(meanframe_path) - - PSNR = [] - before = cv2.VideoCapture(video_path) - success, img = before.read() - - while True: - success, img = before.read() - if not success or img is None: - break - PSNR.append(peak_signal_noise_ratio(meanframe, img)) - before.release() - - plt.plot(PSNR) - plt.show() - - -def main(): - - parser = argparse.ArgumentParser() - parser.add_argument('videopath', type=str, help='path of "normal" video') - - try: - args = parser.parse_args() - get_meanframe(args.videopath, args.videopath[:-4] + '_meanframe.jpg') - except: - parser.print_help() - sys.exit(0) - - -if __name__ == '__main__': - main() diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 741a672..0000000 --- a/requirements.txt +++ /dev/null @@ -1,7 +0,0 @@ -pyzmq==25.1.2 -opencv-python==4.9.0.80 -numpy==1.26.4 -ultralytics==8.1.18 -scikit-image==0.23.2 -matplotlib==3.8.3 -torch \ No newline at end of file diff --git a/run.py b/run.py deleted file mode 100644 index daf3acd..0000000 --- a/run.py +++ /dev/null @@ -1,17 +0,0 @@ -from multiprocessing import Process -import gui -import broker -from cameras import cam -from workers import worker - - -if __name__ == "__main__": - - broker_process = Process(target=broker.main) - worker_process = Process(target=worker.main) - gui_process = Process(target=gui.main) - - worker_process.start() - broker_process.start() - cam.run() # start camera Processes - gui_process.start() \ No newline at end of file diff --git a/src/broker.cpp b/src/broker.cpp new file mode 100644 index 0000000..1ae6f80 --- /dev/null +++ b/src/broker.cpp @@ -0,0 +1,51 @@ +#include "broker.hpp" + +#include + +#include + +#include "config.hpp" +#include "message.hpp" + +namespace va { + +void run_broker() { + const Config& cfg = config(); + + zmq::context_t context(1); + + zmq::socket_t subscriber(context, zmq::socket_type::sub); + for (const auto& s : cfg.sources) { + subscriber.connect(s.port); + } + subscriber.set(zmq::sockopt::subscribe, ""); + + std::vector workers; + workers.reserve(cfg.workers.size()); + for (const auto& w : cfg.workers) { + zmq::socket_t worker(context, zmq::socket_type::push); + worker.connect(w); + workers.push_back(std::move(worker)); + } + + while (true) { + zmq::message_t address; + zmq::message_t contents; + if (!subscriber.recv(address, zmq::recv_flags::none)) { + continue; + } + if (!subscriber.recv(contents, zmq::recv_flags::none)) { + continue; + } + + int idx = decode_address(to_string(address)); + if (idx < 0 || idx >= static_cast(workers.size())) { + continue; + } + + workers[idx].send(address, zmq::send_flags::sndmore); + workers[idx].send(contents, zmq::send_flags::none); + } +} + +} // namespace va diff --git a/src/broker.hpp b/src/broker.hpp new file mode 100644 index 0000000..0279c70 --- /dev/null +++ b/src/broker.hpp @@ -0,0 +1,9 @@ +#pragma once + +namespace va { + +// Subscribes to every source and routes each frame to the worker whose index +// matches the frame's source id. Runs until the process is terminated. +void run_broker(); + +} // namespace va diff --git a/src/camera.cpp b/src/camera.cpp new file mode 100644 index 0000000..5d41acd --- /dev/null +++ b/src/camera.cpp @@ -0,0 +1,69 @@ +#include "camera.hpp" + +#include +#include +#include +#include + +#include +#include + +#include "config.hpp" +#include "message.hpp" + +namespace va { + +namespace { + +// Streams a single video source to a ZeroMQ PUB socket. +void start_source(int id, const Source& source) { + zmq::context_t context(1); + zmq::socket_t publisher(context, zmq::socket_type::pub); + publisher.bind(source.port); + + cv::VideoCapture cap(source.path); + if (!cap.isOpened()) { + std::fprintf(stderr, "camera %d: unable to open %s\n", id, + source.path.c_str()); + return; + } + + int count = 0; + cv::Mat img; + while (true) { + bool success = cap.read(img); + + // Skip a few frames to behave more like a realtime source. + count += 5; + cap.set(cv::CAP_PROP_POS_FRAMES, count); + if (!success || img.empty()) { + break; + } + + std::vector buf; + cv::imencode(".jpg", img, buf); + + std::string address = encode_address(id); + publisher.send(zmq::buffer(address), zmq::send_flags::sndmore); + publisher.send(zmq::buffer(buf.data(), buf.size()), + zmq::send_flags::none); + + // Delay to avoid skipping frames on the subscriber side. + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } +} + +} // namespace + +std::vector run_cameras() { + const Config& cfg = config(); + std::vector threads; + for (std::size_t id = 0; id < cfg.sources.size(); ++id) { + threads.emplace_back(start_source, static_cast(id), + cfg.sources[id]); + std::printf("inited source thread\n"); + } + return threads; +} + +} // namespace va diff --git a/src/camera.hpp b/src/camera.hpp new file mode 100644 index 0000000..5a6378f --- /dev/null +++ b/src/camera.hpp @@ -0,0 +1,13 @@ +#pragma once + +#include +#include + +namespace va { + +// Publishes frames from every configured video source. Each source runs in its +// own thread, binding a ZeroMQ PUB socket and emitting [address, jpeg] frames. +// Returns the spawned threads so the caller can join them. +std::vector run_cameras(); + +} // namespace va diff --git a/src/config.cpp b/src/config.cpp new file mode 100644 index 0000000..8800fbf --- /dev/null +++ b/src/config.cpp @@ -0,0 +1,49 @@ +#include "config.hpp" + +#include +#include +#include + +#include + +namespace va { + +Config load_config(const std::string& path) { + std::ifstream f(path); + if (!f.is_open()) { + throw std::runtime_error("Unable to open config file: " + path); + } + + nlohmann::json j; + f >> j; + + Config cfg; + + for (const auto& s : j.at("sources")) { + Source src; + src.port = s.at("port").get(); + src.path = s.value("path", std::string{}); + src.meanframe = s.value("meanframe", std::string{}); + cfg.sources.push_back(std::move(src)); + } + + for (const auto& w : j.at("workers")) { + cfg.workers.push_back(w.get()); + } + + const auto& models = j.at("models"); + cfg.models.fire = models.value("fire", std::string{}); + cfg.models.forklift = models.value("forklift", std::string{}); + + cfg.gui = j.at("gui").get(); + cfg.psnr_threshold = j.value("psnr_threshold", 25); + + return cfg; +} + +const Config& config(const std::string& path) { + static Config cfg = load_config(path); + return cfg; +} + +} // namespace va diff --git a/src/config.hpp b/src/config.hpp new file mode 100644 index 0000000..5d90f78 --- /dev/null +++ b/src/config.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include +#include + +namespace va { + +struct Source { + std::string port; // ZeroMQ endpoint the source binds to (PUB) + std::string path; // video file / stream path + std::string meanframe; // path to the PSNR reference frame (.jpg) +}; + +struct Models { + std::string fire; + std::string forklift; +}; + +// Application configuration, loaded from config.json. +struct Config { + std::vector sources; + std::vector workers; // ZeroMQ endpoints workers bind to (PULL) + Models models; + std::string gui; // ZeroMQ endpoint the GUI binds to (PULL) + int psnr_threshold = 25; +}; + +// Load configuration from a JSON file. Throws std::runtime_error on failure. +Config load_config(const std::string& path); + +// Access the process-wide configuration, loading it on first use. +const Config& config(const std::string& path = "./config.json"); + +} // namespace va diff --git a/src/detector.cpp b/src/detector.cpp new file mode 100644 index 0000000..b99d02a --- /dev/null +++ b/src/detector.cpp @@ -0,0 +1,131 @@ +#include "detector.hpp" + +#include +#include + +#include +#include + +namespace va { + +YoloDetector::YoloDetector(const std::string& model_path, + std::vector class_names, + float conf_threshold, float nms_threshold, + int input_size) + : class_names_(std::move(class_names)), + conf_threshold_(conf_threshold), + nms_threshold_(nms_threshold), + input_size_(input_size) { + try { + net_ = cv::dnn::readNet(model_path); + net_.setPreferableBackend(cv::dnn::DNN_BACKEND_DEFAULT); + net_.setPreferableTarget(cv::dnn::DNN_TARGET_CPU); + } catch (const cv::Exception& e) { + std::fprintf(stderr, "failed to load model %s: %s\n", + model_path.c_str(), e.what()); + } +} + +std::vector YoloDetector::detect(const cv::Mat& image) const { + std::vector results; + if (net_.empty() || image.empty()) { + return results; + } + + cv::Mat blob = cv::dnn::blobFromImage( + image, 1.0 / 255.0, cv::Size(input_size_, input_size_), cv::Scalar(), + true, false); + net_.setInput(blob); + + std::vector outputs; + net_.forward(outputs, net_.getUnconnectedOutLayersNames()); + if (outputs.empty()) { + return results; + } + + // YOLOv8 ONNX output: shape [1, 4 + num_classes, num_boxes]. + cv::Mat out = outputs[0]; + if (out.dims == 3) { + out = out.reshape(1, out.size[1]); + } + cv::Mat preds = out.t(); // [num_boxes, 4 + num_classes] + + const int num_classes = preds.cols - 4; + const float x_scale = static_cast(image.cols) / input_size_; + const float y_scale = static_cast(image.rows) / input_size_; + + std::vector class_ids; + std::vector confidences; + std::vector boxes; + + for (int i = 0; i < preds.rows; ++i) { + const float* row = preds.ptr(i); + const float* scores = row + 4; + + int class_id = 0; + float max_score = scores[0]; + for (int c = 1; c < num_classes; ++c) { + if (scores[c] > max_score) { + max_score = scores[c]; + class_id = c; + } + } + if (max_score < conf_threshold_) { + continue; + } + + float cx = row[0]; + float cy = row[1]; + float w = row[2]; + float h = row[3]; + int left = static_cast((cx - w / 2.f) * x_scale); + int top = static_cast((cy - h / 2.f) * y_scale); + int width = static_cast(w * x_scale); + int height = static_cast(h * y_scale); + + class_ids.push_back(class_id); + confidences.push_back(max_score); + boxes.emplace_back(left, top, width, height); + } + + std::vector keep; + cv::dnn::NMSBoxes(boxes, confidences, conf_threshold_, nms_threshold_, keep); + + for (int idx : keep) { + Detection det; + det.confidence = confidences[idx]; + int class_id = class_ids[idx]; + if (class_id >= 0 && class_id < static_cast(class_names_.size())) { + det.name = class_names_[class_id]; + } else { + det.name = std::to_string(class_id); + } + const cv::Rect& r = boxes[idx]; + det.box.x1 = static_cast(r.x); + det.box.y1 = static_cast(r.y); + det.box.x2 = static_cast(r.x + r.width); + det.box.y2 = static_cast(r.y + r.height); + results.push_back(std::move(det)); + } + + return results; +} + +std::string detections_to_json(const std::vector& dets) { + nlohmann::json arr = nlohmann::json::array(); + for (const auto& d : dets) { + nlohmann::json j; + j["name"] = d.name; + j["confidence"] = d.confidence; + j["box"] = { + {"x1", d.box.x1}, + {"y1", d.box.y1}, + {"x2", d.box.x2}, + {"y2", d.box.y2}, + }; + arr.push_back(std::move(j)); + } + return arr.dump(); +} + +} // namespace va diff --git a/src/detector.hpp b/src/detector.hpp new file mode 100644 index 0000000..6404b40 --- /dev/null +++ b/src/detector.hpp @@ -0,0 +1,54 @@ +#pragma once + +#include +#include + +#include + +namespace va { + +struct Box { + float x1 = 0.f; + float y1 = 0.f; + float x2 = 0.f; + float y2 = 0.f; +}; + +struct Detection { + std::string name; + float confidence = 0.f; + Box box; +}; + +// YOLOv8 object detector backed by OpenCV's DNN module. +// +// The original project used Ultralytics' PyTorch (.pt) checkpoints. In C++ we +// run the exported ONNX models instead (export once with +// `yolo export model=forklift_8s.pt format=onnx`). The class names default to +// the label attached to the model; override them via the constructor. +class YoloDetector { +public: + YoloDetector(const std::string& model_path, + std::vector class_names, + float conf_threshold = 0.25f, float nms_threshold = 0.45f, + int input_size = 640); + + // Runs inference on a BGR image and returns the surviving detections. + std::vector detect(const cv::Mat& image) const; + + bool ready() const { return !net_.empty(); } + +private: + mutable cv::dnn::Net net_; + std::vector class_names_; + float conf_threshold_; + float nms_threshold_; + int input_size_; +}; + +// Serialises detections to the JSON string consumed by the GUI. The layout +// matches Ultralytics' Results.tojson(): a list of +// {name, class, confidence, box:{x1,y1,x2,y2}}. +std::string detections_to_json(const std::vector& dets); + +} // namespace va diff --git a/src/gui.cpp b/src/gui.cpp new file mode 100644 index 0000000..3dae3aa --- /dev/null +++ b/src/gui.cpp @@ -0,0 +1,138 @@ +#include "gui.hpp" + +#include +#include +#include + +#include +#include +#include + +#include "config.hpp" +#include "message.hpp" + +namespace va { + +namespace { + +// Assembles the source frames into a single canvas laid out on a grid. +class Canvas { +public: + Canvas() { + canvas_ = cv::Mat(height_, width_, CV_8UC3, cv::Scalar(40, 40, 40)); + } + + const cv::Mat& canvas() const { return canvas_; } + + void update_cell(int n, const cv::Mat& img) { + int col = n % grid_w_; + int row = n / grid_w_; + int w_pix = width_ / grid_w_; + int h_pix = height_ / grid_h_; + + int w_start = col * w_pix; + int h_start = row * h_pix; + if (h_start + h_pix > height_ || w_start + w_pix > width_) { + return; + } + + cv::Mat resized; + cv::resize(img, resized, cv::Size(w_pix, h_pix), 0, 0, cv::INTER_CUBIC); + resized.copyTo( + canvas_(cv::Rect(w_start, h_start, w_pix, h_pix))); + } + +private: + int width_ = 1920; + int height_ = 1080; + int grid_w_ = 2; + int grid_h_ = 2; + cv::Mat canvas_; +}; + +} // namespace + +void run_gui() { + const Config& cfg = config(); + + zmq::context_t context(1); + zmq::socket_t puller(context, zmq::socket_type::pull); + puller.bind(cfg.gui); + std::printf("GUI\n"); + + Canvas gui; + + while (true) { + zmq::message_t addr_msg; + zmq::message_t img_msg; + zmq::message_t dets_msg; + zmq::message_t psnr_msg; + if (!puller.recv(addr_msg, zmq::recv_flags::none)) continue; + if (!puller.recv(img_msg, zmq::recv_flags::none)) continue; + if (!puller.recv(dets_msg, zmq::recv_flags::none)) continue; + if (!puller.recv(psnr_msg, zmq::recv_flags::none)) continue; + + int addr = decode_address(to_string(addr_msg)); + std::string psnr = to_string(psnr_msg); + + std::vector buf( + static_cast(img_msg.data()), + static_cast(img_msg.data()) + img_msg.size()); + cv::Mat img = cv::imdecode(buf, cv::IMREAD_COLOR); + if (img.empty()) { + continue; + } + + bool is_fire = false; + try { + nlohmann::json dets = nlohmann::json::parse(to_string(dets_msg)); + for (const auto& d : dets) { + int x1 = static_cast(d.at("box").at("x1").get()); + int y1 = static_cast(d.at("box").at("y1").get()); + int x2 = static_cast(d.at("box").at("x2").get()); + int y2 = static_cast(d.at("box").at("y2").get()); + cv::rectangle(img, cv::Point(x1, y1), cv::Point(x2, y2), + cv::Scalar(255, 0, 0), 2); + std::string name = d.value("name", std::string{}); + if (name == "fire") { + is_fire = true; + } + cv::putText(img, name, cv::Point(x1, y1 - 5), + cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(255, 0, 0), 2, + cv::LINE_AA); + } + } catch (const nlohmann::json::exception& e) { + std::fprintf(stderr, "gui: bad detections json: %s\n", e.what()); + } + + cv::putText(img, "PSNR: " + psnr, cv::Point(50, 50), + cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255, 0), 2, + cv::LINE_AA); + + int psnr_val = 0; + try { + psnr_val = std::stoi(psnr); + } catch (...) { + psnr_val = 0; + } + if (psnr_val <= cfg.psnr_threshold) { + cv::putText(img, "WARNING!", cv::Point(250, 50), + cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 4, + cv::LINE_AA); + } + if (is_fire) { + cv::putText(img, "FIRE!", cv::Point(450, 50), + cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 255), 4, + cv::LINE_AA); + } + + gui.update_cell(addr, img); + cv::imshow("Image", gui.canvas()); + if ((cv::waitKey(10) & 0xFF) == 27) { + cv::destroyAllWindows(); + break; + } + } +} + +} // namespace va diff --git a/src/gui.hpp b/src/gui.hpp new file mode 100644 index 0000000..8a4cd96 --- /dev/null +++ b/src/gui.hpp @@ -0,0 +1,11 @@ +#pragma once + +namespace va { + +// Displays a grid of the latest annotated frames from every source. Pulls +// [address, jpeg, detections_json, psnr] messages from the GUI endpoint and +// renders bounding boxes, PSNR, WARNING and FIRE overlays. Blocks until the +// window is closed with ESC. +void run_gui(); + +} // namespace va diff --git a/src/main.cpp b/src/main.cpp new file mode 100644 index 0000000..041bd7e --- /dev/null +++ b/src/main.cpp @@ -0,0 +1,57 @@ +#include +#include +#include + +#include "broker.hpp" +#include "camera.hpp" +#include "config.hpp" +#include "gui.hpp" +#include "psnr.hpp" +#include "worker.hpp" + +namespace { + +int run_all() { + // Force config to load (and fail fast) before spawning anything. + va::config(); + + std::thread broker(va::run_broker); + std::vector workers = va::run_workers(); + std::vector cameras = va::run_cameras(); + + // The GUI runs on the main thread and controls the lifetime of the app. + va::run_gui(); + + broker.join(); + for (auto& t : workers) t.join(); + for (auto& t : cameras) t.join(); + return 0; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc >= 2 && std::strcmp(argv[1], "psnr") == 0) { + return va::psnr_tool_main(argc, argv); + } + if (argc >= 2 && std::strcmp(argv[1], "broker") == 0) { + va::run_broker(); + return 0; + } + if (argc >= 2 && std::strcmp(argv[1], "gui") == 0) { + va::run_gui(); + return 0; + } + if (argc >= 2 && std::strcmp(argv[1], "workers") == 0) { + auto workers = va::run_workers(); + for (auto& t : workers) t.join(); + return 0; + } + if (argc >= 2 && std::strcmp(argv[1], "cameras") == 0) { + auto cameras = va::run_cameras(); + for (auto& t : cameras) t.join(); + return 0; + } + + return run_all(); +} diff --git a/src/message.hpp b/src/message.hpp new file mode 100644 index 0000000..45d6ad2 --- /dev/null +++ b/src/message.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include +#include + +#include + +namespace va { + +// Address encoding mirrors the original Python implementation, where a source +// id is serialised as a 2-byte big-endian integer (id.to_bytes(2, 'big')) and +// decoded with int.from_bytes(address, 'big'). +inline std::string encode_address(int id) { + std::string s(2, '\0'); + s[0] = static_cast((id >> 8) & 0xFF); + s[1] = static_cast(id & 0xFF); + return s; +} + +inline int decode_address(const std::string& s) { + int value = 0; + for (unsigned char c : s) { + value = (value << 8) | c; + } + return value; +} + +inline std::string to_string(const zmq::message_t& msg) { + return std::string(static_cast(msg.data()), msg.size()); +} + +} // namespace va diff --git a/src/psnr.cpp b/src/psnr.cpp new file mode 100644 index 0000000..7933633 --- /dev/null +++ b/src/psnr.cpp @@ -0,0 +1,91 @@ +#include "psnr.hpp" + +#include + +#include + +namespace va { + +double psnr(const cv::Mat& reference, const cv::Mat& image) { + if (reference.empty() || image.empty() || + reference.size() != image.size() || + reference.type() != image.type()) { + return 0.0; + } + + cv::Mat diff; + cv::absdiff(reference, image, diff); + diff.convertTo(diff, CV_32F); + diff = diff.mul(diff); + + cv::Scalar s = cv::sum(diff); + double sse = s[0] + s[1] + s[2]; + + double total = static_cast(reference.total()) * reference.channels(); + double mse = sse / total; + if (mse <= 1e-10) { + return 100.0; // effectively identical + } + + return 10.0 * std::log10((255.0 * 255.0) / mse); +} + +bool build_meanframe(const std::string& video_path, + const std::string& meanframe_path) { + cv::VideoCapture cap(video_path); + if (!cap.isOpened()) { + std::fprintf(stderr, "unable to open video: %s\n", video_path.c_str()); + return false; + } + + cv::Mat img; + cv::Mat avg; + long frames = 0; + while (cap.read(img)) { + if (img.empty()) { + break; + } + cv::Mat imgf; + img.convertTo(imgf, CV_64F); + if (avg.empty()) { + avg = cv::Mat::zeros(imgf.size(), imgf.type()); + } + avg += imgf; + ++frames; + } + cap.release(); + + if (frames == 0 || avg.empty()) { + std::fprintf(stderr, "no frames read from: %s\n", video_path.c_str()); + return false; + } + + avg /= static_cast(frames); + cv::Mat out; + avg.convertTo(out, CV_8U); + return cv::imwrite(meanframe_path, out); +} + +int psnr_tool_main(int argc, char** argv) { + if (argc < 3) { + std::printf("usage: videoanalytics psnr \n"); + std::printf(" creates _meanframe.jpg\n"); + return 0; + } + + std::string video_path = argv[2]; + std::string out_path = video_path; + auto dot = out_path.find_last_of('.'); + if (dot != std::string::npos) { + out_path = out_path.substr(0, dot); + } + out_path += "_meanframe.jpg"; + + if (!build_meanframe(video_path, out_path)) { + return 1; + } + std::printf("meanframe written to %s\n", out_path.c_str()); + return 0; +} + +} // namespace va diff --git a/src/psnr.hpp b/src/psnr.hpp new file mode 100644 index 0000000..4142956 --- /dev/null +++ b/src/psnr.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include + +#include + +namespace va { + +// Peak signal-to-noise ratio between two same-size 8-bit images (in dB). +// Returns a large finite value when the images are identical. +double psnr(const cv::Mat& reference, const cv::Mat& image); + +// Builds the mean frame of a "normal" video (no emergency) and writes it as a +// JPEG. Mirrors the reference-frame generation of the original psnr.py tool. +// Returns true on success. +bool build_meanframe(const std::string& video_path, + const std::string& meanframe_path); + +// CLI entry point for the meanframe tool: `videoanalytics psnr `. +int psnr_tool_main(int argc, char** argv); + +} // namespace va diff --git a/src/worker.cpp b/src/worker.cpp new file mode 100644 index 0000000..9269b83 --- /dev/null +++ b/src/worker.cpp @@ -0,0 +1,98 @@ +#include "worker.hpp" + +#include +#include +#include +#include + +#include +#include + +#include "config.hpp" +#include "detector.hpp" +#include "message.hpp" +#include "psnr.hpp" + +namespace va { + +void worker_process(const std::string& pull_addr, + const std::string& push_addr) { + const Config& cfg = config(); + + // Default class names for the shipped models. Adjust to match the labels of + // your exported ONNX models if they differ. + YoloDetector forklift_model(cfg.models.forklift, {"forklift"}); + YoloDetector fire_model(cfg.models.fire, {"fire"}); + + zmq::context_t context(1); + zmq::socket_t puller(context, zmq::socket_type::pull); + puller.bind(pull_addr); + + zmq::socket_t pusher(context, zmq::socket_type::push); + pusher.connect(push_addr); + + while (true) { + zmq::message_t address; + zmq::message_t contents; + if (!puller.recv(address, zmq::recv_flags::none)) { + continue; + } + if (!puller.recv(contents, zmq::recv_flags::none)) { + continue; + } + + std::string addr_str = to_string(address); + int source_id = decode_address(addr_str); + + std::vector buf( + static_cast(contents.data()), + static_cast(contents.data()) + contents.size()); + cv::Mat frame = cv::imdecode(buf, cv::IMREAD_COLOR); + if (frame.empty()) { + continue; + } + + cv::Mat meanframe; + if (source_id >= 0 && source_id < static_cast(cfg.sources.size())) { + meanframe = cv::imread(cfg.sources[source_id].meanframe); + } + + // Run the two detectors and PSNR concurrently, mirroring the original + // threaded worker. + auto f_forklift = std::async(std::launch::async, [&] { + return forklift_model.detect(frame); + }); + auto f_fire = std::async(std::launch::async, [&] { + return fire_model.detect(frame); + }); + auto f_psnr = std::async(std::launch::async, [&] { + return meanframe.empty() ? 0.0 : psnr(meanframe, frame); + }); + + std::vector dets = f_forklift.get(); + std::vector fire = f_fire.get(); + double res_psnr = f_psnr.get(); + dets.insert(dets.end(), fire.begin(), fire.end()); + + std::string dets_json = detections_to_json(dets); + std::string psnr_str = std::to_string(static_cast(std::lround(res_psnr))); + + std::printf("PSNR %d %f\n", source_id, res_psnr); + + pusher.send(zmq::buffer(addr_str), zmq::send_flags::sndmore); + pusher.send(contents, zmq::send_flags::sndmore); + pusher.send(zmq::buffer(dets_json), zmq::send_flags::sndmore); + pusher.send(zmq::buffer(psnr_str), zmq::send_flags::none); + } +} + +std::vector run_workers() { + const Config& cfg = config(); + std::vector threads; + for (const auto& worker : cfg.workers) { + threads.emplace_back(worker_process, worker, cfg.gui); + } + return threads; +} + +} // namespace va diff --git a/src/worker.hpp b/src/worker.hpp new file mode 100644 index 0000000..166a24d --- /dev/null +++ b/src/worker.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include +#include +#include + +namespace va { + +// Runs a single worker: pulls frames from pull_addr, runs the fire/forklift +// detectors and PSNR, then forwards [address, jpeg, detections_json, psnr] to +// push_addr. Blocks until the process is terminated. +void worker_process(const std::string& pull_addr, const std::string& push_addr); + +// Spawns one worker thread per configured worker endpoint (all pushing to the +// GUI endpoint). Returns the spawned threads so the caller can join them. +std::vector run_workers(); + +} // namespace va diff --git a/workers/worker.py b/workers/worker.py deleted file mode 100644 index c120fb1..0000000 --- a/workers/worker.py +++ /dev/null @@ -1,102 +0,0 @@ -import zmq -import numpy as np -import cv2 -import torch -import json -from ultralytics import YOLO -from multiprocessing import Process -from threading import Thread -from skimage.metrics import peak_signal_noise_ratio -from config import config - - -class ThreadRet(Thread): - """Custom Thread class allowing to return values from threads""" - - def __init__(self, group=None, target=None, name=None, - args=(), kwargs={}, Verbose=None): - Thread.__init__(self, group, target, name, args, kwargs) - self._return = None - - def run(self): - if self._target is not None: - self._return = self._target(*self._args, **self._kwargs) - - def join(self, *args): - Thread.join(self, *args) - return self._return - - -def predict(model, image): - """get YOLO prediction - - Args: - model (YOLO): YOLOv8 model - image (MatLike): cv2 image - - Returns: - str: results of json-like string type - """ - result = model(image)[0] - return result.tojson() - - -def worker_process(pull_addr: str, push_addr: str): - """Make worker process - - Args: - pull_addr (str): address to get images from - push_addr (str): address to send images to - """ - - if torch.cuda.is_available(): - torch.cuda.set_device(0) - - forklift_model = YOLO(config['models']['forklift']) - fire_model = YOLO(config['models']['fire']) - - context = zmq.Context() - puller = context.socket(zmq.PULL) - puller.bind(pull_addr) - - pusher = context.socket(zmq.PUSH) - pusher.connect(push_addr) - - while True: - [address, contents] = puller.recv_multipart() - - encoded_frame = np.frombuffer(contents, dtype=np.uint8) - frame = cv2.imdecode(encoded_frame, cv2.IMREAD_COLOR) - meanframe = cv2.imread(config['sources'][int.from_bytes(address)]['meanframe']) - - t_forklift = ThreadRet(target=predict, args=(forklift_model, frame)) - t_fire = ThreadRet(target=predict, args=(fire_model, frame)) - t_psnr = ThreadRet(target=peak_signal_noise_ratio, args=(meanframe, frame)) - t_forklift.start() - t_fire.start() - t_psnr.start() - res_forklift = json.loads(t_forklift.join()) - res_fire = json.loads(t_fire.join()) - res_psnr = t_psnr.join() - - res = res_forklift + res_fire - res = str(res).replace('\'', '\"').encode() - - print(f'PSNR {int.from_bytes(address)} {res_psnr}') - - data = [address, contents, res, str(round(res_psnr)).encode()] - pusher.send_multipart(data) - - - # We never get here but clean up anyhow - pusher.close() - context.term() - - -def main(): - for worker in config['workers']: - Process(target=worker_process, args=(worker, config['gui'])).start() - - -if __name__ == "__main__": - main() \ No newline at end of file