Added some extra functionality to lazy array

.github/workflows/codecov.yml

@@ -58,7 +58,7 @@ jobs:
  python -m pip install .
  - name: Run coverage
  run: |
- coverage run -m pytest || true 
+ coverage run -m pytest || true
  - name: Upload Coverage to Codecov
  uses: codecov/codecov-action@v4
  with:

.spin/cmds.py

@@ -1,5 +1,4 @@
 import os
-import site
 import subprocess
 import webbrowser
 from pathlib import Path

examples/annotations_to_mask.py

@@ -2,6 +2,7 @@
 """This code provides an example of how to convert annotations to a mask."""
 import json
 from pathlib import Path
+
 import numpy as np
 import PIL.Image
 
@@ -48,17 +49,8 @@
 print(curr_mask)
 print(np.asarray(curr_mask).shape)
 
-import time
-start_time = time.time()
-# Let's grab the polygons
-
 mask = LUT[region.polygons.to_mask().numpy()]
 
-print(f"Time lazy: {time.time() - start_time}")
-
-start_time = time.time()
-mask = LUT[region.polygons.to_mask_no_lazy()]
-print(f"Time eager: {time.time() - start_time}")
 
 PIL.Image.fromarray(mask).save("mask.png")
 

src/geometry.cpp

@@ -4,10 +4,10 @@
 
 #include "geometry/base.h"
 #include "geometry/box.h"
-#include "geometry/lazy_array.h"
 #include "geometry/collection.h"
 #include "geometry/exceptions.h"
 #include "geometry/factory.h"
+#include "geometry/lazy_array.h"
 #include "geometry/point.h"
 #include "geometry/polygon.h"
 #include "geometry/region.h"
@@ -150,9 +150,7 @@ PYBIND11_MODULE(_geometry, m) {
 
  py::class_<PolygonCollection, std::shared_ptr<PolygonCollection>>(m, "PolygonCollection")
  .def("get_geometries", &PolygonCollection::getGeometries)
- .def("to_mask", &PolygonCollection::toMask, py::arg("default_value") = 0)
- .def("to_mask_no_lazy", &PolygonCollection::toMaskNonLazy, py::arg("default_value") = 0);
-
+ .def("to_mask", &PolygonCollection::toMask, py::arg("default_value") = 0);
 
  py::class_<AnnotationRegion, std::shared_ptr<AnnotationRegion>>(m, "AnnotationRegion")
  .def_property_readonly("polygons", &AnnotationRegion::getPolygons)

src/geometry/lazy_array.h

@@ -14,34 +14,78 @@ class LazyArray {
  public:
  using ComputeFunction = std::function<py::array_t<T>()>;
 
- LazyArray(ComputeFunction compute_func) : compute_func_(std::move(compute_func)), computed_(false) {}
+ LazyArray(ComputeFunction compute_func, std::vector<py::ssize_t> shape)
+ : compute_func_(std::move(compute_func)), computed_(false), shape_(std::move(shape)) {}
 
- py::array_t<T> numpy() {
+ py::array_t<T> numpy() const {
  if (!computed_) {
  data_ = compute_func_();
  computed_ = true;
  }
  return data_;
  }
 
- py::array_t<T> operator*() { return numpy(); }
+ py::array_t<T> operator*() const { return numpy(); }
 
- // Changed this method to return py::array_t<T> directly
- py::array_t<T> py_numpy() { return numpy(); }
+ py::array_t<T> py_numpy() const { return numpy(); }
+
+ std::vector<py::ssize_t> shape() const { return shape_; }
+
+ LazyArray<T> transpose(const std::vector<py::ssize_t> &axes = {}) const {
+ std::vector<py::ssize_t> new_shape(shape_);
+ if (axes.empty()) {
+ std::reverse(new_shape.begin(), new_shape.end());
+ } else {
+ for (size_t i = 0; i < axes.size(); ++i) {
+ new_shape[i] = shape_[axes[i]];
+ }
+ }
+
+ return LazyArray<T>(
+ [this, axes]() { return this->numpy().attr("transpose")(axes).template cast<py::array_t<T>>(); }, new_shape);
+ }
+
+ LazyArray<T> reshape(const std::vector<py::ssize_t> &new_shape) const {
+ return LazyArray<T>([this, new_shape]() { return this->numpy().reshape(new_shape); }, new_shape);
+ }
+
+ LazyArray<T> operator+(const LazyArray<T> &other) const {
+ return LazyArray<T>([this, &other]() { return this->numpy() + other.numpy(); }, shape_);
+ }
+
+ LazyArray<T> operator-(const LazyArray<T> &other) const {
+ return LazyArray<T>([this, &other]() { return this->numpy() - other.numpy(); }, shape_);
+ }
+
+ LazyArray<T> multiply(const LazyArray<T> &other) const {
+ return LazyArray<T>([this, &other]() { return this->numpy() * other.numpy(); }, shape_);
+ }
+
+ LazyArray<T> operator/(const LazyArray<T> &other) const {
+ return LazyArray<T>([this, &other]() { return this->numpy() / other.numpy(); }, shape_);
+ }
 
  private:
  ComputeFunction compute_func_;
- py::array_t<T> data_;
- bool computed_;
+ mutable py::array_t<T> data_;
+ mutable bool computed_;
+ std::vector<py::ssize_t> shape_;
 };
 
 template <typename T>
 void declare_lazy_array(py::module &m, const std::string &type_name) {
  py::class_<LazyArray<T>>(m, type_name.c_str())
- .def(py::init<typename LazyArray<T>::ComputeFunction>())
+ .def(py::init<typename LazyArray<T>::ComputeFunction, std::vector<py::ssize_t>>())
  .def("numpy", &LazyArray<T>::py_numpy)
+ .def("shape", &LazyArray<T>::shape)
+ .def("transpose", &LazyArray<T>::transpose, py::arg("axes") = std::vector<py::ssize_t>())
+ .def("reshape", &LazyArray<T>::reshape)
  .def("__array__", &LazyArray<T>::py_numpy)
- .def("__repr__", [](const LazyArray<T> &) { return "<LazyArray: use numpy() or __array__() to compute>"; });
+ .def("__repr__", [](const LazyArray<T> &) { return "<LazyArray: use numpy() or __array__() to compute>"; })
+ .def("__add__", &LazyArray<T>::operator+)
+ .def("__sub__", &LazyArray<T>::operator-)
+ .def("__mul__", &LazyArray<T>::multiply)
+ .def("__truediv__", &LazyArray<T>::operator/);
 }
 
-#endif // DLUP_GEOMETRY_LAZY_ARRAY_H
+#endif // DLUP_GEOMETRY_LAZY_ARRAY_H

src/geometry/polygon_collection.h

@@ -28,36 +28,26 @@ class PolygonCollection {
  return py_objects;
  }
 
- py::array_t<int> toMaskNonLazy(int default_value = 0) const {
- auto mask = generateMaskFromAnnotations(polygons_, mask_size_, default_value);
- std::cout << "Outside lambda - Number of polygons: " << polygons_.size() << std::endl;
-
- int width = std::get<0>(mask_size_);
- int height = std::get<1>(mask_size_);
-
- return py::array_t<int>({height, width}, mask->data());
-}
-
-LazyArray<int> toMask(int default_value = 0) const {
+ LazyArray<int> toMask(int default_value = 0) const {
  // Capture polygons_ and mask_size_ by value
  auto polygons_copy = polygons_;
  auto mask_size_copy = mask_size_;
 
- return LazyArray<int>([polygons_copy, mask_size_copy, default_value]() {
-  auto mask = generateMaskFromAnnotations(polygons_copy, mask_size_copy, default_value);
- int width = std::get<0>(mask_size_copy);
- int height = std::get<1>(mask_size_copy);
- return py::array_t<int>({height, width}, mask->data());
- });
-}
-
-
-
-
+ // Provide the shape as the second argument to the LazyArray constructor
+ return LazyArray<int>(
+ [polygons_copy, mask_size_copy, default_value]() {
+  auto mask = generateMaskFromAnnotations(polygons_copy, mask_size_copy, default_value);
+  int width = std::get<0>(mask_size_copy);
+  int height = std::get<1>(mask_size_copy);
+ return py::array_t<int>({height, width}, mask->data());
+ },
+ {std::get<1>(mask_size_copy), std::get<0>(mask_size_copy)} // Provide the shape explicitly
+ );
+ }
 
-private:
-std::vector<std::shared_ptr<Polygon>> polygons_;
-std::tuple<int, int> mask_size_;
+ private:
+ std::vector<std::shared_ptr<Polygon>> polygons_;
+ std::tuple<int, int> mask_size_;
 };
 
-#endif // DLUP_POLYGON_COLLECTION_H
+#endif // DLUP_POLYGON_COLLECTION_H