mirror of
https://github.com/slsdetectorgroup/aare.git
synced 2026-10-31 01:11:30 +01:00
Erik Fröjdh
bf0aa4a3bf
- Python extension. Originally left out since it didn't matter but good practice for future - to_string.hpp - this is the important one
229 lines
8.7 KiB
C++
229 lines
8.7 KiB
C++
// SPDX-License-Identifier: MPL-2.0
|
|
#pragma once
|
|
#include "aare/ClusterCollector.hpp"
|
|
#include "aare/ClusterFileSink.hpp"
|
|
#include "aare/ClusterFinder.hpp"
|
|
#include "aare/ClusterFinderMT.hpp"
|
|
#include "aare/ClusterVector.hpp"
|
|
#include "aare/NDView.hpp"
|
|
#include "aare/Pedestal.hpp"
|
|
#include "np_helper.hpp"
|
|
|
|
#include <cstdint>
|
|
#include <filesystem>
|
|
#include <pybind11/pybind11.h>
|
|
#include <pybind11/stl.h>
|
|
#include <pybind11/stl_bind.h>
|
|
|
|
namespace py = pybind11;
|
|
|
|
using namespace aare;
|
|
|
|
template <typename Type, uint8_t ClusterSizeX, uint8_t ClusterSizeY,
|
|
typename CoordType = uint16_t>
|
|
void define_ClusterVector(py::module &m, const std::string &typestr) {
|
|
using ClusterType = Cluster<Type, ClusterSizeX, ClusterSizeY, CoordType>;
|
|
auto class_name = fmt::format("ClusterVector_{}", typestr);
|
|
|
|
py::class_<ClusterVector<
|
|
Cluster<Type, ClusterSizeX, ClusterSizeY, CoordType>, void>>(
|
|
m, class_name.c_str(), R"doc(
|
|
A contiguous, move-only container of fixed-size clusters.
|
|
|
|
The class supports the Python buffer protocol, so ``numpy.array`` can
|
|
either copy its data or create a zero-copy view. A zero-copy view is
|
|
valid only while the ClusterVector's underlying allocation and size
|
|
remain unchanged.
|
|
)doc",
|
|
py::buffer_protocol())
|
|
|
|
.def(py::init(), R"doc(
|
|
Create an empty ClusterVector with frame number 0 and space reserved
|
|
for at least 1024 clusters.
|
|
)doc")
|
|
|
|
.def(
|
|
"__call__",
|
|
[](ClusterVector<ClusterType> &self, py::array_t<bool> mask) {
|
|
return self(make_view_1d(mask));
|
|
},
|
|
py::arg("mask").noconvert(), R"doc(
|
|
Return a filtered copy of this ClusterVector.
|
|
|
|
Parameters
|
|
----------
|
|
mask : numpy.ndarray
|
|
One-dimensional, writable, C-contiguous array with dtype
|
|
``numpy.bool_`` and one element per cluster.
|
|
|
|
Returns
|
|
-------
|
|
ClusterVector
|
|
Selected clusters in their original order. The frame number is
|
|
preserved.
|
|
)doc")
|
|
|
|
.def(
|
|
"push_back",
|
|
[](ClusterVector<ClusterType> &self, const ClusterType &cluster) {
|
|
self.push_back(cluster);
|
|
},
|
|
py::arg("cluster"), R"doc(
|
|
Append one cluster.
|
|
|
|
Notes
|
|
-----
|
|
Do not call this method while a zero-copy NumPy view of the
|
|
ClusterVector exists. Reallocation invalidates the view, while an
|
|
append without reallocation leaves its shape unchanged.
|
|
)doc")
|
|
|
|
.def(
|
|
"sum",
|
|
[](ClusterVector<ClusterType> &self) {
|
|
auto *vec = new std::vector<Type>(self.sum());
|
|
return return_vector(vec);
|
|
},
|
|
R"doc(
|
|
Return the sum of all pixels in each cluster.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
One value per cluster, in container order and with the cluster
|
|
pixel dtype.
|
|
)doc")
|
|
.def(
|
|
"sum_2x2",
|
|
[](ClusterVector<ClusterType> &self) {
|
|
auto *vec = new std::vector<Sum_index_pair<Type, corner>>(
|
|
self.sum_2x2());
|
|
|
|
return return_vector(vec);
|
|
},
|
|
R"doc(
|
|
Return the highest-sum center-adjacent 2x2 subcluster for each
|
|
cluster.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
Structured array with ``sum`` and ``index`` fields. Indices are
|
|
0 for top-left, 1 for top-right, 2 for bottom-left, and 3 for
|
|
bottom-right, relative to the cluster center.
|
|
)doc")
|
|
.def_property_readonly("size", &ClusterVector<ClusterType>::size,
|
|
"Number of stored clusters.")
|
|
.def("empty", &ClusterVector<ClusterType>::empty,
|
|
"Return True when no clusters are stored.")
|
|
.def("item_size", &ClusterVector<ClusterType>::item_size,
|
|
"Return the size in bytes of one stored cluster, including "
|
|
"padding.")
|
|
.def_property_readonly(
|
|
"fmt",
|
|
[typestr]([[maybe_unused]] ClusterVector<ClusterType> &self) {
|
|
return fmt_format<ClusterType>();
|
|
},
|
|
"PEP 3118 format string for one stored cluster.")
|
|
|
|
.def_property_readonly("cluster_size_x",
|
|
&ClusterVector<ClusterType>::cluster_size_x,
|
|
"Cluster size in the x dimension.")
|
|
.def_property_readonly("cluster_size_y",
|
|
&ClusterVector<ClusterType>::cluster_size_y,
|
|
"Cluster size in the y dimension.")
|
|
.def_property_readonly("capacity",
|
|
&ClusterVector<ClusterType>::capacity,
|
|
"Number of clusters that fit without "
|
|
"reallocation.")
|
|
.def_property("frame_number", &ClusterVector<ClusterType>::frame_number,
|
|
&ClusterVector<ClusterType>::set_frame_number,
|
|
"Signed 32-bit frame number; 0 can indicate clusters "
|
|
"from multiple frames.")
|
|
.def_buffer(
|
|
[typestr](ClusterVector<ClusterType> &self) -> py::buffer_info {
|
|
return py::buffer_info(
|
|
self.data(), /* Pointer to buffer */
|
|
self.item_size(), /* Size of one scalar */
|
|
fmt_format<ClusterType>(), /* Format descriptor */
|
|
1, /* Number of dimensions */
|
|
{self.size()}, /* Buffer dimensions */
|
|
{self.item_size()} /* Strides (in bytes) for each index */
|
|
);
|
|
});
|
|
|
|
// Free functions using ClusterVector
|
|
m.def(
|
|
"hitmap",
|
|
[](std::array<size_t, 2> image_size, ClusterVector<ClusterType> &cv) {
|
|
// Create a numpy array to hold the hitmap
|
|
// The shape of the array is (image_size[0], image_size[1])
|
|
// note that the python array is passed as [row, col] which
|
|
// is the opposite of the clusters [x,y]
|
|
py::array_t<int32_t> hitmap(image_size);
|
|
auto r = hitmap.mutable_unchecked<2>();
|
|
|
|
// Initialize hitmap to 0
|
|
for (py::ssize_t i = 0; i < r.shape(0); i++)
|
|
for (py::ssize_t j = 0; j < r.shape(1); j++)
|
|
r(i, j) = 0;
|
|
|
|
// Loop over the clusters and increment the hitmap
|
|
// Skip out of bound clusters
|
|
for (const auto &cluster : cv) {
|
|
auto x = cluster.x;
|
|
auto y = cluster.y;
|
|
if (x < image_size[1] && y < image_size[0])
|
|
r(cluster.y, cluster.x) += 1;
|
|
}
|
|
|
|
return hitmap;
|
|
},
|
|
py::arg("image_size"), py::arg("clusters"), R"doc(
|
|
Count cluster centers into an ``int32`` image whose shape is given by
|
|
``image_size`` as ``(rows, columns)``. Element ``[y, x]`` contains the
|
|
number of cluster centers at that coordinate. Out-of-bounds centers are
|
|
ignored.
|
|
)doc");
|
|
}
|
|
|
|
template <typename Type, uint8_t ClusterSizeX, uint8_t ClusterSizeY,
|
|
typename CoordType = uint16_t>
|
|
void define_2x2_reduction(py::module &m) {
|
|
m.def(
|
|
"reduce_to_2x2",
|
|
[](const ClusterVector<
|
|
Cluster<Type, ClusterSizeX, ClusterSizeY, CoordType>> &cv) {
|
|
return new ClusterVector<Cluster<Type, 2, 2, CoordType>>(
|
|
reduce_to_2x2(cv));
|
|
},
|
|
R"doc(
|
|
Reduce every cluster to its highest-sum center-adjacent 2x2 block.
|
|
Returns a new ClusterVector; cluster order, coordinates, and frame
|
|
number are preserved. Input pixel data is interpreted in row-major
|
|
order.
|
|
)doc",
|
|
py::arg("clustervector"));
|
|
}
|
|
|
|
template <typename Type, uint8_t ClusterSizeX, uint8_t ClusterSizeY,
|
|
typename CoordType = uint16_t>
|
|
void define_3x3_reduction(py::module &m) {
|
|
|
|
m.def(
|
|
"reduce_to_3x3",
|
|
[](const ClusterVector<
|
|
Cluster<Type, ClusterSizeX, ClusterSizeY, CoordType>> &cv) {
|
|
return new ClusterVector<Cluster<Type, 3, 3, CoordType>>(
|
|
reduce_to_3x3(cv));
|
|
},
|
|
R"doc(
|
|
Reduce every cluster to the 3x3 block around its center index.
|
|
Both input dimensions must be at least 3, and at least one must be
|
|
greater than 3.
|
|
Returns a new ClusterVector; cluster order, coordinates, and frame
|
|
number are preserved.
|
|
)doc",
|
|
py::arg("clustervector"));
|
|
}
|