diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml new file mode 100644 index 0000000..10c4bac --- /dev/null +++ b/.github/workflows/build.yml @@ -0,0 +1,127 @@ +name: Build + +on: + push: + branches: [main, dev] + pull_request: + branches: [main, dev] + workflow_dispatch: + +# Cancel an in-flight run when a newer commit lands on the same ref. +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + build: + name: ${{ matrix.label }} + runs-on: ${{ matrix.os }} + timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + include: + - os: ubuntu-22.04 + label: linux-x86_64 + - os: windows-2022 + label: windows-x86_64 + + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + submodules: recursive + + # ------------------------------------------------------------------- + # Linux: Microsoft "prod" apt repo for the Azure Kinect SDK (libk4a + + # k4arecord + the dlopen'd depth engine via k4a-tools), plus OpenCV and + # the GLFW/GTK system deps. + # + # The k4a packages are published only under the Ubuntu 18.04 ("bionic") + # prod feed; they install cleanly on 22.04 by adding that feed and + # pointing the suite at "bionic" (the feed has no "jammy" suite, so + # $(lsb_release -cs) would 404 here). The packages carry an interactive + # EULA, so we preseed debconf before apt runs. + # ------------------------------------------------------------------- + - name: Install Linux dependencies + if: runner.os == 'Linux' + env: + DEBIAN_FRONTEND: noninteractive + ACCEPT_EULA: 'Y' + run: | + set -euxo pipefail + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + ca-certificates curl gnupg lsb-release \ + build-essential cmake ninja-build pkg-config \ + libopencv-dev \ + libgl1-mesa-dev libglu1-mesa-dev \ + libxinerama-dev libxcursor-dev libxi-dev libxrandr-dev \ + libwayland-dev libxkbcommon-dev \ + libusb-1.0-0-dev libudev-dev \ + libgtk-3-dev + + # Microsoft prod apt repo (18.04 / bionic feed) for the Azure Kinect SDK. + # Modern signed-by keyring rather than the deprecated apt-key. + sudo install -m 0755 -d /etc/apt/keyrings + curl -fsSL https://packages.microsoft.com/keys/microsoft.asc \ + | sudo gpg --dearmor -o /etc/apt/keyrings/microsoft.gpg + sudo chmod 0644 /etc/apt/keyrings/microsoft.gpg + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/microsoft.gpg] https://packages.microsoft.com/ubuntu/18.04/prod bionic main" \ + | sudo tee /etc/apt/sources.list.d/microsoft-prod.list + sudo apt-get update + + # Preseed the EULA so libk4a / k4a-tools install non-interactively. + # Hash is the published "accepted-eula-hash" for the 1.4 packages; if + # it is wrong the preseed is silently ignored and apt blocks on the + # prompt until the job times out. + echo 'libk4a1.4 libk4a1.4/accepted-eula-hash string 0f5d5c5de396e4fee4c0753a21fee0c1ed726cf0316204edda484f08cb266d76' \ + | sudo debconf-set-selections + echo 'k4a-tools k4a-tools/accepted-eula-hash string 0f5d5c5de396e4fee4c0753a21fee0c1ed726cf0316204edda484f08cb266d76' \ + | sudo debconf-set-selections + + # libk4a1.4 ships libk4a.so + libk4arecord.so + the dlopen'd + # libdepthengine.so.2.0; -dev adds the headers + CMake/pkg-config that + # find_package(k4a) needs; k4a-tools is belt-and-suspenders for the + # depth engine. Pin matching versions. + sudo apt-get install -y --no-install-recommends \ + libk4a1.4=1.4.2 libk4a1.4-dev=1.4.2 k4a-tools=1.4.2 + + # ------------------------------------------------------------------- + # Windows: chocolatey for OpenCV only. The Azure Kinect SDK is VENDORED + # in the repo (extern/k4a/win32) and resolved by CMake from there, so + # there is NO k4a install step. CMake's POST_BUILD copies the 3 k4a DLLs + # next to the exe. + # ------------------------------------------------------------------- + - name: Install Windows dependencies + if: runner.os == 'Windows' + timeout-minutes: 25 + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + + Write-Host "Installing OpenCV via chocolatey..." + choco install opencv --version=4.10.0 -y --no-progress + $opencvDir = "C:\tools\opencv\build" + if (-not (Test-Path "$opencvDir\OpenCVConfig.cmake")) { + $opencvDir = Get-ChildItem -Path "C:\tools\opencv" -Recurse -Filter "OpenCVConfig.cmake" | Select-Object -First 1 -ExpandProperty Directory | Select-Object -ExpandProperty FullName + } + "OpenCV_DIR=$opencvDir" | Out-File -FilePath $env:GITHUB_ENV -Append + "$opencvDir\x64\vc16\bin" | Out-File -FilePath $env:GITHUB_PATH -Append + + # ------------------------------------------------------------------- + # Configure + build. Windows is pinned to the x64 VS generator so we + # never accidentally pick up a 32-bit toolchain (the OpenCV and vendored + # k4a libs are x64-only). /MT linkage on Windows + prebuilt OpenCV (/MD) + # relies on the /FORCE linker flag already set in CMakeLists. + # ------------------------------------------------------------------- + - name: Configure (Windows x64) + if: runner.os == 'Windows' + run: cmake -S . -B build -G "Visual Studio 17 2022" -A x64 -DCMAKE_BUILD_TYPE=Release + + - name: Configure (Linux) + if: runner.os == 'Linux' + run: cmake -S . -B build -DCMAKE_BUILD_TYPE=Release + + - name: Build + run: cmake --build build --config Release --parallel diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..cf5200c --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,352 @@ +name: Release + +on: + push: + tags: + - 'v*' + +permissions: + contents: write + +jobs: + build: + name: ${{ matrix.label }} + runs-on: ${{ matrix.os }} + timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + include: + - os: ubuntu-22.04 + label: linux-x86_64 + - os: windows-2022 + label: windows-x86_64 + + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + submodules: recursive + + # ----------------------------------------------------------------- + # Platform-specific dependency install (kept in sync with build.yml). + # patchelf is added on Linux for the $ORIGIN rpath rewrite in packaging. + # ----------------------------------------------------------------- + - name: Install Linux dependencies + if: runner.os == 'Linux' + env: + DEBIAN_FRONTEND: noninteractive + ACCEPT_EULA: 'Y' + run: | + set -euxo pipefail + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + ca-certificates curl gnupg lsb-release \ + build-essential cmake ninja-build pkg-config patchelf \ + libopencv-dev \ + libgl1-mesa-dev libglu1-mesa-dev \ + libxinerama-dev libxcursor-dev libxi-dev libxrandr-dev \ + libwayland-dev libxkbcommon-dev \ + libusb-1.0-0-dev libudev-dev \ + libgtk-3-dev + + # See build.yml — Microsoft prod (18.04 / bionic) feed for the Azure Kinect SDK. + sudo install -m 0755 -d /etc/apt/keyrings + curl -fsSL https://packages.microsoft.com/keys/microsoft.asc \ + | sudo gpg --dearmor -o /etc/apt/keyrings/microsoft.gpg + sudo chmod 0644 /etc/apt/keyrings/microsoft.gpg + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/microsoft.gpg] https://packages.microsoft.com/ubuntu/18.04/prod bionic main" \ + | sudo tee /etc/apt/sources.list.d/microsoft-prod.list + sudo apt-get update + + echo 'libk4a1.4 libk4a1.4/accepted-eula-hash string 0f5d5c5de396e4fee4c0753a21fee0c1ed726cf0316204edda484f08cb266d76' \ + | sudo debconf-set-selections + echo 'k4a-tools k4a-tools/accepted-eula-hash string 0f5d5c5de396e4fee4c0753a21fee0c1ed726cf0316204edda484f08cb266d76' \ + | sudo debconf-set-selections + + sudo apt-get install -y --no-install-recommends \ + libk4a1.4=1.4.2 libk4a1.4-dev=1.4.2 k4a-tools=1.4.2 + + # Windows: OpenCV via choco; k4a is vendored in-repo, no install needed. + - name: Install Windows dependencies + if: runner.os == 'Windows' + timeout-minutes: 25 + shell: pwsh + run: | + $ErrorActionPreference = 'Stop' + + choco install opencv --version=4.10.0 -y --no-progress + $opencvDir = "C:\tools\opencv\build" + if (-not (Test-Path "$opencvDir\OpenCVConfig.cmake")) { + $opencvDir = Get-ChildItem -Path "C:\tools\opencv" -Recurse -Filter "OpenCVConfig.cmake" | Select-Object -First 1 -ExpandProperty Directory | Select-Object -ExpandProperty FullName + } + "OpenCV_DIR=$opencvDir" | Out-File -FilePath $env:GITHUB_ENV -Append + "$opencvDir\x64\vc16\bin" | Out-File -FilePath $env:GITHUB_PATH -Append + + # ----------------------------------------------------------------- + # Configure + build. Windows pinned to x64 VS generator (the vendored + # k4a libs and OpenCV are x64-only). + # ----------------------------------------------------------------- + - name: Configure (Windows x64) + if: runner.os == 'Windows' + run: cmake -S . -B build -G "Visual Studio 17 2022" -A x64 -DCMAKE_BUILD_TYPE=Release + + - name: Configure (Linux) + if: runner.os == 'Linux' + run: cmake -S . -B build -DCMAKE_BUILD_TYPE=Release + + - name: Build + run: cmake --build build --config Release --parallel + + # ----------------------------------------------------------------- + # Package: produce a platform-native installer. + # Linux -> .deb (dpkg-deb, depends on system libgtk-3/libgl/libusb) + # Windows -> NSIS .exe installer (Start menu shortcut, Add/Remove entry) + # ----------------------------------------------------------------- + - name: Derive version from tag + id: ver + shell: bash + run: | + # ref_name on a tag push is "v1.2.3" — strip leading 'v' to get "1.2.3". + VERSION="${GITHUB_REF_NAME#v}" + # Releases triggered via workflow_dispatch / branch push won't have a v-tag; + # use a 0.0.0+sha fallback so the packaging step still works. + if [[ "$VERSION" == "$GITHUB_REF_NAME" && ! "$GITHUB_REF_NAME" =~ ^[0-9] ]]; then + VERSION="0.0.0-${GITHUB_SHA::7}" + fi + echo "version=$VERSION" >> "$GITHUB_OUTPUT" + echo "Resolved version: $VERSION" + + - name: Package (Linux .deb) + if: runner.os == 'Linux' + env: + VERSION: ${{ steps.ver.outputs.version }} + run: | + set -euxo pipefail + PKG="ir-tracking-app" + STAGE="$RUNNER_TEMP/deb-staging" + rm -rf "$STAGE" + mkdir -p "$STAGE/DEBIAN" "$STAGE/usr/bin" "$STAGE/usr/lib/$PKG" \ + "$STAGE/usr/share/applications" + + # Binary + bundled libs under /usr/lib// with $ORIGIN rpath. + cp build/ir-tracking-app "$STAGE/usr/lib/$PKG/" + + # NEEDED libs picked up via ldd: k4a, k4arecord, opencv, libusb. + # (We deliberately bundle the vendor libs so the app runs without the + # full Azure Kinect SDK installed on the target.) + ldd build/ir-tracking-app \ + | awk '/(libk4a|libk4arecord|libopencv|libusb-1\.0)/ { print $3 }' \ + | grep -v '^$' | sort -u \ + | while read -r lib; do cp -L "$lib" "$STAGE/usr/lib/$PKG/"; done + + # libdepthengine.so.2.0 is dlopen'd by libk4a at runtime, so it is + # NOT a NEEDED entry and ldd does not list it — copy it explicitly. + DEPTHENGINE="$(find /usr/lib/x86_64-linux-gnu -name 'libdepthengine.so.2.0' 2>/dev/null | head -n1)" + if [[ -z "$DEPTHENGINE" ]]; then + echo "ERROR: libdepthengine.so.2.0 not found — k4a depth will fail at runtime." >&2 + exit 1 + fi + cp -L "$DEPTHENGINE" "$STAGE/usr/lib/$PKG/" + + # Also bundle the k4a .so explicitly in case ldd reported a versioned + # symlink target outside the awk filter (defensive; cp -n avoids clobbering). + for so in $(find /usr/lib/x86_64-linux-gnu -maxdepth 2 \ + \( -name 'libk4a.so*' -o -name 'libk4arecord.so*' \) 2>/dev/null); do + cp -Ln "$so" "$STAGE/usr/lib/$PKG/" || true + done + + patchelf --set-rpath '$ORIGIN' "$STAGE/usr/lib/$PKG/ir-tracking-app" + + # Launcher shim on PATH. We set LD_LIBRARY_PATH (not just the binary's + # $ORIGIN rpath) because libk4a dlopen()s libdepthengine.so.2.0 by + # soname at runtime, and that lookup uses the loader's default search + # path, not the binary's rpath. Pointing LD_LIBRARY_PATH at the bundle + # dir lets the depth engine resolve next to libk4a. + cat > "$STAGE/usr/bin/$PKG" < "$STAGE/usr/share/applications/$PKG.desktop" < "$STAGE/DEBIAN/control" < + Homepage: https://github.com/stytim/AzureKinect-ToolTracker + Depends: libgtk-3-0, libgl1, libusb-1.0-0 + Installed-Size: $INSTALLED_KB + Description: Azure Kinect IR retro-reflective marker tool tracker + Tracks passive IR sphere markers from an Azure Kinect (live capture or + recorded .mkv playback) and publishes the tool pose over UDP. See the + GitHub README for usage. + EOF + + OUT="${PKG}_${VERSION}_amd64.deb" + dpkg-deb --build --root-owner-group "$STAGE" "$OUT" + ls -lh "$OUT" + dpkg-deb -I "$OUT" + + - name: Package (Windows .exe installer) + if: runner.os == 'Windows' + shell: pwsh + env: + VERSION: ${{ steps.ver.outputs.version }} + run: | + $ErrorActionPreference = 'Stop' + + # Stage every file that needs to end up in C:\Program Files\\ + $stage = "$env:RUNNER_TEMP\nsis-payload" + if (Test-Path $stage) { Remove-Item -Recurse -Force $stage } + New-Item -ItemType Directory -Force -Path $stage | Out-Null + + $exe = "build\Release\ir-tracking-app.exe" + if (-not (Test-Path $exe)) { $exe = "build\ir-tracking-app.exe" } + Copy-Item $exe -Destination $stage + + # The 3 vendored Azure Kinect DLLs (k4a.dll, k4arecord.dll, + # depthengine_2_0.dll) were copied next to the exe by CMake's + # WIN32 POST_BUILD step — grab them from the build dir. + $buildDir = Split-Path $exe -Parent + foreach ($dll in @("k4a.dll", "k4arecord.dll", "depthengine_2_0.dll")) { + $src = Join-Path $buildDir $dll + if (-not (Test-Path $src)) { throw "Expected vendored DLL not found next to exe: $src" } + Copy-Item $src -Destination $stage + } + + # OpenCV runtime (release opencv_world only; skip the *d.dll debug build). + $opencvBin = Join-Path $env:OpenCV_DIR "x64\vc16\bin" + Get-ChildItem -Path $opencvBin -Filter "opencv_world*.dll" ` + | Where-Object { $_.Name -notmatch "d\.dll$" } ` + | Copy-Item -Destination $stage + + # NSIS comes preinstalled on windows-2022 runners under C:\Program Files (x86)\NSIS. + $makensis = "C:\Program Files (x86)\NSIS\makensis.exe" + if (-not (Test-Path $makensis)) { + choco install nsis -y --no-progress + } + if (-not (Test-Path $makensis)) { + throw "NSIS (makensis.exe) not found after install attempt." + } + + $version = $env:VERSION + # NSIS's OutFile is resolved relative to the .nsi file's directory, + # so use an absolute path anchored at the workspace root — that + # way the upload-artifact glob can find it. + $outFileName = "ir-tracking-app-$version-${{ matrix.label }}-setup.exe" + $outFileAbs = Join-Path (Get-Location).Path $outFileName + $outFileNsi = $outFileAbs.Replace('\','\\') + $nsiPath = "$env:RUNNER_TEMP\installer.nsi" + $stageNsi = $stage.Replace('\','\\') + + # NSIS script — install to Program Files\IR Tracking App, Start Menu + # shortcut + uninstaller + Add/Remove Programs entry. + $nsi = @" + !define APPNAME "IR Tracking App" + !define APPID "ir-tracking-app" + !define VERSION "$version" + !define COMPANY "Medivis" + !define DESCRIPTION "Azure Kinect IR retro-reflective marker tool tracker" + !define REGKEY "Software\Microsoft\Windows\CurrentVersion\Uninstall\`${APPID}" + + Name "`${APPNAME} `${VERSION}" + OutFile "$outFileNsi" + Unicode true + InstallDir "`$PROGRAMFILES64\`${APPNAME}" + InstallDirRegKey HKLM "Software\`${APPNAME}" "InstallDir" + RequestExecutionLevel admin + ShowInstDetails show + ShowUninstDetails show + + Page directory + Page instfiles + UninstPage uninstConfirm + UninstPage instfiles + + Section "Install" + SetOutPath "`$INSTDIR" + File /r "$stageNsi\\*" + + CreateDirectory "`$SMPROGRAMS\`${APPNAME}" + CreateShortcut "`$SMPROGRAMS\`${APPNAME}\`${APPNAME}.lnk" "`$INSTDIR\ir-tracking-app.exe" + CreateShortcut "`$SMPROGRAMS\`${APPNAME}\Uninstall.lnk" "`$INSTDIR\Uninstall.exe" + + WriteUninstaller "`$INSTDIR\Uninstall.exe" + + WriteRegStr HKLM "Software\`${APPNAME}" "InstallDir" "`$INSTDIR" + WriteRegStr HKLM "`${REGKEY}" "DisplayName" "`${APPNAME}" + WriteRegStr HKLM "`${REGKEY}" "DisplayVersion" "`${VERSION}" + WriteRegStr HKLM "`${REGKEY}" "Publisher" "`${COMPANY}" + WriteRegStr HKLM "`${REGKEY}" "DisplayIcon" '"`$INSTDIR\ir-tracking-app.exe"' + WriteRegStr HKLM "`${REGKEY}" "UninstallString" '"`$INSTDIR\Uninstall.exe"' + WriteRegDWORD HKLM "`${REGKEY}" "NoModify" 1 + WriteRegDWORD HKLM "`${REGKEY}" "NoRepair" 1 + SectionEnd + + Section "Uninstall" + Delete "`$INSTDIR\Uninstall.exe" + RMDir /r "`$INSTDIR" + Delete "`$SMPROGRAMS\`${APPNAME}\`${APPNAME}.lnk" + Delete "`$SMPROGRAMS\`${APPNAME}\Uninstall.lnk" + RMDir "`$SMPROGRAMS\`${APPNAME}" + DeleteRegKey HKLM "`${REGKEY}" + DeleteRegKey HKLM "Software\`${APPNAME}" + SectionEnd + "@ + + Set-Content -Path $nsiPath -Value $nsi -Encoding UTF8 + & $makensis $nsiPath + if ($LASTEXITCODE -ne 0) { throw "makensis failed (exit $LASTEXITCODE)." } + Get-Item $outFileAbs | Format-List Length, Name + + # ----------------------------------------------------------------- + # Upload the installer for the release job to collect. + # ----------------------------------------------------------------- + - name: Upload artifact + uses: actions/upload-artifact@v4 + with: + name: ir-tracking-app-${{ matrix.label }} + path: | + *.deb + *-setup.exe + if-no-files-found: error + retention-days: 7 + + release: + name: Publish GitHub Release + needs: build + runs-on: ubuntu-latest + steps: + - name: Download all artifacts + uses: actions/download-artifact@v4 + with: + path: artifacts + merge-multiple: true + + - name: List artifacts + run: ls -lh artifacts + + - name: Create GitHub Release + uses: softprops/action-gh-release@v2 + with: + files: artifacts/* + generate_release_notes: true + fail_on_unmatched_files: true diff --git a/CMakeLists.txt b/CMakeLists.txt index 8f12626..6af5ac3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -65,9 +65,17 @@ add_library(IRToolTracking STATIC ${TRACKER_SOURCE_FILES}) # Specify C++17 for IRToolTracking library target_compile_features(IRToolTracking PUBLIC cxx_std_17) -# Link RealSense SDK and OpenCV with the IRToolTracking library +# Link Azure Kinect SDK and OpenCV with the IRToolTracking library target_link_libraries(IRToolTracking PUBLIC Eigen3 k4a::k4a opencv_core opencv_imgproc opencv_imgcodecs opencv_video) +# The app uses k4arecord (k4a_playback_* for .mkv playback). On Windows the vendored +# extern/k4a target already links k4arecord into k4a::k4a; on Linux find_package(k4a) +# exposes it as a separate k4a::k4arecord target, so link it there to resolve the +# k4a_playback_* symbols. +if(NOT WIN32 AND TARGET k4a::k4arecord) + target_link_libraries(IRToolTracking PUBLIC k4a::k4arecord) +endif() + ################################################################################ # Dependencies for the main application diff --git a/include/IRKalmanFilter.h b/include/IRKalmanFilter.h index f925a87..818bf33 100644 --- a/include/IRKalmanFilter.h +++ b/include/IRKalmanFilter.h @@ -52,10 +52,17 @@ class IRToolKalmanFilter cv::Mat R = cv::Mat::eye(3, 3, CV_32F) * m_fMeasurementNoise; // measurement noise m_filter.measurementNoiseCov = R; - // Initialize the state estimate (x) and the error covariance matrix (P) - cv::Mat x = cv::Mat::zeros(6, 1, CV_32F); // initial state is all zeros + // Seed position from the first measurement so the filter doesn't have + // to spend several frames pulling itself from (0,0,0) up to the marker. + // Velocity starts at zero — the next correct() will refine it. + cv::Mat x = cv::Mat::zeros(6, 1, CV_32F); + x.at(0, 0) = value[0]; + x.at(1, 0) = value[1]; + x.at(2, 0) = value[2]; + cv::Mat P = cv::Mat::eye(6, 6, CV_32F); // initial error covariance is identity matrix m_filter.statePre = x; + m_filter.statePost = x.clone(); m_filter.errorCovPost = P; m_bInitialized = true; diff --git a/include/IRStructs.h b/include/IRStructs.h index 3dec9e4..54abf54 100644 --- a/include/IRStructs.h +++ b/include/IRStructs.h @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -10,6 +11,13 @@ #include "IRKalmanFilter.h" +// Bucket sphere radii at 0.1 mm so map keys are no longer +// vulnerable to bit-exact float equality when the same nominal value is +// produced by different expressions (UI input vs. calibration output). +inline int SphereRadiusKey(float radius_mm) { + return static_cast(std::lround(radius_mm * 10.0f)); +} + struct Side { int id_from{ 0 }; @@ -30,7 +38,7 @@ struct AHATFrame { double timestamp; cv::Mat device_pose; cv::Mat cvAbImage; - uint16_t* pDepth; + std::vector pDepth; uint32_t depthWidth; uint32_t depthHeight; }; @@ -41,9 +49,9 @@ struct ProcessedAHATFrame cv::Mat device_pose; uint num_spheres; cv::Mat3f spheres_xyd; - std::map spheres_xyz_per_mm; - std::map> ordered_sides_per_mm; - std::map map_per_mm; + std::map spheres_xyz_per_mm; + std::map> ordered_sides_per_mm; + std::map map_per_mm; }; struct ToolResult @@ -119,6 +127,4 @@ struct IRTrackedTool cv::Vec3f cur_position_cheap{}; std::vector unfiltered_sphere_positions; double timestamp{ 0 }; - - bool tracking_finished = true; }; \ No newline at end of file diff --git a/include/IRToolTrack.h b/include/IRToolTrack.h index f4ad10f..a7275c7 100644 --- a/include/IRToolTrack.h +++ b/include/IRToolTrack.h @@ -2,7 +2,9 @@ #include #include +#include #include +#include #include #include @@ -30,6 +32,8 @@ class IRToolTracker m_pRealSenseToolTracking = pRealSenseToolTracking; } + ~IRToolTracker(); + void AddFrame(void* pAbImage, void* pDepth, uint32_t depthWidth, uint32_t depthHeight, cv::Mat _pose, double _timestamp); bool AddTool(cv::Mat3f spheres, float sphere_radius, std::string identifier, uint min_visible_spheres, float lowpass_rotation, float lowpass_position); @@ -43,7 +47,7 @@ class IRToolTracker void SetThreshold(int threshold); void SetMinMaxSize(int min, int max); - const cv::Mat& GetProcessedFrame(); + cv::Mat GetProcessedFrame(); void StopTracking(); @@ -64,24 +68,26 @@ class IRToolTracker void regionGrowing(const cv::Mat &depthImage, const std::vector &contour, const cv::Point &startPoint, cv::Mat &visited, double &averageDepth, int &validPixelCount); - bool ProcessFrame(AHATFrame *rawFrame, ProcessedAHATFrame &result); - - void TrackTool(IRTrackedTool &tool, ProcessedAHATFrame &frame, ToolResultContainer &result); + bool ProcessFrame(AHATFrame& rawFrame, ProcessedAHATFrame& result); - void UnionSegmentation(ToolResultContainer* raw_solutions, int num_tools, ProcessedAHATFrame frame); + void TrackTool(IRTrackedTool &tool, const ProcessedAHATFrame &frame, ToolResultContainer &result); - cv::Mat MatchPointsKabsch(IRTrackedTool tool, ProcessedAHATFrame frame, std::vector sphere_ids, std::vector occluded_nodes); + void UnionSegmentation(std::vector &raw_solutions, const ProcessedAHATFrame &frame); - cv::Mat FlipTransformRightLeft(cv::Mat hololens_transform); + cv::Mat MatchPointsKabsch(IRTrackedTool &tool, const ProcessedAHATFrame &frame, const std::vector &sphere_ids, const std::vector &occluded_nodes); void ConstructMap(cv::Mat3f spheres_xyz, int num_spheres, cv::Mat& result_map, std::vector& result_ordered_sides); - bool m_bShouldStop = false; + 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; - AHATFrame* m_CurrentFrame = nullptr; + std::unique_ptr m_CurrentFrame; std::mutex m_MutexCurFrame; std::map m_ToolIndexMapping; @@ -89,8 +95,8 @@ class IRToolTracker float m_fToleranceSide = 4.0f; float m_fToleranceAvg = 4.0f; - bool m_bIsCurrentlyTracking = false; - bool m_bIsCurrentlyCalibrating = false; + std::atomic m_bIsCurrentlyTracking{false}; + std::atomic m_bIsCurrentlyCalibrating{false}; std::shared_ptr m_TrackingThread; std::shared_ptr m_CalibrationThread; @@ -98,8 +104,13 @@ class IRToolTracker double m_lTrackedTimestamp = 0; std::mutex mtx_frames; + // Tracking-thread-only scratch buffer drawn into during each ProcessFrame/MatchPointsKabsch. + cv::Mat m_WorkingFrame; + // Published copy snapshot to viewers; only touched under mtx_frames. cv::Mat m_ProcessedFrame; + void PublishWorkingFrame(); + IRToolTracking* m_pRealSenseToolTracking; ushort m_Threshold = 3000; @@ -108,6 +119,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 0e9d4c3..262a886 100644 --- a/include/IRToolTracking.h +++ b/include/IRToolTracking.h @@ -10,8 +10,6 @@ #include #include -#include -#include #include #include "IRToolTrack.h" @@ -32,39 +30,6 @@ inline void JoinThread(std::shared_ptr& th) -class FrameQueue { -private: - std::queue queue; - std::mutex mutex; - std::condition_variable cond; - -public: - cv::Mat lastframe; - void push(cv::Mat frame) { - std::lock_guard lock(mutex); - queue.push(frame); - cond.notify_one(); - } - - cv::Mat pop() { - std::unique_lock lock(mutex); - cond.wait(lock, [this]{ return !queue.empty(); }); - while (queue.size() > 1) { - queue.pop(); - } - auto frame = std::move(queue.front()); - lastframe = frame.clone(); - queue.pop(); - return frame; - } - - bool empty() { - std::lock_guard lock(mutex); - return queue.empty(); - } -}; - - class IRToolTracking { public: IRToolTracking(); diff --git a/include/ROMParser.h b/include/ROMParser.h index 3d40b0d..9d14af1 100644 --- a/include/ROMParser.h +++ b/include/ROMParser.h @@ -4,6 +4,8 @@ #include #include #include +#include +#include class ROMParser { public: @@ -31,8 +33,28 @@ class ROMParser { int num_markers = 4; void parse_rom_data(const std::vector& rom_data) { - num_markers = static_cast(rom_data[28]); - int pos = 72; + constexpr int kMarkerCountOffset = 28; + constexpr int kMarkerDataOffset = 72; + constexpr int kMarkerStride = 12; + + if (rom_data.size() <= kMarkerCountOffset) { + std::cerr << "ROM file too small to contain marker count." << std::endl; + num_markers = 0; + return; + } + + const int candidate_count = static_cast(rom_data[kMarkerCountOffset]); + const std::size_t required = static_cast(kMarkerDataOffset) + + static_cast(candidate_count) * kMarkerStride; + if (rom_data.size() < required) { + std::cerr << "ROM file too small for declared marker count (" << candidate_count + << ")." << std::endl; + num_markers = 0; + return; + } + + num_markers = candidate_count; + int pos = kMarkerDataOffset; for (int i = 0; i < num_markers; ++i) { float x, y, z; std::memcpy(&x, &rom_data[pos], sizeof(float)); @@ -51,7 +73,7 @@ class ROMParser { marker_positions.push_back(y); marker_positions.push_back(z); - pos += 12; + pos += kMarkerStride; } } }; \ No newline at end of file diff --git a/include/ViewerWindow.h b/include/ViewerWindow.h index 9e5bd15..881531c 100644 --- a/include/ViewerWindow.h +++ b/include/ViewerWindow.h @@ -6,7 +6,8 @@ #include #include #include -#include +#include +#include #include "IRToolTracking.h" class ViewerWindow { @@ -23,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 { @@ -42,32 +47,52 @@ 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; std::shared_ptr processingThread; std::shared_ptr udpThread; - bool udpEnabled = false; + std::atomic udpEnabled{false}; std::shared_ptr udpReceiveThread; - bool multiEnabled = false; + std::atomic multiEnabled{false}; std::shared_ptr csvThread; - bool csvEnabled = false; + std::atomic csvEnabled{false}; std::map extrinsics; 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 @@ -93,16 +118,19 @@ class ViewerWindow { NanoSocket receiveSocket; NanoAddress sendAddress = {}; NanoAddress receiveAddress = {}; - bool m_connected; - bool m_receiveconnected; - char ipAddress[16] = "127.0.0.1"; + 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; int frequency = 100; int recordFrequency = 10; int duration = 20; std::string csvFileName = "tracking_data.csv"; - bool finishedRecord = false; + std::atomic finishedRecord{false}; }; #endif // VIEWER_WINDOW_H diff --git a/src/IRToolTrack.cpp b/src/IRToolTrack.cpp index f287c2b..3336194 100644 --- a/src/IRToolTrack.cpp +++ b/src/IRToolTrack.cpp @@ -2,21 +2,47 @@ #include "IRToolTracking.h" #include #include +#include +#include -#define DISABLE_LOWPASS FALSE -#define DISABLE_KALMAN FALSE + +IRToolTracker::~IRToolTracker() +{ + m_bShouldStop = true; + if (m_TrackingThread && m_TrackingThread->joinable()) { + try { m_TrackingThread->join(); } catch (const std::system_error&) {} + } + if (m_CalibrationThread && m_CalibrationThread->joinable()) { + try { m_CalibrationThread->join(); } catch (const std::system_error&) {} + } + // m_CurrentFrame is a unique_ptr; its destructor frees any pending frame. +} + + +// Compile-time toggles for the smoothing filters. Override with -D… if you +// want to disable a filter without editing the source. +#ifndef DISABLE_LOWPASS +#define DISABLE_LOWPASS 0 +#endif +#ifndef DISABLE_KALMAN +#define DISABLE_KALMAN 0 +#endif 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; @@ -27,70 +53,54 @@ 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; - m_MutexCurFrame.lock(); - if (m_CurrentFrame == nullptr) { - m_MutexCurFrame.unlock(); + std::unique_ptr rawFrame; + { + std::lock_guard lock(m_MutexCurFrame); + rawFrame = std::move(m_CurrentFrame); + } + if (!rawFrame) { std::this_thread::sleep_for(std::chrono::milliseconds(5)); continue; } - - //Copy pointer to frame - AHATFrame* rawFrame = m_CurrentFrame; - m_CurrentFrame = nullptr; - m_MutexCurFrame.unlock(); - - int current_num_tools = m_Tools.size(); - ToolResultContainer* raw_results = new ToolResultContainer[current_num_tools]; ProcessedAHATFrame processedFrame; - - if (!ProcessFrame(rawFrame, processedFrame)) { + if (!ProcessFrame(*rawFrame, processedFrame)) { continue; } - - //std::vector tool_track_threads(current_num_tools); - for (int i = 0; i < current_num_tools; i++) { - IRTrackedTool tool = m_Tools.at(i); - if (!tool.tracking_finished) - continue; + const int current_num_tools = static_cast(m_Tools.size()); + std::vector raw_results(current_num_tools); - ToolResultContainer result{ i, std::vector() }; - - //tool_track_threads.at(i) = std::thread(&IRToolTracker::TrackTool, this, tool, processedFrame, result); - TrackTool(tool, processedFrame, result); - raw_results[i] = result; - //ProcessEnvFrame(processedFrame, result); + for (int i = 0; i < current_num_tools; i++) { + raw_results[i].tool_id = i; + TrackTool(m_Tools.at(i), processedFrame, raw_results[i]); } - //for (auto & thread : tool_track_threads) { - // thread.join(); - //} - UnionSegmentation(raw_results, current_num_tools, processedFrame); - - delete[] raw_results; - //TODO: make sure i didnt create a memory leak here + UnionSegmentation(raw_results, processedFrame); + PublishWorkingFrame(); } m_bIsCurrentlyTracking = false; } -void IRToolTracker::TrackTool(IRTrackedTool &tool, ProcessedAHATFrame &frame, ToolResultContainer &result) +void IRToolTracker::TrackTool(IRTrackedTool &tool, const ProcessedAHATFrame &frame, ToolResultContainer &result) { - tool.tracking_finished = false; if (frame.num_spheres < tool.min_visible_spheres) { //Not enough spheres for the tool are available - tool.tracking_finished = true; return; } std::vector eligible_sides; - - auto it_sides = frame.ordered_sides_per_mm.find(tool.sphere_radius); - std::vector frame_ordered_sides = it_sides->second; - auto it_map = frame.map_per_mm.find(tool.sphere_radius); + const int radius_key = SphereRadiusKey(tool.sphere_radius); + auto it_sides = frame.ordered_sides_per_mm.find(radius_key); + 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 @@ -122,7 +132,6 @@ void IRToolTracker::TrackTool(IRTrackedTool &tool, ProcessedAHATFrame &frame, To } if (eligible_sides.size() == 0 && max_occluded_spheres == 0) { - tool.tracking_finished = true; return; } if (eligible_sides.size() != 0) @@ -136,7 +145,6 @@ void IRToolTracker::TrackTool(IRTrackedTool &tool, ProcessedAHATFrame &frame, To } if (eligible_sides.size() == 0 && max_occluded_spheres == 0) { - tool.tracking_finished = true; return; } @@ -232,98 +240,92 @@ void IRToolTracker::TrackTool(IRTrackedTool &tool, ProcessedAHATFrame &frame, To search_list.push_back(search_entry{ searched_ids_new, curr.combined_error + error_new, curr.num_sides + error_counter, curr.occluded_nodes_tool}); } } - tool.tracking_finished = true; return; } -void IRToolTracker::UnionSegmentation(ToolResultContainer* raw_solutions, int num_tools, ProcessedAHATFrame frame) { - int* tool_solutions = new int[num_tools]; +void IRToolTracker::UnionSegmentation(std::vector &raw_solutions, const ProcessedAHATFrame &frame) { std::vector unique_solutions; + const int num_tools = static_cast(raw_solutions.size()); for (int i = 0; i < num_tools; i++) { - tool_solutions[i] = 0; - ToolResultContainer tool_results = raw_solutions[i]; - - //std::cout << "Tool " << i << " has " << tool_results.candidates.size() << " candidates" << std::endl; - if (tool_results.candidates.size() == 0) + std::vector &candidates = raw_solutions[i].candidates; + if (candidates.empty()) continue; - std::vector ordered_candidates = tool_results.candidates; - std::sort(ordered_candidates.begin(), ordered_candidates.end(), &ToolResult::compare); + std::sort(candidates.begin(), candidates.end(), &ToolResult::compare); - for (ToolResult candidate : ordered_candidates) + for (ToolResult candidate : candidates) { candidate.tool_id = i; - unique_solutions.push_back(candidate); - tool_solutions[i]++; + unique_solutions.push_back(std::move(candidate)); } } std::sort(unique_solutions.begin(), unique_solutions.end(), &ToolResult::compare); - while (unique_solutions.size() > 0) + 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) { - ToolResult current = unique_solutions.front(); - int cur_toolid = current.tool_id; - unique_solutions.erase(unique_solutions.begin()); - cv::Mat result = MatchPointsKabsch(m_Tools[cur_toolid], frame, current.sphere_ids, current.occluded_nodes); + if (claimed_tools[current.tool_id]) + continue; + + bool overlap = false; + for (int sid : current.sphere_ids) { + if (used_spheres.count(sid) > 0) { + overlap = true; + break; + } + } + if (overlap) + continue; + + 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(cur_toolid).cur_transform = result.clone(); - m_Tools.at(cur_toolid).timestamp = frame.timestamp; + committed.emplace_back(current.tool_id, result.clone()); } - std::vector remaining_unique_solutions; - for (ToolResult next_check : unique_solutions) { - if (next_check.tool_id == cur_toolid) - continue; - - bool used = false; - for (auto cursphere : current.sphere_ids) - { - if (used) - { - break; - } - for (auto nexsphere : next_check.sphere_ids) - { - if (cursphere == nexsphere) - { - used = true; - break; - } + claimed_tools[current.tool_id] = true; + used_spheres.insert(current.sphere_ids.begin(), current.sphere_ids.end()); + } - } - } - // std::vector intersection; - //std::set_intersection(current.sphere_ids.begin(), current.sphere_ids.end(), next_check.sphere_ids.begin(), next_check.sphere_ids.end(), intersection.begin()); - //if (intersection.size() > 0) - // continue; - if (used) - { - continue; - } - remaining_unique_solutions.push_back(next_check); + { + 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; } - unique_solutions = remaining_unique_solutions; + m_lTrackedTimestamp = frame.timestamp; } - delete[] tool_solutions; - m_lTrackedTimestamp = frame.timestamp; - return; } -cv::Mat IRToolTracker::MatchPointsKabsch(IRTrackedTool tool, ProcessedAHATFrame frame, std::vector sphere_ids, std::vector occluded_nodes) { - int num_points = tool.num_spheres-occluded_nodes.size(); +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(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; @@ -363,13 +365,13 @@ cv::Mat IRToolTracker::MatchPointsKabsch(IRTrackedTool tool, ProcessedAHATFrame m_pRealSenseToolTracking->ProjectPointToPixel(xyz,uv); cv::Point center(uv[0], uv[1]); - cv::drawMarker(m_ProcessedFrame, center, cv::Scalar(0, 255, 0), cv::MARKER_CROSS, 20, 2); + cv::drawMarker(m_WorkingFrame, center, cv::Scalar(0, 255, 0), cv::MARKER_CROSS, 20, 2); cv::Mat sphere_world_mat = hololens_pose_mm * sphere_frame_mat; cv::Vec3f sphere_world = cv::Vec3f(sphere_world_mat.at(0, 0), sphere_world_mat.at(1, 0), sphere_world_mat.at(2, 0)); //Filter the resulting world position -#if !DEBUG_NO_FILTER && !DISABLE_KALMAN +#if !DISABLE_KALMAN sphere_world = tool.sphere_kalman_filters.at(tool_node_id).FilterData(sphere_world); #endif tool_node_id++; @@ -439,8 +441,6 @@ cv::Mat IRToolTracker::MatchPointsKabsch(IRTrackedTool tool, ProcessedAHATFrame transform_matrix.at(2, 3) = t.at(2, 0) / 1000.f; transform_matrix.at(3, 3) = 1.f; - // transform_matrix = FlipTransformRightLeft(transform_matrix); - //Copy translation and convert mm to m cv::Vec3f position; position[0] = transform_matrix.at(0, 3); @@ -460,7 +460,7 @@ cv::Mat IRToolTracker::MatchPointsKabsch(IRTrackedTool tool, ProcessedAHATFrame Eigen::Quaternionf rotation(quat[3], quat[0], quat[1], quat[2]); -#if !DISABLE_LOWPASS && !DEBUG_NO_FILTER +#if !DISABLE_LOWPASS { Eigen::Quaternionf rotation_old(tool.cur_transform.at(6, 0), tool.cur_transform.at(3, 0), tool.cur_transform.at(4, 0), tool.cur_transform.at(5, 0)); rotation = rotation_old.slerp(tool.lowpass_factor_rotation, rotation); @@ -489,19 +489,6 @@ cv::Mat IRToolTracker::MatchPointsKabsch(IRTrackedTool tool, ProcessedAHATFrame } -cv::Mat IRToolTracker::FlipTransformRightLeft(cv::Mat transform_rhs) -{ - //Bring to unity coordinate system - cv::Mat flipz = cv::Mat::ones(4, 4, CV_32F); - flipz.at(2, 0) = -1.f; - flipz.at(0, 2) = -1.f; - flipz.at(2, 1) = -1.f; - flipz.at(1, 2) = -1.f; - flipz.at(2, 3) = -1.f; - cv::Mat transform_lhs = transform_rhs.mul(flipz); - return transform_lhs; -} - void IRToolTracker::ConstructMap(cv::Mat3f spheres_xyz, int num_spheres, cv::Mat& map, std::vector& ordered_sides) { for (int i = 0; i < num_spheres; i++) { @@ -536,18 +523,17 @@ void IRToolTracker::AddFrame(void* pAbImage, void* pDepth, uint32_t depthWidth, cv::Mat cvAbImage_origin(cv::Size(depthWidth, depthHeight), CV_16UC1, (void*)pAbImage); cv::Mat cvAbImage = cvAbImage_origin.clone(); - m_MutexCurFrame.lock(); - - if (m_CurrentFrame != nullptr) { - delete[] m_CurrentFrame->pDepth; - delete m_CurrentFrame; - } - - m_CurrentFrame = new AHATFrame { _timestamp, _pose, cvAbImage, new uint16_t[depthWidth * depthHeight], depthWidth, depthHeight }; - memcpy(m_CurrentFrame->pDepth, pDepth, depthWidth * depthHeight * sizeof(uint16_t)); - - m_MutexCurFrame.unlock(); - + auto frame = std::make_unique(); + frame->timestamp = _timestamp; + frame->device_pose = _pose; + frame->cvAbImage = std::move(cvAbImage); + frame->depthWidth = depthWidth; + frame->depthHeight = depthHeight; + frame->pDepth.resize(static_cast(depthWidth) * depthHeight); + std::memcpy(frame->pDepth.data(), pDepth, frame->pDepth.size() * sizeof(uint16_t)); + + std::lock_guard lock(m_MutexCurFrame); + m_CurrentFrame = std::move(frame); } void IRToolTracker::SetThreshold(int threshold) @@ -555,10 +541,20 @@ void IRToolTracker::SetThreshold(int threshold) m_Threshold = threshold; } -const cv::Mat& IRToolTracker::GetProcessedFrame() +cv::Mat IRToolTracker::GetProcessedFrame() { std::lock_guard lock(mtx_frames); - return m_ProcessedFrame; + if (m_ProcessedFrame.empty()) + return cv::Mat(); + return m_ProcessedFrame.clone(); +} + +void IRToolTracker::PublishWorkingFrame() +{ + if (m_WorkingFrame.empty()) + return; + std::lock_guard lock(mtx_frames); + m_WorkingFrame.copyTo(m_ProcessedFrame); } void IRToolTracker::SetMinMaxSize(int min, int max) @@ -619,7 +615,7 @@ void IRToolTracker::regionGrowing(const cv::Mat &depthImage, const std::vectorcvAbImage, thresholdedIR, m_Threshold, 255, cv::THRESH_BINARY); + cv::threshold(rawFrame.cvAbImage, thresholdedIR, m_Threshold, 255, cv::THRESH_BINARY); - thresholdedIR.convertTo(rawFrame->cvAbImage, CV_8UC1); + thresholdedIR.convertTo(rawFrame.cvAbImage, CV_8UC1); - cv::cvtColor(rawFrame->cvAbImage.clone(), m_ProcessedFrame, cv::COLOR_GRAY2RGB); + cv::cvtColor(rawFrame.cvAbImage.clone(), m_WorkingFrame, cv::COLOR_GRAY2RGB); - cv::Mat depthImage(rawFrame->depthHeight, rawFrame->depthWidth, CV_16U, rawFrame->pDepth); + cv::Mat depthImage(rawFrame.depthHeight, rawFrame.depthWidth, CV_16U, rawFrame.pDepth.data()); double perimeter, circ; static constexpr double PI = 3.141592653589793238462; - int areaCount = cv::connectedComponentsWithStats(rawFrame->cvAbImage, labels, stats, centroids, 8); + int areaCount = cv::connectedComponentsWithStats(rawFrame.cvAbImage, labels, stats, centroids, 8); for (int i = 1; i < areaCount; ++i) { auto area = stats.at(i, cv::CC_STAT_AREA); @@ -705,16 +701,14 @@ bool IRToolTracker::ProcessFrame(AHATFrame* rawFrame, ProcessedAHATFrame &result if (num_spheres < 3) { //If theres less than 3 points visible, theres no tool to track - //Free memory - delete[] rawFrame->pDepth; - delete rawFrame; + //Caller owns rawFrame via unique_ptr, so no manual cleanup needed. return false; } //Create 3d coordinates for every possible sphere size - std::map> ordered_sides_per_mm; - std::map map_per_mm; - std::map spheres_xyz_per_mm; + std::map> ordered_sides_per_mm; + std::map map_per_mm; + std::map spheres_xyz_per_mm; if (!m_bIsCurrentlyCalibrating) @@ -722,7 +716,8 @@ bool IRToolTracker::ProcessFrame(AHATFrame* rawFrame, ProcessedAHATFrame &result for (IRTrackedTool tool : m_Tools) { float cur_radius = tool.sphere_radius; - if (!(spheres_xyz_per_mm.find(cur_radius) == spheres_xyz_per_mm.end())) { + const int radius_key = SphereRadiusKey(cur_radius); + if (spheres_xyz_per_mm.find(radius_key) != spheres_xyz_per_mm.end()) { //We already created this map continue; } @@ -746,15 +741,16 @@ bool IRToolTracker::ProcessFrame(AHATFrame* rawFrame, ProcessedAHATFrame &result ConstructMap(spheres_xyz, num_spheres, map, ordered_sides); - ordered_sides_per_mm.insert({ cur_radius, ordered_sides }); - map_per_mm.insert({ cur_radius, map }); - spheres_xyz_per_mm.insert({ cur_radius, spheres_xyz }); + ordered_sides_per_mm.insert({ radius_key, ordered_sides }); + map_per_mm.insert({ radius_key, map }); + spheres_xyz_per_mm.insert({ radius_key, spheres_xyz }); } } else { float cur_radius = m_fCalibrationSphereRadius; + const int radius_key = SphereRadiusKey(cur_radius); cv::Mat3f spheres_xyz = spheres.clone(); spheres_xyz.forEach( [&](cv::Vec3f& xyz, const int* position) -> void { @@ -773,7 +769,7 @@ bool IRToolTracker::ProcessFrame(AHATFrame* rawFrame, ProcessedAHATFrame &result // Convert point coordinates to cv::Point cv::Point center(uv[0], uv[1]); - cv::drawMarker(m_ProcessedFrame, center, cv::Scalar(255, 255, 0), cv::MARKER_CROSS, 20, 2); + cv::drawMarker(m_WorkingFrame, center, cv::Scalar(255, 255, 0), cv::MARKER_CROSS, 20, 2); } ); @@ -783,33 +779,31 @@ bool IRToolTracker::ProcessFrame(AHATFrame* rawFrame, ProcessedAHATFrame &result ConstructMap(spheres_xyz, 4, map, ordered_sides); - ordered_sides_per_mm.insert({ cur_radius, ordered_sides }); - map_per_mm.insert({ cur_radius, map }); - spheres_xyz_per_mm.insert({ cur_radius, spheres_xyz }); + ordered_sides_per_mm.insert({ radius_key, ordered_sides }); + map_per_mm.insert({ radius_key, map }); + spheres_xyz_per_mm.insert({ radius_key, spheres_xyz }); } - result.timestamp = rawFrame->timestamp; - result.device_pose = rawFrame->device_pose; + result.timestamp = rawFrame.timestamp; + result.device_pose = rawFrame.device_pose; result.num_spheres = static_cast(num_spheres); result.spheres_xyd = spheres.clone(); result.spheres_xyz_per_mm = spheres_xyz_per_mm; result.ordered_sides_per_mm = ordered_sides_per_mm; result.map_per_mm = map_per_mm; - - //Free memory - delete[] rawFrame->pDepth; - delete rawFrame; - return true; } 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) { @@ -838,8 +832,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) { @@ -864,15 +860,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; @@ -886,8 +885,14 @@ bool IRToolTracker::RemoveTool(std::string identifier) bool IRToolTracker::RemoveAllTools() { - m_Tools.clear(); - m_ToolIndexMapping.clear(); + if (m_bIsCurrentlyTracking) { + StopTracking(); + } + { + std::lock_guard lock(m_ToolsMutex); + m_Tools.clear(); + m_ToolIndexMapping.clear(); + } return true; } @@ -895,6 +900,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; } @@ -923,6 +933,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; @@ -932,6 +944,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; @@ -958,26 +974,26 @@ void IRToolTracker::CalibrateTool() std::vector processedFrames; processedFrames.reserve(MAX_CALIBRATION_FRAMES); - while (!m_bShouldStop) { - m_MutexCurFrame.lock(); - if (m_CurrentFrame == nullptr) { - m_MutexCurFrame.unlock(); + // 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); + rawFrame = std::move(m_CurrentFrame); + } + if (!rawFrame) { std::this_thread::sleep_for(std::chrono::milliseconds(5)); continue; } - - - //Copy pointer to frame - AHATFrame* rawFrame = m_CurrentFrame; - m_CurrentFrame = nullptr; - m_MutexCurFrame.unlock(); ProcessedAHATFrame processedFrame; - - if (!ProcessFrame(rawFrame, processedFrame)) { + if (!ProcessFrame(*rawFrame, processedFrame)) { continue; } processedFrames.push_back(processedFrame); + PublishWorkingFrame(); // If enough data is collected, perform calibration if (processedFrames.size() == MAX_CALIBRATION_FRAMES) { @@ -986,21 +1002,58 @@ 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 - NUM_CALIBRATION_SPHERES = processedFrames[0].spheres_xyz_per_mm.at(m_fCalibrationSphereRadius).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(m_fCalibrationSphereRadius); + auto it_sides = frame.ordered_sides_per_mm.find(calib_radius_key); + 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; - - auto it_spheres_xyz = frame.spheres_xyz_per_mm.find(m_fCalibrationSphereRadius); 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(); @@ -1033,6 +1086,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 c86f84b..604b298 100644 --- a/src/IRToolTracking.cpp +++ b/src/IRToolTracking.cpp @@ -24,6 +24,7 @@ void IRToolTracking::queryDevices() { size_t serial_number_length = sizeof(serial_number); k4a_device_get_serialnum(device, serial_number, &serial_number_length); k4a_device_close(device); + device = NULL; std::stringstream ss; ss << "[" << i << "] " << "K4A Device" << " (" << serial_number << ")"; @@ -37,12 +38,14 @@ void IRToolTracking::initializeFromFile(const std::string& file) { if (result != K4A_RESULT_SUCCEEDED) { std::cerr << "Failed to open file " << file << std::endl; + Terminated = true; return; } result = k4a_playback_get_calibration(playback, &calibration); if (result != K4A_RESULT_SUCCEEDED) { std::cerr << "Failed to get calibration" << std::endl; + Terminated = true; return; } Terminated = false; @@ -62,6 +65,8 @@ void IRToolTracking::initialize(int index, int width, int height) { { std::cout << "Failed to open device" << std::endl; k4a_device_close(device); + device = NULL; + Terminated = true; return; } @@ -73,6 +78,8 @@ void IRToolTracking::initialize(int index, int width, int height) { { std::cout << "Failed to get calibration" << std::endl; k4a_device_close(device); + device = NULL; + Terminated = true; return; } @@ -120,6 +127,7 @@ void IRToolTracking::processStreams() { if (K4A_RESULT_SUCCEEDED != k4a_device_start_cameras(device, &config)) { std::cout << "Failed to start device" << std::endl; + Terminated = true; return; } } @@ -132,6 +140,12 @@ void IRToolTracking::processStreams() { int frame_width = 0; int frame_height = 0; + // Everything in the capture loop can throw (cv::Exception on a malformed frame, or + // a k4a wrapper error if the device disconnects mid-stream). This runs on a worker + // thread, so an escaping exception would call std::terminate() and abort the whole + // app. Catch it, flag the session terminated, and return so the GUI thread can join + // and shut down cleanly. + try { // Continuously capture frames and process them while (!Terminated) { @@ -147,6 +161,7 @@ void IRToolTracking::processStreams() { break; case K4A_STREAM_RESULT_EOF: std::cout << "Reached end of the file" << std::endl; + Terminated = true; return; } } @@ -164,6 +179,8 @@ void IRToolTracking::processStreams() { case K4A_WAIT_RESULT_FAILED: std::cout << "Failed to read a capture" << std::endl; k4a_device_close(device); + device = NULL; + Terminated = true; return; } } @@ -181,6 +198,7 @@ void IRToolTracking::processStreams() { if (ir_image == NULL) { std::cout<<"IR16 None"<IsTracking() || m_IRToolTracker->IsCalibrating()) && timestamp > m_latestTrackedFrame) + if (m_IRToolTracker != nullptr) { - // Create a 4x4 identity matrix - cv::Mat pose = cv::Mat::eye(4, 4, CV_32F); - m_IRToolTracker->AddFrame(left_frame_image.data, depth_frame_image.data, left_frame_image.cols, left_frame_image.rows, pose ,timestamp); - m_latestTrackedFrame = playFromFile ? -1 : timestamp; + if ((m_IRToolTracker->IsTracking() || m_IRToolTracker->IsCalibrating()) && timestamp > m_latestTrackedFrame) + { + // Create a 4x4 identity matrix + cv::Mat pose = cv::Mat::eye(4, 4, CV_32F); + m_IRToolTracker->AddFrame(left_frame_image.data, depth_frame_image.data, left_frame_image.cols, left_frame_image.rows, pose ,timestamp); + m_latestTrackedFrame = playFromFile ? -1 : timestamp; + } + + // Publish the latest preview under the frame mutex so getNextIRFrame + // (GUI thread) never races on the pixel buffer. + cv::Mat preview = m_IRToolTracker->GetProcessedFrame(); + { + std::lock_guard lock(mtx_frames); + trackingFrame = preview; + } } - - trackingFrame = m_IRToolTracker->GetProcessedFrame(); k4a_image_release(depth_image); k4a_image_release(ir_image); k4a_capture_release(capture); @@ -224,18 +251,30 @@ void IRToolTracking::processStreams() { std::this_thread::sleep_for(std::chrono::milliseconds(33)); } } + } catch (const std::exception &e) { + std::cerr << "K4A streaming stopped: " << e.what() << std::endl; + Terminated = true; + } } void IRToolTracking::shutdown() { - // Clean up resources as necessary - if (device != NULL) - { - k4a_device_close(device); + // Clean up resources as necessary. Null the handles after closing so a second + // shutdown() (ViewerWindow now always calls it, plus the destructor) cannot + // double-close, and guard the teardown so a disconnected device can't throw here. + try { + if (device != NULL) + { + k4a_device_close(device); + device = NULL; + } + if (playback != NULL) + { + k4a_playback_close(playback); + playback = NULL; + } + } catch (const std::exception &e) { + std::cerr << "Error during K4A shutdown: " << e.what() << std::endl; } - if (playback != NULL) - { - k4a_playback_close(playback); - } trackingFrame.release(); depthFrame.release(); } diff --git a/src/ViewerWindow.cpp b/src/ViewerWindow.cpp index 50a9e8d..5066e4f 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.at("numSpheres").get(); + tool.spherePositions = toolJson.at("spherePositions").get>(); + tool.sphereRadius = toolJson.at("sphereRadius").get(); + tool.toolName = toolJson.at("toolName").get(); - Tool tool; - tool.numSpheres = toolJson["numSpheres"].get(); - tool.spherePositions = toolJson["spherePositions"].get>(); - tool.sphereRadius = toolJson["sphereRadius"].get(); - tool.toolName = toolJson["toolName"].get(); + // 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; + } - if (!file) + 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,20 +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) + if (_socket < 0) { std::cerr << "Failed to create a socket." << std::endl; + ReleaseSocketInit(); return false; } @@ -212,6 +251,8 @@ bool ViewerWindow::Connect(NanoSocket& _socket, NanoAddress& address, const char if (nanosockets_address_set_ip(&address, "127.0.0.1")) { 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)) @@ -268,12 +320,20 @@ void ViewerWindow::UdpReceiveThreadFunction() } + std::vector buffer(sizeof(TrackingData)); while (multiEnabled) { + // Wait up to 50 ms for a packet so the loop doesn't burn a core when idle. + const int ready = nanosockets_poll(receiveSocket, 50); + if (ready <= 0) { + continue; + } + NanoAddress sender; - std::vector buffer(sizeof(TrackingData)); - //toolTransforms.clear(); - 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)); @@ -285,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 @@ -338,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; @@ -356,15 +417,17 @@ 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; - std::copy(tool_transform.begin(), tool_transform.begin() + 3, data.position); + // 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; data.timestamp = GetCurrentUnixTimestamp(); @@ -385,7 +448,7 @@ void ViewerWindow::UdpThreadFunction() } } - int sleepDurationMs = 1000 / frequency; // Convert frequency to sleep duration in milliseconds + int sleepDurationMs = 1000 / std::max(frequency, 1); std::this_thread::sleep_for(std::chrono::milliseconds(sleepDurationMs)); } @@ -438,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(); @@ -459,7 +522,7 @@ void ViewerWindow::WriteToCSV() } } - int sleepDurationMs = 1000 / recordFrequency; + int sleepDurationMs = 1000 / std::max(recordFrequency, 1); std::this_thread::sleep_for(std::chrono::milliseconds(sleepDurationMs)); } @@ -558,8 +621,26 @@ 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 + // doesn't keep matching ghosts. + for (int i = numTools; i < static_cast(tools.size()); ++i) + { + if (tools[i].isAdded) + { + tracker.RemoveToolDefinition(tools[i].toolName); + } + } + } if (static_cast(tools.size()) != numTools) { tools.resize(numTools); @@ -569,6 +650,7 @@ void ViewerWindow::Render() { ImGui::SetNextWindowSize(ImVec2(windowWidth, 0.0f)); ImGui::Begin("Tool Definitions", nullptr, overlayFlags); + isToolAdded = false; for (int toolIdx = 0; toolIdx < numTools; ++toolIdx) { tools[toolIdx].toolName = tools[toolIdx].toolName == "Tool" ? "Tool" + std::to_string(toolIdx + 1) : tools[toolIdx].toolName; @@ -585,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) { @@ -600,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 @@ -620,9 +705,21 @@ void ViewerWindow::Render() { { 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); @@ -637,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())) @@ -647,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(); @@ -659,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) + { + 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 = tool_definition.size() / 3; + 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 the capture loop 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(); @@ -679,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); @@ -733,18 +881,24 @@ void ViewerWindow::Render() { ImGui::SameLine(); 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(); - if (ImGui::Checkbox("UDP", &udpEnabled)) { - if (udpEnabled) - { - udpThread = std::make_shared(&ViewerWindow::UdpThreadFunction, this); - } - else + bool udpEnabledTmp = udpEnabled.load(); + if (ImGui::Checkbox("UDP", &udpEnabledTmp)) { - JoinThread(udpThread); + udpEnabled.store(udpEnabledTmp); + if (udpEnabledTmp) + { + udpThread = std::make_shared(&ViewerWindow::UdpThreadFunction, this); + } + else + { + JoinThread(udpThread); + } } } ImGui::End(); @@ -753,20 +907,25 @@ void ViewerWindow::Render() { ImGui::SetNextWindowSize(ImVec2(300, 0.0f)); ImGui::SetNextWindowPos(ImVec2(740, 80), ImGuiCond_FirstUseEver); ImGui::Begin("Multi-Camera Settings", nullptr, overlayFlags); - if (ImGui::Checkbox("Multi-Camera", &multiEnabled)) { - if (multiEnabled) - { - udpReceiveThread = std::make_shared(&ViewerWindow::UdpReceiveThreadFunction, this); - } - else + bool multiEnabledTmp = multiEnabled.load(); + if (ImGui::Checkbox("Multi-Camera", &multiEnabledTmp)) { - JoinThread(udpReceiveThread); + multiEnabled.store(multiEnabledTmp); + if (multiEnabledTmp) + { + udpReceiveThread = std::make_shared(&ViewerWindow::UdpReceiveThreadFunction, this); + } + else + { + JoinThread(udpReceiveThread); + } } } 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 @@ -792,17 +951,21 @@ void ViewerWindow::Render() { NFD_Quit(); } ImGui::SameLine(); - if (ImGui::Checkbox("Record", &csvEnabled)) { - if (csvEnabled) + bool csvEnabledTmp = csvEnabled.load(); + if (ImGui::Checkbox("Record", &csvEnabledTmp)) { - csvThread = std::make_shared(&ViewerWindow::WriteToCSV, this); - } - else - { - JoinThread(csvThread); - } - } + csvEnabled.store(csvEnabledTmp); + if (csvEnabledTmp) + { + csvThread = std::make_shared(&ViewerWindow::WriteToCSV, this); + } + else + { + JoinThread(csvThread); + } + } + } ImGui::End(); recordFrequency = std::min(std::max(recordFrequency, 1), 90); duration = std::max(duration, 1); @@ -864,7 +1027,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; } } @@ -899,6 +1063,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() ); @@ -923,11 +1110,17 @@ void ViewerWindow::Shutdown() { Terminated = true; udpEnabled = false; multiEnabled = false; + csvEnabled = false; + tracker.StopToolTracking(); + tracker.StopToolCalibration(); JoinThread(processingThread); JoinThread(udpThread); JoinThread(udpReceiveThread); JoinThread(csvThread); - if (tracker.IsTrackingTools() || tracker.IsCalibratingTool()) - tracker.shutdown(); + + // shutdown() is idempotent (k4a handles are nulled after close), 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 device. + tracker.shutdown(); } \ No newline at end of file