diff --git a/image_analysis/RotationSpotAccumulator.cpp b/image_analysis/RotationSpotAccumulator.cpp index 2ac876ac7..39324b013 100644 --- a/image_analysis/RotationSpotAccumulator.cpp +++ b/image_analysis/RotationSpotAccumulator.cpp @@ -2,4 +2,114 @@ // Created by jungfrau on 2/4/26. // -#include "RotationSpotAccumulator.h" \ No newline at end of file +#include "RotationSpotAccumulator.h" + +RotationSpot3D::RotationSpot3D(SpotToSave s, float phi_deg, int64_t image) { + x = s.x * s.intensity; + y = s.y * s.intensity; + phi = phi_deg * s.intensity; + intensity = s.intensity; + maxc = s.maxc; + first_image = image; + last_image = image; +} + +void RotationSpot3D::Add(const SpotToSave& s, float phi_deg, int64_t image) { + x += s.x * s.intensity; + y += s.y * s.intensity; + phi += phi_deg * s.intensity; + intensity += s.intensity; + maxc = std::max(maxc, s.maxc); + ice_ring = ice_ring || s.ice_ring; + first_image = std::min(first_image, image); + last_image = std::max(last_image, image); +} + +float RotationSpot3D::calcX() const { + return x / intensity; +} + +float RotationSpot3D::calcY() const { + return y / intensity; +} + +float RotationSpot3D::calcPhi() const { + return phi / intensity; +} + +float RotationSpot3D::calcI() const { + return intensity; +} + +int64_t RotationSpot3D::calcMAXC() const { + return maxc; +} + +int64_t RotationSpot3D::getLastImage() const { + return last_image; +} + +void RotationSpotAccumulator::AddImage(int64_t image, const std::vector& in_spots) { + std::lock_guard lock(mutex_); + spots[image] = in_spots; +} + + +std::vector RotationSpotAccumulator::FinalizeAll() { + std::lock_guard lock(mutex_); + + int64_t image0 = INT64_MIN; + + std::vector completed_; + std::vector pending_; + + for (auto& [image, v]: spots) { + float image_angle = axis_.GetIncrement_deg() * image + axis_.GetWedge_deg() / 2.0f; + + std::vector tmp; + + if (image == image0 + 1) { + for (const auto &p: pending_) { + if (p.getLastImage() != image - 1) { + completed_.push_back(p); + } else { + tmp.push_back(p); + } + } + + for (const auto &s: v) { + bool matched = false; + + for (int i = 0; i < tmp.size(); i++) { + float dist_x = s.x - tmp[i].calcX(); + float dist_y = s.y - tmp[i].calcY(); + + if (std::fabs(dist_x * dist_x + dist_y * dist_y) < config_.xy_tolerance_pxl_sq) { + tmp[i].Add(s, image_angle, image); + matched = true; + break; + } + } + + if (!matched) + tmp.emplace_back(s, image_angle, image); + } + + + } else { + for (const auto &p: pending_) + completed_.push_back(p); + + for (const auto &s: v) + tmp.emplace_back(s, image_angle, image); + + } + pending_ = tmp; + image0 = image; + } + + for (const auto &p: pending_) + completed_.push_back(p); + + return completed_; +} diff --git a/image_analysis/RotationSpotAccumulator.h b/image_analysis/RotationSpotAccumulator.h index 66af67b92..72d2239f0 100644 --- a/image_analysis/RotationSpotAccumulator.h +++ b/image_analysis/RotationSpotAccumulator.h @@ -17,85 +17,33 @@ #include "../common/DiffractionGeometry.h" // Forward declaration -struct RotationSpot3D; - -// Spatial grid cell key -struct GridCell { - int32_t cx, cy; - - bool operator==(const GridCell& o) const { - return cx == o.cx && cy == o.cy; - } -}; - -struct GridCellHash { - size_t operator()(const GridCell& c) const { - return std::hash{}((static_cast(c.cx) << 32) | - static_cast(c.cy)); - } -}; - -// A pending spot that may still be merged -struct PendingSpot { - SpotToSave spot; - size_t merged_spot_id; // ID of the RotationSpot3D it belongs to (or SIZE_MAX if new) -}; - -// Accumulated 3D spot data -struct RotationSpot3D { - // Intensity-weighted centroid - float x = 0, y = 0; - float phi_centroid_deg = 0; +class RotationSpot3D { + float x = 0; + float y = 0; + float phi = 0; float intensity = 0; int64_t maxc = 0; - float d_A = 0; + bool ice_ring = false; int64_t first_image = INT64_MAX; int64_t last_image = INT64_MIN; - bool ice_ring = false; +public: + RotationSpot3D(SpotToSave s, float phi_deg, int64_t image); + void Add(const SpotToSave& s, float phi_deg, int64_t image); - // For weighted averaging - double sum_x_I = 0, sum_y_I = 0, sum_phi_I = 0; - - // Images that contributed to this spot - std::unordered_set contributing_images; - - void AddObservation(const SpotToSave& s, float phi_deg) { - sum_x_I += s.x * s.intensity; - sum_y_I += s.y * s.intensity; - sum_phi_I += phi_deg * s.intensity; - intensity += s.intensity; - - if (intensity > 0) { - x = static_cast(sum_x_I / intensity); - y = static_cast(sum_y_I / intensity); - phi_centroid_deg = static_cast(sum_phi_I / intensity); - } - - maxc = std::max(maxc, s.maxc); - d_A = s.d_A; // Will be recalculated at finalization - ice_ring = ice_ring || s.ice_ring; - - first_image = std::min(first_image, s.image); - last_image = std::max(last_image, s.image); - contributing_images.insert(s.image); - } - - bool IsComplete(int64_t min_confirmed_image, int64_t max_image_gap) const { - // A spot is complete when we've received all images that could - // possibly contribute (i.e., gap after last_image is confirmed) - return min_confirmed_image > last_image + max_image_gap; - } + float calcX() const; + float calcY() const; + float calcPhi() const; + float calcI() const; + int64_t calcMAXC() const; + int64_t getLastImage() const; }; class RotationSpotAccumulator { public: struct Config { - float xy_tolerance_pxl = 2.0f; - int64_t max_image_gap = 1; - float grid_cell_size = 10.0f; // pixels - int64_t window_size = 10; // images to keep in pending buffer + float xy_tolerance_pxl_sq = 2.0f * 2.0f; }; private: @@ -105,226 +53,17 @@ private: const GoniometerAxis& axis_; Config config_; - // Per-image spatial grid of pending spots - // image_number -> (grid_cell -> list of spot indices in pending_spots_) - std::map, GridCellHash>> image_grids_; - - // All pending spots (indexed by ID) - std::vector pending_spots_; - std::vector free_list_; // Recycled IDs - - // Accumulated 3D spots (indexed by ID) - std::vector spots_3d_; - std::vector free_3d_list_; - - // Track which images we've seen - std::unordered_set received_images_; - int64_t min_received_image_ = INT64_MAX; - int64_t max_received_image_ = INT64_MIN; - - // Completed spots ready for output - std::vector completed_; - - GridCell GetGridCell(float x, float y) const { - return GridCell{ - static_cast(std::floor(x / config_.grid_cell_size)), - static_cast(std::floor(y / config_.grid_cell_size)) - }; - } - - // Get neighboring cells (3x3 neighborhood) - std::vector GetNeighborCells(const GridCell& c) const { - std::vector neighbors; - neighbors.reserve(9); - for (int dx = -1; dx <= 1; ++dx) { - for (int dy = -1; dy <= 1; ++dy) { - neighbors.push_back({c.cx + dx, c.cy + dy}); - } - } - return neighbors; - } - - size_t AllocatePendingSpot() { - if (!free_list_.empty()) { - size_t id = free_list_.back(); - free_list_.pop_back(); - return id; - } - size_t id = pending_spots_.size(); - pending_spots_.emplace_back(); - return id; - } - - size_t Allocate3DSpot() { - if (!free_3d_list_.empty()) { - size_t id = free_3d_list_.back(); - free_3d_list_.pop_back(); - spots_3d_[id] = RotationSpot3D{}; - return id; - } - size_t id = spots_3d_.size(); - spots_3d_.emplace_back(); - return id; - } - - float GetPhiCenter(int64_t image) const { - return axis_.GetAngle_deg(image) + axis_.GetWedge_deg() / 2.0f; - } - - // Find a matching 3D spot in adjacent images - std::optional FindMatch(const SpotToSave& spot) { - GridCell cell = GetGridCell(spot.x, spot.y); - auto neighbors = GetNeighborCells(cell); - - float best_dist_sq = config_.xy_tolerance_pxl * config_.xy_tolerance_pxl; - std::optional best_match; - - // Search adjacent images - for (int64_t delta = -config_.max_image_gap; delta <= config_.max_image_gap; ++delta) { - if (delta == 0) continue; // Skip same image - - int64_t adj_image = spot.image + delta; - auto grid_it = image_grids_.find(adj_image); - if (grid_it == image_grids_.end()) continue; - - for (const auto& nc : neighbors) { - auto cell_it = grid_it->second.find(nc); - if (cell_it == grid_it->second.end()) continue; - - for (size_t pending_id : cell_it->second) { - const auto& ps = pending_spots_[pending_id]; - float dx = ps.spot.x - spot.x; - float dy = ps.spot.y - spot.y; - float dist_sq = dx * dx + dy * dy; - - if (dist_sq < best_dist_sq) { - best_dist_sq = dist_sq; - best_match = ps.merged_spot_id; - } - } - } - } - - return best_match; - } - - void FlushOldImages(int64_t current_min_image) { - // Remove images that are too old to merge - int64_t cutoff = current_min_image - config_.window_size; - - std::vector to_remove; - for (auto& [img, grid] : image_grids_) { - if (img < cutoff) { - to_remove.push_back(img); - // Free pending spots from this image - for (auto& [cell, ids] : grid) { - for (size_t id : ids) { - free_list_.push_back(id); - } - } - } - } - - for (int64_t img : to_remove) { - image_grids_.erase(img); - } - - // Check which 3D spots are now complete - for (size_t i = 0; i < spots_3d_.size(); ++i) { - auto& s = spots_3d_[i]; - if (s.intensity > 0 && s.IsComplete(cutoff, config_.max_image_gap)) { - // Finalize resolution at centroid - s.d_A = geom_.PxlToRes(s.x, s.y); - completed_.push_back(std::move(s)); - s = RotationSpot3D{}; // Clear - free_3d_list_.push_back(i); - } - } - } + std::map> spots; public: RotationSpotAccumulator(const DiffractionGeometry& geom, const GoniometerAxis& axis, - Config config = {}) + Config config) : geom_(geom), axis_(axis), config_(config) {} // Add spots from one image (thread-safe, handles out-of-order) - void AddImage(int64_t image, const std::vector& spots) { - std::lock_guard lock(mutex_); - - // Track received images - received_images_.insert(image); - min_received_image_ = std::min(min_received_image_, image); - max_received_image_ = std::max(max_received_image_, image); - - float phi_center = GetPhiCenter(image); - - // Create grid for this image if needed - auto& grid = image_grids_[image]; - - for (const auto& spot : spots) { - // Look for match in adjacent images - auto match = FindMatch(spot); - - size_t spot_3d_id; - if (match.has_value()) { - spot_3d_id = *match; - } else { - spot_3d_id = Allocate3DSpot(); - } - - // Add observation to 3D spot - spots_3d_[spot_3d_id].AddObservation(spot, phi_center); - - // Store pending spot for future matching - size_t pending_id = AllocatePendingSpot(); - pending_spots_[pending_id] = PendingSpot{ - .spot = spot, - .merged_spot_id = spot_3d_id - }; - - // Add to spatial grid - GridCell cell = GetGridCell(spot.x, spot.y); - grid[cell].push_back(pending_id); - } - - // Periodically flush old data - if (received_images_.size() % 10 == 0) { - FlushOldImages(min_received_image_); - } - } - - // Get completed spots (empties the completed buffer) - std::vector GetCompleted() { - std::lock_guard lock(mutex_); - return std::move(completed_); - } + void AddImage(int64_t image, const std::vector& in_spots); // Finalize all remaining spots (call at end of dataset) - std::vector FinalizeAll() { - std::lock_guard lock(mutex_); - - std::vector result = std::move(completed_); - - for (auto& s : spots_3d_) { - if (s.intensity > 0) { - s.d_A = geom_.PxlToRes(s.x, s.y); - result.push_back(std::move(s)); - } - } - - // Clear state - spots_3d_.clear(); - pending_spots_.clear(); - image_grids_.clear(); - free_list_.clear(); - free_3d_list_.clear(); - - return result; - } - - size_t PendingCount() const { - std::lock_guard lock(mutex_); - return spots_3d_.size() - free_3d_list_.size(); - } + std::vector FinalizeAll(); }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 248d7daca..79a6b38ba 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -64,6 +64,7 @@ ADD_EXECUTABLE(jfjoch_test RotationIndexerTest.cpp TopPixelsTest.cpp HKLKeyTest.cpp + RotationScanAccumulatorTest.cpp ) target_link_libraries(jfjoch_test Catch2WithMain JFJochBroker JFJochReceiver JFJochReader JFJochWriter JFJochImageAnalysis JFJochCommon JFJochHLSSimulation JFJochPreview) diff --git a/tests/RotationScanAccumulatorTest.cpp b/tests/RotationScanAccumulatorTest.cpp new file mode 100644 index 000000000..d8f10fd20 --- /dev/null +++ b/tests/RotationScanAccumulatorTest.cpp @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#include +#include "../image_analysis//RotationSpotAccumulator.h" + +TEST_CASE("RotationSpotAccumulator") { + DiffractionGeometry geometry; + geometry.BeamX_pxl(1000).BeamY_pxl(1000).DetectorDistance_mm(100); + + GoniometerAxis axis("U", 0, 0.1, Coord{1, 0, 0}, std::nullopt); + + RotationSpotAccumulator::Config config{}; + RotationSpotAccumulator accumulator(geometry, axis, config); + + accumulator.AddImage(2, { + SpotToSave{.x = 500, .y = 500, .intensity = 100}, + SpotToSave{.x = 101, .y = 101, .intensity = 50} + }); + + accumulator.AddImage(0, { + SpotToSave{.x = 100, .y = 100, .intensity = 50}, + SpotToSave{.x = 200, .y = 200, .intensity = 50} + }); + + accumulator.AddImage(1, { + SpotToSave{.x = 500, .y = 500, .intensity = 100}, + SpotToSave{.x = 101, .y = 100, .intensity = 100} + }); + + accumulator.AddImage(4, { + SpotToSave{.x = 500, .y = 500, .intensity = 100} + }); + + auto v = accumulator.FinalizeAll(); + + struct Expect { + float x, y, phi, I; + }; + + std::vector expected = { + {200.0f, 200.0f, 0.05f, 50.0f}, + {500.0f, 500.0f, 0.2f, 200.0f}, + {100.75f, 100.25f, 0.15f, 200.0f}, + {500.0f, 500.0f, 0.45f, 100.0f} + }; + + auto key = [](const Expect& e) { return std::pair(e.x, e.y); }; + + std::vector got; + got.reserve(v.size()); + for (const auto& s : v) { + got.push_back({s.calcX(), s.calcY(), s.calcPhi(), s.calcI()}); + } + + std::sort(expected.begin(), expected.end(), + [&](const Expect& a, const Expect& b){ return key(a) < key(b); }); + std::sort(got.begin(), got.end(), + [&](const Expect& a, const Expect& b){ return key(a) < key(b); }); + + REQUIRE(got.size() == expected.size()); + for (size_t i = 0; i < got.size(); ++i) { + CHECK(got[i].x == Catch::Approx(expected[i].x)); + CHECK(got[i].y == Catch::Approx(expected[i].y)); + CHECK(got[i].phi == Catch::Approx(expected[i].phi)); + CHECK(got[i].I == Catch::Approx(expected[i].I)); + } +}