diff --git a/include/IRToolTrack.h b/include/IRToolTrack.h index 1104736..2584d0f 100644 --- a/include/IRToolTrack.h +++ b/include/IRToolTrack.h @@ -75,6 +75,10 @@ class IRToolTracker std::atomic m_bShouldStop{false}; std::vector m_Tools; + // Guards m_Tools, m_ToolIndexMapping and the per-tool cur_transform/timestamp + // against concurrent access (GUI thread add/remove vs. UDP/CSV/GUI readers via + // GetToolTransform vs. the tracking thread writing results in UnionSegmentation). + std::mutex m_ToolsMutex; std::unique_ptr m_CurrentFrame; std::mutex m_MutexCurFrame; @@ -108,6 +112,9 @@ class IRToolTracker float m_fCalibrationSphereRadius = 6.5f; const int MAX_CALIBRATION_FRAMES = 100; + // Hard upper bound on how many spheres a calibration may resolve. A tool with + // more than this is rejected as unsuccessful rather than fed downstream. + static constexpr int MAX_CALIBRATION_SPHERES = 6; int NUM_CALIBRATION_SPHERES = 4; std::vector markerPositions; diff --git a/include/IRToolTracking.h b/include/IRToolTracking.h index a94a176..c6e85fc 100644 --- a/include/IRToolTracking.h +++ b/include/IRToolTracking.h @@ -99,6 +99,9 @@ class IRToolTracking { bool playFromFile = false; bool intrinsics_found = false; + // True only between a successful pipeline.start() and the matching shutdown(). + // Guards against rs2::pipeline::stop() throwing when start() never succeeded. + bool pipeline_started = false; std::vector deviceNames; diff --git a/include/ViewerWindow.h b/include/ViewerWindow.h index 43b7440..25d6ba3 100644 --- a/include/ViewerWindow.h +++ b/include/ViewerWindow.h @@ -24,6 +24,10 @@ class ViewerWindow { public: void Initialize(const std::string& file); void Shutdown(); + // Make destruction safe regardless of how Render()/main() unwound: stop every + // worker and join it, so no joinable std::thread is ever destroyed (which would + // call std::terminate()). + ~ViewerWindow() { Shutdown(); } bool IsTerminated() const { @@ -43,10 +47,24 @@ class ViewerWindow { bool Connect(NanoSocket& _socket, NanoAddress& address, const char* host, int port, bool& _connected); void Disconnect(NanoSocket& _socket, bool& _connected); + // Reference-counted nanosockets init/deinit so the library is initialized exactly + // once while one or more channels (send/receive) are active. + bool EnsureSocketInit(); + void ReleaseSocketInit(); + + // Thread-safe snapshot of the current tool names, so worker threads never index the + // GUI-owned `tools` vector while it is being resized/edited on the GUI thread. + std::vector SnapshotToolNames(); void GetSerialNumber(); Eigen::Matrix4f TrackingDataToMatrix(const TrackingData &data); + // Validation bounds. + static constexpr int MAX_TOOLS = 32; // upper bound for "Number of Tools" + static constexpr int MIN_SPHERES = 4; // lower bound for spheres per tool + static constexpr int MAX_SPHERES = 20; // safety cap for manual entry / ROM + static constexpr int MAX_CALIB_SPHERES = 6; // a calibration must not exceed this + std::atomic Terminated = ATOMIC_VAR_INIT(false); std::shared_ptr Thread; @@ -66,9 +84,15 @@ class ViewerWindow { GLFWwindow* window = nullptr; GLuint texture = 0, dtexture = 0; std::vector tools; + // Guards `tools` (resized/edited by the GUI thread, read by UDP/CSV worker threads). + std::mutex toolsMutex; std::map toolTransforms; std::mutex secondaryDataMutex; + // Set when a calibration produced an invalid result (too many spheres / NaN / + // no data); triggers the "Tool calibration unsuccessful" popup on the next frame. + bool showCalibrationError = false; + int numTools = 0; // Default number of tools std::string toolName = "Tool1"; // Default tool name int toolId = 0; // Default tool id @@ -96,8 +120,11 @@ class ViewerWindow { NanoSocket receiveSocket; NanoAddress sendAddress = {}; NanoAddress receiveAddress = {}; - bool m_connected; - bool m_receiveconnected; + bool m_connected = false; + bool m_receiveconnected = false; + // Reference count + guard for nanosockets_initialize()/deinitialize(). + std::mutex socketInitMutex; + int socketInitCount = 0; char ipAddress[16] = "127.0.0.1"; int m_port = 12345; int m_receiveport = 12345; diff --git a/src/IRToolTrack.cpp b/src/IRToolTrack.cpp index 3446bae..64e6e03 100644 --- a/src/IRToolTrack.cpp +++ b/src/IRToolTrack.cpp @@ -1,6 +1,7 @@ #include "IRToolTrack.h" #include "IRToolTracking.h" #include +#include IRToolTracker::~IRToolTracker() @@ -29,13 +30,17 @@ IRToolTracker::~IRToolTracker() cv::Mat IRToolTracker::GetToolTransform(std::string identifier) { + std::lock_guard lock(m_ToolsMutex); if (m_Tools.size() == 0 || m_ToolIndexMapping.count(identifier) == 0) return cv::Mat::zeros(8, 1, CV_32F); auto it = m_ToolIndexMapping.find(identifier); int index = it->second; - cv::Mat transform = m_Tools.at(index).cur_transform; + // Deep copy so callers (UDP/CSV/GUI threads) never alias the live buffer that + // the tracking thread keeps writing, and so the visibility reset below does not + // corrupt the tool's stored transform. + cv::Mat transform = m_Tools.at(index).cur_transform.clone(); if (m_lTrackedTimestamp != m_Tools.at(index).timestamp) { transform.at(7, 0) = 0.f; @@ -46,8 +51,9 @@ cv::Mat IRToolTracker::GetToolTransform(std::string identifier) void IRToolTracker::TrackTools() { + // m_bIsCurrentlyTracking is set true by StartTracking() before this thread starts, + // and cleared on exit below, so callers observe the tracking state synchronously. while (!m_bShouldStop) { - m_bIsCurrentlyTracking = true; std::unique_ptr rawFrame; { std::lock_guard lock(m_MutexCurFrame); @@ -88,9 +94,11 @@ void IRToolTracker::TrackTool(IRTrackedTool &tool, const ProcessedAHATFrame &fra const int radius_key = SphereRadiusKey(tool.sphere_radius); auto it_sides = frame.ordered_sides_per_mm.find(radius_key); - std::vector frame_ordered_sides = it_sides->second; - auto it_map = frame.map_per_mm.find(radius_key); + // No precomputed map for this tool's sphere radius in this frame: nothing to match. + if (it_sides == frame.ordered_sides_per_mm.end() || it_map == frame.map_per_mm.end()) + return; + std::vector frame_ordered_sides = it_sides->second; cv::Mat frame_map = it_map->second; //Find the set of eligible side to start with - aka sides that have similar length to first side of tool @@ -257,6 +265,12 @@ void IRToolTracker::UnionSegmentation(std::vector &raw_solu std::vector claimed_tools(num_tools, false); std::set used_spheres; + // Collect the winning poses first (MatchPointsKabsch is heavy and must run without + // the lock), then publish them all together under a single lock so readers + // (GetToolTransform) observe the per-tool transforms/timestamps and + // m_lTrackedTimestamp updated atomically — no torn read, no visibility glitch. + std::vector> committed; + for (const ToolResult ¤t : unique_solutions) { if (claimed_tools[current.tool_id]) @@ -275,26 +289,40 @@ void IRToolTracker::UnionSegmentation(std::vector &raw_solu cv::Mat result = MatchPointsKabsch(m_Tools[current.tool_id], frame, current.sphere_ids, current.occluded_nodes); if (result.at(7, 0) == 1.f) { - m_Tools.at(current.tool_id).cur_transform = result.clone(); - m_Tools.at(current.tool_id).timestamp = frame.timestamp; + committed.emplace_back(current.tool_id, result.clone()); } claimed_tools[current.tool_id] = true; used_spheres.insert(current.sphere_ids.begin(), current.sphere_ids.end()); } - m_lTrackedTimestamp = frame.timestamp; + + { + std::lock_guard lock(m_ToolsMutex); + for (auto &c : committed) + { + m_Tools.at(c.first).cur_transform = c.second; + m_Tools.at(c.first).timestamp = frame.timestamp; + } + m_lTrackedTimestamp = frame.timestamp; + } } cv::Mat IRToolTracker::MatchPointsKabsch(IRTrackedTool &tool, const ProcessedAHATFrame &frame, const std::vector &sphere_ids, const std::vector &occluded_nodes) { int num_points = static_cast(tool.num_spheres) - static_cast(occluded_nodes.size()); + // A degenerate match (no usable points) cannot yield a pose; report invisible. + if (num_points <= 0) + return cv::Mat::zeros(8, 1, CV_32F); + + auto it_spheres_xyz = frame.spheres_xyz_per_mm.find(SphereRadiusKey(tool.sphere_radius)); + if (it_spheres_xyz == frame.spheres_xyz_per_mm.end()) + return cv::Mat::zeros(8, 1, CV_32F); + cv::Mat3f frame_spheres_xyz = it_spheres_xyz->second; + cv::Mat p = cv::Mat(num_points, 3, CV_32F); cv::Mat q = cv::Mat(num_points, 3, CV_32F); cv::Vec3f p_center = cv::Vec3f(0.f); cv::Vec3f q_center = cv::Vec3f(0.f); - auto it_spheres_xyz = frame.spheres_xyz_per_mm.find(SphereRadiusKey(tool.sphere_radius)); - cv::Mat3f frame_spheres_xyz = it_spheres_xyz->second; - cv::Mat hololens_pose_mm = frame.device_pose.clone(); hololens_pose_mm.at(0, 3) = hololens_pose_mm.at(0, 3) * 1000.f; @@ -676,8 +704,11 @@ bool IRToolTracker::ProcessFrame(AHATFrame &rawFrame, ProcessedAHATFrame &result bool IRToolTracker::AddTool(cv::Mat3f spheres, float sphere_radius, std::string identifier, uint min_visible_spheres, float lowpass_rotation, float lowpass_position) { //Do we already have this tool? - if (m_ToolIndexMapping.count(identifier) > 0) - return false; + { + std::lock_guard lock(m_ToolsMutex); + if (m_ToolIndexMapping.count(identifier) > 0) + return false; + } bool restartTracking = false; if (m_bIsCurrentlyTracking) { @@ -706,8 +737,11 @@ bool IRToolTracker::AddTool(cv::Mat3f spheres, float sphere_radius, std::string tool.sphere_kalman_filters.push_back(IRToolKalmanFilter()); } //Add to map so we can find the tool with name - m_ToolIndexMapping.insert({ identifier, m_Tools.size() }); - m_Tools.push_back(tool); + { + std::lock_guard lock(m_ToolsMutex); + m_ToolIndexMapping.insert({ identifier, m_Tools.size() }); + m_Tools.push_back(tool); + } std::cout<<"Added "<second; + { + std::lock_guard lock(m_ToolsMutex); + if (m_ToolIndexMapping.count(identifier) == 0) + return false; + } bool restartTracking = false; if (m_bIsCurrentlyTracking) { @@ -732,15 +765,18 @@ bool IRToolTracker::RemoveTool(std::string identifier) StopTracking(); } - m_ToolIndexMapping.erase(identifier); - std::vector oldTools(m_Tools); - std::map oldMapping(m_ToolIndexMapping); - m_Tools.clear(); - m_ToolIndexMapping.clear(); - for (auto pair : oldMapping) { - m_ToolIndexMapping.insert({ pair.first, m_Tools.size() }); - m_Tools.push_back(oldTools.at(pair.second)); + std::lock_guard lock(m_ToolsMutex); + m_ToolIndexMapping.erase(identifier); + std::vector oldTools(m_Tools); + std::map oldMapping(m_ToolIndexMapping); + m_Tools.clear(); + m_ToolIndexMapping.clear(); + for (auto pair : oldMapping) + { + m_ToolIndexMapping.insert({ pair.first, m_Tools.size() }); + m_Tools.push_back(oldTools.at(pair.second)); + } } std::cout << "Removed " << identifier << std::endl; @@ -757,8 +793,11 @@ bool IRToolTracker::RemoveAllTools() if (m_bIsCurrentlyTracking) { StopTracking(); } - m_Tools.clear(); - m_ToolIndexMapping.clear(); + { + std::lock_guard lock(m_ToolsMutex); + m_Tools.clear(); + m_ToolIndexMapping.clear(); + } return true; } @@ -766,6 +805,11 @@ bool IRToolTracker::StartTracking() { if (m_bIsCurrentlyTracking || m_Tools.size() == 0) return false; m_bShouldStop = false; + // Mark tracking active BEFORE spawning the thread. Otherwise there is a window + // where IsTracking()==false while the thread is already reading m_Tools, during + // which AddTool/RemoveTool would skip StopTracking() and mutate m_Tools out from + // under the running tracking thread (use-after-free on vector reallocation). + m_bIsCurrentlyTracking = true; m_TrackingThread = std::make_shared(&IRToolTracker::TrackTools, this); return true; } @@ -794,6 +838,8 @@ Eigen::Vector3f cvVec3fToEigen(const cv::Vec3f &cvVec) } Eigen::Vector3f calculateStdDev(const std::vector& points, const Eigen::Vector3f& mean) { + if (points.empty()) + return Eigen::Vector3f::Constant(std::numeric_limits::quiet_NaN()); Eigen::Vector3f sumSq(0, 0, 0); for (const auto& point : points) { Eigen::Vector3f diff = point - mean; @@ -803,6 +849,10 @@ Eigen::Vector3f calculateStdDev(const std::vector& points, cons } Eigen::Vector3f calculateMean(const std::vector& points) { + // Empty input would divide by zero; return NaN so the calibration is flagged + // invalid by the caller instead of silently producing a (0,0,0) marker. + if (points.empty()) + return Eigen::Vector3f::Constant(std::numeric_limits::quiet_NaN()); Eigen::Vector3f sum(0, 0, 0); for (const auto& point : points) { sum += point; @@ -829,7 +879,10 @@ void IRToolTracker::CalibrateTool() std::vector processedFrames; processedFrames.reserve(MAX_CALIBRATION_FRAMES); - while (!m_bShouldStop) { + // Exit if asked to stop OR if the streaming pipeline terminated (e.g. no device + // connected / device disconnected), otherwise this loop would spin forever + // waiting for frames that will never arrive. + while (!m_bShouldStop && !m_pRealSenseToolTracking->IsTerminated()) { std::unique_ptr rawFrame; { std::lock_guard lock(m_MutexCurFrame); @@ -854,28 +907,64 @@ void IRToolTracker::CalibrateTool() } } + // markerPositions is the calibration result consumed by the GUI. Clear it up front + // so that every early-return below leaves an empty result, which the GUI treats as + // "calibration unsuccessful". + markerPositions.clear(); + + const int calib_radius_key = SphereRadiusKey(m_fCalibrationSphereRadius); + + // No frames captured (no device / disconnected / stopped before any frame): there + // is nothing to calibrate. Bail out instead of dereferencing processedFrames[0]. + if (processedFrames.empty()) + { + m_bIsCurrentlyCalibrating = false; + return; + } + // Perform calibration, loop through processedFrames std::vector> markerPoints; - // Define the number of calibration spheres based on size of frame_spheres_xyz in the first frame - const int calib_radius_key = SphereRadiusKey(m_fCalibrationSphereRadius); - NUM_CALIBRATION_SPHERES = processedFrames[0].spheres_xyz_per_mm.at(calib_radius_key).size().height; + // Define the number of calibration spheres from the first frame's detected blobs. + auto it_first = processedFrames[0].spheres_xyz_per_mm.find(calib_radius_key); + if (it_first == processedFrames[0].spheres_xyz_per_mm.end()) + { + m_bIsCurrentlyCalibrating = false; + return; + } + NUM_CALIBRATION_SPHERES = it_first->second.size().height; + + // A tool needs at least 3 spheres to define a coordinate frame, and we refuse to + // calibrate more than MAX_CALIBRATION_SPHERES. Either case is reported as failure. + if (NUM_CALIBRATION_SPHERES < 3 || NUM_CALIBRATION_SPHERES > MAX_CALIBRATION_SPHERES) + { + m_bIsCurrentlyCalibrating = false; + return; + } markerPoints.resize(NUM_CALIBRATION_SPHERES); for (ProcessedAHATFrame frame : processedFrames) { auto it_sides = frame.ordered_sides_per_mm.find(calib_radius_key); - std::vector frame_ordered_sides = it_sides->second; - auto it_spheres_xyz = frame.spheres_xyz_per_mm.find(calib_radius_key); + if (it_sides == frame.ordered_sides_per_mm.end() || + it_spheres_xyz == frame.spheres_xyz_per_mm.end()) + continue; + std::vector frame_ordered_sides = it_sides->second; cv::Mat3f frame_spheres_xyz = it_spheres_xyz->second; + // Only use frames whose detected-sphere count matches the reference frame; + // otherwise the fixed-size indexing below (up to NUM_CALIBRATION_SPHERES) would + // read past the end of a frame that detected fewer blobs. + if (frame_spheres_xyz.size().height != NUM_CALIBRATION_SPHERES || frame_ordered_sides.empty()) + continue; + //Side shortest = frame_ordered_sides.front(); Side longest = frame_ordered_sides.back(); if (std::isnan(longest.distance)) { - continue; + continue; } int tempMarker1ID = longest.id_from; @@ -902,6 +991,10 @@ void IRToolTracker::CalibrateTool() return a.second < b.second; }); + // Need at least one off-axis marker to resolve the third coordinate. + if (distanceToLine.empty()) + continue; + int marker3ID = distanceToLine.front().first; distanceToLine.erase(distanceToLine.begin()); diff --git a/src/IRToolTracking.cpp b/src/IRToolTracking.cpp index 8051108..5bfdfa8 100644 --- a/src/IRToolTracking.cpp +++ b/src/IRToolTracking.cpp @@ -23,22 +23,29 @@ void IRToolTracking::queryDevices() { } void IRToolTracking::initializeFromFile(const std::string& file) { - // Get the width and height of the frames from the file - rs2::pipeline pipe(ctx); - rs2::config cfg; - cfg.enable_device_from_file(file); - pipe.start(cfg); - auto profile = pipe.get_active_profile(); - auto depth_profile = profile.get_stream(RS2_STREAM_DEPTH).as(); - frame_width = depth_profile.width(); - frame_height = depth_profile.height(); - pipe.stop(); - - // Load the configuration from file - config.enable_device_from_file(file); - intrinsics_found = false; - Terminated = false; - playFromFile = true; + // Runs on the GUI thread; a missing/corrupt recording makes rs2 throw. Catch it so + // the app stays alive (Terminated=true makes the subsequent processStreams bail). + try { + // Get the width and height of the frames from the file + rs2::pipeline pipe(ctx); + rs2::config cfg; + cfg.enable_device_from_file(file); + pipe.start(cfg); + auto profile = pipe.get_active_profile(); + auto depth_profile = profile.get_stream(RS2_STREAM_DEPTH).as(); + frame_width = depth_profile.width(); + frame_height = depth_profile.height(); + pipe.stop(); + + // Load the configuration from file + config.enable_device_from_file(file); + intrinsics_found = false; + Terminated = false; + playFromFile = true; + } catch (const std::exception &e) { + std::cerr << "Failed to open recording '" << file << "': " << e.what() << std::endl; + Terminated = true; + } } void IRToolTracking::initialize(int index, int width, int height) { @@ -49,33 +56,40 @@ void IRToolTracking::initialize(int index, int width, int height) { return; } - dev = devices[index]; - //dev.hardware_reset(); - std::string model_name = dev.get_info(RS2_CAMERA_INFO_NAME); - if (model_name == "Intel RealSense D415") { - irThreshold = 100; - minSize = 10; - maxSize = 300; - } - else if (model_name == "Intel RealSense D435") { - irThreshold = 180; - minSize = 10; - maxSize = 370; - } - SetThreshold(irThreshold); - SetMinMaxSize(minSize, maxSize); + // Querying/configuring a device handle throws rs2::error if the camera was + // unplugged after enumeration. This runs on the GUI thread, so catch it. + try { + dev = devices[index]; + //dev.hardware_reset(); + std::string model_name = dev.get_info(RS2_CAMERA_INFO_NAME); + if (model_name == "Intel RealSense D415") { + irThreshold = 100; + minSize = 10; + maxSize = 300; + } + else if (model_name == "Intel RealSense D435") { + irThreshold = 180; + minSize = 10; + maxSize = 370; + } + SetThreshold(irThreshold); + SetMinMaxSize(minSize, maxSize); - frame_width = width; - frame_height = height; - intrinsics_found = false; - // Configure the RealSense pipeline + frame_width = width; + frame_height = height; + intrinsics_found = false; + // Configure the RealSense pipeline - config.enable_device(dev.get_info(RS2_CAMERA_INFO_SERIAL_NUMBER)); + config.enable_device(dev.get_info(RS2_CAMERA_INFO_SERIAL_NUMBER)); - config.enable_stream(RS2_STREAM_INFRARED, 1, frame_width, frame_height, RS2_FORMAT_Y8, 90); - config.enable_stream(RS2_STREAM_DEPTH, frame_width, frame_height, RS2_FORMAT_Z16, 90); + config.enable_stream(RS2_STREAM_INFRARED, 1, frame_width, frame_height, RS2_FORMAT_Y8, 90); + config.enable_stream(RS2_STREAM_DEPTH, frame_width, frame_height, RS2_FORMAT_Z16, 90); - Terminated = false; + Terminated = false; + } catch (const std::exception &e) { + std::cerr << "Failed to initialize device: " << e.what() << std::endl; + Terminated = true; + } } void IRToolTracking::StartToolCalibration() @@ -148,16 +162,23 @@ void IRToolTracking::processStreams() { // can join us cleanly instead of having the whole process killed. try { profile = pipeline.start(config); + pipeline_started = true; } catch (const rs2::error &e) { std::cerr << "Error during RealSense pipeline start: " << e.what() << std::endl; Terminated = true; return; } + + // Everything below can throw rs2::error if the device disconnects mid-stream (or + // OpenCV throws on a malformed frame). This runs on a worker thread, so an escaping + // exception would call std::terminate() and kill the whole app. Catch it, flag the + // session as terminated, and return so the GUI thread can join and shut down cleanly. + try { // Warm up the device - for (int i = 0; i < 50; i++) { + for (int i = 0; i < 50 && !Terminated; i++) { pipeline.wait_for_frames(); } - + // if not on Mac #if !defined(__APPLE__) { @@ -215,12 +236,25 @@ void IRToolTracking::processStreams() { } } } + } catch (const std::exception &e) { + std::cerr << "RealSense streaming stopped: " << e.what() << std::endl; + Terminated = true; + } } void IRToolTracking::shutdown() { - // Clean up resources as necessary - pipeline.stop(); - config.disable_all_streams(); + // Clean up resources as necessary. stop() throws if the pipeline was never + // started (e.g. start() failed), so only stop a pipeline we actually started. + // Guard the whole teardown so a disconnected device can't throw out of here. + try { + if (pipeline_started) { + pipeline.stop(); + } + config.disable_all_streams(); + } catch (const std::exception &e) { + std::cerr << "Error during RealSense shutdown: " << e.what() << std::endl; + } + pipeline_started = false; trackingFrame.release(); depthFrame.release(); } diff --git a/src/ViewerWindow.cpp b/src/ViewerWindow.cpp index b0b1f83..718168f 100644 --- a/src/ViewerWindow.cpp +++ b/src/ViewerWindow.cpp @@ -6,6 +6,8 @@ #include #include #include +#include +#include #include #include #include @@ -70,30 +72,46 @@ bool ViewerWindow::LoadToolDefinition() { if (entry.path().extension() == ".json") { - std::ifstream file(entry.path()); - if (!file) + // Parse defensively: a malformed or hand-edited file must be skipped, not + // crash startup. nlohmann throws on parse errors / missing keys / type + // mismatches; catch per-file so one bad file doesn't take down the rest. + try { - std::cerr << "Failed to open " << entry.path() << std::endl; - continue; - } + std::ifstream file(entry.path()); + if (!file) + { + std::cerr << "Failed to open " << entry.path() << std::endl; + continue; + } - nlohmann::json toolJson; - file >> toolJson; + nlohmann::json toolJson; + file >> toolJson; - Tool tool; - tool.numSpheres = toolJson["numSpheres"].get(); - tool.spherePositions = toolJson["spherePositions"].get>(); - tool.sphereRadius = toolJson["sphereRadius"].get(); - tool.toolName = toolJson["toolName"].get(); + Tool tool; + tool.numSpheres = toolJson.at("numSpheres").get(); + tool.spherePositions = toolJson.at("spherePositions").get>(); + tool.sphereRadius = toolJson.at("sphereRadius").get(); + tool.toolName = toolJson.at("toolName").get(); - if (!file) + // Reject inconsistent definitions (huge/negative counts, mismatched + // arrays, non-finite radius) before they reach any resize/indexing. + if (tool.numSpheres < MIN_SPHERES || tool.numSpheres > MAX_SPHERES || + tool.spherePositions.size() != static_cast(tool.numSpheres) * 3 || + !std::isfinite(tool.sphereRadius) || tool.sphereRadius <= 0.f) + { + std::cerr << "Skipping invalid tool definition: " << entry.path() << std::endl; + continue; + } + + std::cout << "Loaded tool definition: " << tool.toolName << std::endl; + tools.push_back(std::move(tool)); + numTools += 1; + } + catch (const std::exception &e) { - std::cerr << "Failed to read data from " << entry.path() << std::endl; + std::cerr << "Failed to parse " << entry.path() << ": " << e.what() << std::endl; continue; } - std::cout << "Loaded tool definition: " << tool.toolName << std::endl; - tools.push_back(std::move(tool)); - numTools += 1; } } @@ -188,24 +206,41 @@ Eigen::Matrix4f ViewerWindow::TrackingDataToMatrix(const TrackingData& data) return matrix; } -bool ViewerWindow::Connect(NanoSocket& _socket, NanoAddress& address, const char *host, int port, bool& _connected) { - if (udpEnabled != multiEnabled) +bool ViewerWindow::EnsureSocketInit() +{ + std::lock_guard lock(socketInitMutex); + if (socketInitCount == 0) { - if (nanosockets_initialize()) + if (nanosockets_initialize()) // non-zero == failure { - std::cerr << "Error initializing socket library" << std::endl; return false; } } + ++socketInitCount; + return true; +} + +void ViewerWindow::ReleaseSocketInit() +{ + std::lock_guard lock(socketInitMutex); + if (socketInitCount > 0 && --socketInitCount == 0) + { + nanosockets_deinitialize(); + } +} + +bool ViewerWindow::Connect(NanoSocket& _socket, NanoAddress& address, const char *host, int port, bool& _connected) { + if (!EnsureSocketInit()) + { + std::cerr << "Error initializing socket library" << std::endl; + return false; + } _socket = nanosockets_create(1024, 1024); if (_socket < 0) { std::cerr << "Failed to create a socket." << std::endl; - if (!multiEnabled && !udpEnabled) - { - nanosockets_deinitialize(); - } + ReleaseSocketInit(); return false; } @@ -217,10 +252,7 @@ bool ViewerWindow::Connect(NanoSocket& _socket, NanoAddress& address, const char { std::cerr<<"Error setting default address"< ViewerWindow::SnapshotToolNames() +{ + std::lock_guard lock(toolsMutex); + std::vector names; + names.reserve(tools.size()); + for (const auto &t : tools) + { + names.push_back(t.toolName); + } + return names; +} + void ViewerWindow::UdpReceiveThreadFunction() { if (!Connect(receiveSocket, receiveAddress, "0.0.0.0", m_receiveport, m_receiveconnected)) @@ -292,7 +330,10 @@ void ViewerWindow::UdpReceiveThreadFunction() } NanoAddress sender; - if (nanosockets_receive(receiveSocket, &sender, buffer.data(), buffer.size()) > 0) + const int received = nanosockets_receive(receiveSocket, &sender, buffer.data(), buffer.size()); + // Only accept a datagram that is exactly one TrackingData; a truncated or + // oversized packet must not be reinterpreted as a valid struct. + if (received == static_cast(sizeof(TrackingData))) { TrackingData data; std::memcpy(&data, buffer.data(), sizeof(TrackingData)); @@ -304,12 +345,13 @@ void ViewerWindow::UdpReceiveThreadFunction() if (tracker.IsTrackingTools()) { - for (size_t id = 0; id < tools.size(); ++id) + std::vector toolNames = SnapshotToolNames(); + for (size_t id = 0; id < toolNames.size(); ++id) { if (data.toolId == static_cast(id + 1)) { // Always update extrinsics in case camera is moved - std::vector tool_transform = tracker.GetToolTransform(tools[id].toolName); + std::vector tool_transform = tracker.GetToolTransform(toolNames[id]); if (!tool_transform.empty() && !std::isnan(tool_transform[0]) && tool_transform[7] != 0) { // tool_transform to matrix @@ -357,10 +399,10 @@ void ViewerWindow::UdpThreadFunction() { if (tracker.IsTrackingTools()) { - for (size_t id = 0; id < tools.size(); ++id) + std::vector toolNames = SnapshotToolNames(); + for (size_t id = 0; id < toolNames.size(); ++id) { - const auto &tool = tools[id]; - std::vector tool_transform = tracker.GetToolTransform(tool.toolName); + std::vector tool_transform = tracker.GetToolTransform(toolNames[id]); bool validPrimaryData = !tool_transform.empty() && !std::isnan(tool_transform[0]) && tool_transform[7] != 0; @@ -375,14 +417,16 @@ void ViewerWindow::UdpThreadFunction() tool_transform = { secondaryData(0, 3), secondaryData(1, 3), secondaryData(2, 3), q.x(), q.y(), q.z(), q.w() }; - toolTransforms.clear(); + // Consume only this tool's entry, not the whole map. + toolTransforms.erase(id); validSecondaryData = true; } } if (validPrimaryData || validSecondaryData) { - TrackingData data; + // Value-initialize so no uninitialized padding goes on the wire. + TrackingData data{}; std::copy(tool_transform.begin(), tool_transform.begin() + 3, data.position); std::copy(tool_transform.begin() + 3, tool_transform.end(), data.quaternion); data.toolId = static_cast(id) + 1; @@ -457,10 +501,10 @@ void ViewerWindow::WriteToCSV() break; // Exit loop } } - for (size_t id = 0; id < tools.size(); ++id) + std::vector toolNames = SnapshotToolNames(); + for (size_t id = 0; id < toolNames.size(); ++id) { - const auto& tool = tools[id]; - std::vector tool_transform = tracker.GetToolTransform(tool.toolName); + std::vector tool_transform = tracker.GetToolTransform(toolNames[id]); if (!tool_transform.empty() && !std::isnan(tool_transform[0]) && tool_transform[7] != 0) { double unixTimestamp = GetCurrentUnixTimestamp(); @@ -577,8 +621,14 @@ void ViewerWindow::Render() { ImGui::InputInt("Number of Tools", &numTools); ImGui::End(); - // Adjust the size of the tools vector based on numTools - numTools = std::max(numTools, 1); + // Hold toolsMutex while we resize/edit `tools` so the UDP/CSV worker threads + // never index it mid-modification. Released right after the calibration-result + // block below. + std::unique_lock toolsLock(toolsMutex); + + // Adjust the size of the tools vector based on numTools. Clamp to a sane range: + // an unbounded value from the InputInt would make tools.resize() allocate wildly. + numTools = std::clamp(numTools, 1, MAX_TOOLS); if (static_cast(tools.size()) > numTools) { // Unregister any tools that are about to be dropped so the tracker @@ -617,8 +667,8 @@ void ViewerWindow::Render() { tmpName[MAX_TOOL_NAME_LENGTH - 1] = '\0'; // Ensure null-termination ImGui::InputText(("Tool Name##" + std::to_string(toolIdx)).c_str(), tmpName, MAX_TOOL_NAME_LENGTH); tools[toolIdx].toolName = tmpName; - tools[toolIdx].numSpheres = std::max(tools[toolIdx].numSpheres, 4); - tools[toolIdx].spherePositions.resize(tools[toolIdx].numSpheres * 3); + tools[toolIdx].numSpheres = std::clamp(tools[toolIdx].numSpheres, MIN_SPHERES, MAX_SPHERES); + tools[toolIdx].spherePositions.resize(static_cast(tools[toolIdx].numSpheres) * 3); for (int i = 0; i < tools[toolIdx].numSpheres; ++i) { @@ -632,9 +682,12 @@ void ViewerWindow::Render() { if (ImGui::Button(("Add Tool Definition##" + std::to_string(toolIdx)).c_str())) { - tracker.AddToolDefinition(tools[toolIdx].numSpheres, tools[toolIdx].spherePositions, tools[toolIdx].sphereRadius, tools[toolIdx].toolName); - SaveToolDefinition(tools[toolIdx]); - tools[toolIdx].isAdded = true; + // Only persist/mark added if the tracker actually accepted it. + if (tracker.AddToolDefinition(tools[toolIdx].numSpheres, tools[toolIdx].spherePositions, tools[toolIdx].sphereRadius, tools[toolIdx].toolName)) + { + SaveToolDefinition(tools[toolIdx]); + tools[toolIdx].isAdded = true; + } } } else @@ -649,12 +702,24 @@ void ViewerWindow::Render() { } if (!tools[toolIdx].isAdded) - { + { ImGui::SameLine(); + // Calibration needs a live stream (a selected device or a recording). + // Without one, the old code still spawned worker threads and left the + // app wedged "calibrating" — and a second click overwrote a joinable + // std::thread, aborting the process. Also block re-entry while any + // session is already running. + const bool canCalibrate = (selectedDeviceIndex != -1 || recordedFile != "") && + !calibrationInitiated && !tracker.IsTrackingTools() && + !tracker.IsCalibratingTool(); + ImGui::BeginDisabled(!canCalibrate); if (ImGui::Button(("Calibrate Tool##" + std::to_string(toolIdx)).c_str())) { std::cout << "Calibrating tool " << toolIdx + 1 << std::endl; + // Ensure any previous streaming thread is fully joined before we + // replace the handle (destroying a joinable std::thread aborts). + JoinThread(processingThread); if (recordedFile != "") { tracker.initializeFromFile(recordedFile); @@ -669,6 +734,7 @@ void ViewerWindow::Render() { processingThread = std::make_shared([this]() { this->tracker.processStreams(); }); } + ImGui::EndDisabled(); ImGui::SameLine(); if (ImGui::Button(("Load ROM##" + std::to_string(toolIdx)).c_str())) @@ -679,8 +745,22 @@ void ViewerWindow::Render() { nfdresult_t result = NFD_OpenDialog(&path, filterItem, 1, NULL); if (result == NFD_OKAY && path) { ROMParser parser(path); - tools[toolIdx].numSpheres = parser.GetNumMarkers(); - tools[toolIdx].spherePositions = parser.GetMarkerPositions(); + const int romMarkers = parser.GetNumMarkers(); + const std::vector& romPositions = parser.GetMarkerPositions(); + // Only accept a ROM whose marker count is in range and whose + // position array matches; otherwise the next-frame resize and + // indexing would be inconsistent or oversized. + if (romMarkers >= MIN_SPHERES && romMarkers <= MAX_SPHERES && + romPositions.size() == static_cast(romMarkers) * 3) + { + tools[toolIdx].numSpheres = romMarkers; + tools[toolIdx].spherePositions = romPositions; + } + else + { + std::cerr << "Ignoring ROM with unsupported marker count: " + << romMarkers << std::endl; + } NFD_FreePath(path); } NFD_Quit(); @@ -691,14 +771,47 @@ void ViewerWindow::Render() { ImGui::End(); - if (!tracker.IsCalibratingTool() && calibrationInitiated) + const bool finishCalibration = (!tracker.IsCalibratingTool() && calibrationInitiated); + if (finishCalibration) { std::vector tool_definition = tracker.GetToolDefinition(); - if (!tool_definition.empty() && !std::isnan(tool_definition[0]) && tool_definition[0] != -1) + + // Validate before accepting: a whole number of spheres in + // [MIN_SPHERES, MAX_CALIB_SPHERES], every coordinate finite (no NaN/inf), + // and not the "no result" sentinel. Anything else means calibration failed. + // Note the lower bound matches MIN_SPHERES so an accepted count is never + // bumped (and corrupted) by the per-frame sphere-count clamp above. + bool validCalibration = false; + const int numFloats = static_cast(tool_definition.size()); + if (numFloats > 0 && numFloats % 3 == 0 && tool_definition[0] != -1) { - tools[toolId].numSpheres = tool_definition.size() / 3; + const int sphereCount = numFloats / 3; + if (sphereCount >= MIN_SPHERES && sphereCount <= MAX_CALIB_SPHERES) + { + validCalibration = std::all_of(tool_definition.begin(), tool_definition.end(), + [](float v) { return std::isfinite(v); }); + } + } + + if (validCalibration && toolId >= 0 && toolId < static_cast(tools.size())) + { + tools[toolId].numSpheres = numFloats / 3; tools[toolId].spherePositions = tool_definition; } + else + { + // Discard the result and prompt the user to recalibrate. + showCalibrationError = true; + } + } + + // Done mutating `tools` for this frame; release before any blocking teardown + // (JoinThread can block inside wait_for_frames up to the frame timeout, and we + // must not stall the UDP/CSV workers on toolsMutex while it does). + toolsLock.unlock(); + + if (finishCalibration) + { tracker.StopToolCalibration(); JoinThread(processingThread); tracker.shutdown(); @@ -711,6 +824,9 @@ void ViewerWindow::Render() { ImGui::Begin("Tracking Control", nullptr, overlayFlags); if (!tracker.IsTrackingTools()) { if (ImGui::Button("Start Tracking")) { + // Make sure any previous streaming thread is fully joined before we + // replace the handle (destroying a joinable std::thread aborts). + JoinThread(processingThread); if (recordedFile != "") { tracker.initializeFromFile(recordedFile); @@ -773,8 +889,10 @@ void ViewerWindow::Render() { ImGui::SetNextItemWidth(110); ImGui::InputText("IP Address", ipAddress, sizeof(ipAddress)); ImGui::SameLine(); - ImGui::SetNextItemWidth(50); + ImGui::SetNextItemWidth(50); ImGui::InputInt("Port", &m_port, 0, 0, ImGuiInputTextFlags_CharsDecimal); + // Keep within the valid UDP port range; the value is cast to uint16_t later. + m_port = std::clamp(m_port, 1, 65535); ImGui::SetNextItemWidth(110); ImGui::InputInt("Frequency", &frequency); ImGui::SameLine(); @@ -817,6 +935,7 @@ void ViewerWindow::Render() { ImGui::SameLine(); ImGui::SetNextItemWidth(50); ImGui::InputInt("Receive Port", &m_receiveport, 0, 0, ImGuiInputTextFlags_CharsDecimal); + m_receiveport = std::clamp(m_receiveport, 1, 65535); ImGui::End(); // GUI for CSV recording @@ -918,7 +1037,8 @@ void ViewerWindow::Render() { tool_transform = { secondaryData(0, 3), secondaryData(1, 3), secondaryData(2, 3), q.x(), q.y(), q.z(), q.w() }; - toolTransforms.clear(); + // Consume only this tool's entry, not the whole map. + toolTransforms.erase(i); validSecondaryData = true; } } @@ -953,6 +1073,29 @@ void ViewerWindow::Render() { ImGui::End(); } + // Calibration failure popup: shown when a calibration produced an invalid + // result (too many spheres, NaN/inf positions, or no data captured). + if (showCalibrationError) + { + ImGui::OpenPopup("Calibration Failed"); + showCalibrationError = false; + } + { + ImVec2 center = ImGui::GetMainViewport()->GetCenter(); + ImGui::SetNextWindowPos(center, ImGuiCond_Appearing, ImVec2(0.5f, 0.5f)); + } + if (ImGui::BeginPopupModal("Calibration Failed", nullptr, ImGuiWindowFlags_AlwaysAutoResize)) + { + ImGui::Text("Tool calibration unsuccessful."); + ImGui::Text("Please recalibrate the tool."); + ImGui::Separator(); + if (ImGui::Button("OK", ImVec2(120, 0))) + { + ImGui::CloseCurrentPopup(); + } + ImGui::EndPopup(); + } + ImGui::Render(); ImGui_ImplOpenGL3_RenderDrawData( ImGui::GetDrawData() ); @@ -979,8 +1122,6 @@ void ViewerWindow::Shutdown() { multiEnabled = false; csvEnabled = false; - const bool pipelineRunning = tracker.IsTrackingTools() || tracker.IsCalibratingTool(); - tracker.StopToolTracking(); tracker.StopToolCalibration(); JoinThread(processingThread); @@ -988,6 +1129,8 @@ void ViewerWindow::Shutdown() { JoinThread(udpReceiveThread); JoinThread(csvThread); - if (pipelineRunning) - tracker.shutdown(); + // shutdown() is idempotent and guarded by pipeline_started, so always call it. + // (Gating on IsTracking/IsCalibrating could miss the window where calibration + // already finished but its completion teardown hadn't run, leaking the pipeline.) + tracker.shutdown(); } \ No newline at end of file