bindings.cpp

#include <stdexcept>

#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>

#include "filters.hpp"

namespace py = pybind11;

static inline void validate_kernel(int kernel_size) {
  if (kernel_size <= 0) {
    throw std::invalid_argument("kernel_size must be > 0");
  }
  if ((kernel_size % 2) == 0) {
    throw std::invalid_argument("kernel_size must be odd");
  }
}

static py::array gaussian_blur(py::array img, int kernel_size = 5,float sigma = 2.0f) {
  validate_kernel(kernel_size);

  py::buffer_info info = img.request();
  if (info.format != py::format_descriptor<std::uint8_t>::format()) {
    throw std::invalid_argument("gaussian_blur expects uint8 numpy array");
  }

  int height, width;
  if (info.ndim == 2) {
    height = static_cast<int>(info.shape[0]);
    width = static_cast<int>(info.shape[1]);
  } else if (info.ndim == 3) {
    if (info.shape[2] != 1) {
      throw std::invalid_argument(
          "gaussian_blur only supports grayscale images (1 channel)");
    }
    height = static_cast<int>(info.shape[0]);
    width = static_cast<int>(info.shape[1]);
  } else {
    throw std::invalid_argument(
        "gaussian_blur expects a 2D or 3D(1-channel) array");
  }

  if (height <= 0 || width <= 0) {
    throw std::invalid_argument("gaussian_blur: empty image");
  }

  // Force contiguous C-order view (copy if needed)
  py::array_t<std::uint8_t, py::array::c_style | py::array::forcecast> in(img);
  py::buffer_info in_info = in.request();

  // Allocate output with same shape as input
  py::array out;
  if (info.ndim == 2) {
    out = py::array_t<std::uint8_t>(
        std::vector<py::ssize_t>{(py::ssize_t)height, (py::ssize_t)width});
  } else {
    out = py::array_t<std::uint8_t>(
        std::vector<py::ssize_t>{(py::ssize_t)height, (py::ssize_t)width, 1});
  }
  py::buffer_info out_info = out.request();

  const auto *in_ptr = static_cast<const std::uint8_t *>(in_info.ptr);
  auto *out_ptr = static_cast<std::uint8_t *>(out_info.ptr);

  gaussian_blurs(in_ptr, out_ptr, height, width, kernel_size, sigma);
  return out;
}

static py::tuple find_contours_hierarchy(py::array img) {
  py::buffer_info info = img.request();
  if (info.format != py::format_descriptor<std::uint8_t>::format()) {
    throw std::invalid_argument(
        "find_contours_hierarchy expects a uint8 numpy array");
  }

  if (info.ndim != 2) {
    throw std::invalid_argument("find_contours_hierarchy expects a 2D array");
  }

  int height = static_cast<int>(info.shape[0]);
  int width = static_cast<int>(info.shape[1]);

  py::array_t<std::uint8_t, py::array::c_style | py::array::forcecast> in(img);
  const auto *in_ptr = static_cast<const std::uint8_t *>(in.request().ptr);

  ContourResult res = get_contours(in_ptr, height, width);

  py::list py_contours;
  for (const auto &contour : res.contours) {
    size_t n_pts = contour.size() / 2;
    auto py_contour = py::array_t<int32_t>({n_pts, (size_t)1, (size_t)2});
    auto pt_ptr = py_contour.mutable_data();
    size_t i = 0;
    while (i < contour.size()) {
      pt_ptr[i] = static_cast<int32_t>(contour[i]);
      ++i;
    }
    py_contours.append(py_contour);
  }

  size_t num_found = res.contours.size();
  auto py_hierarchy = py::array_t<int32_t>({(size_t)1, num_found, (size_t)4});
  auto h_ptr = py_hierarchy.mutable_data();

  size_t i = 0;
  while (i < num_found) {
    size_t base_idx = (i + 1) * 4;
    int p_idx = res.hierarchy[base_idx + 0];
    int c_idx = res.hierarchy[base_idx + 1];
    int next_idx = res.hierarchy[base_idx + 2];
    int prev_idx = res.hierarchy[base_idx + 3];

    int32_t final_next = (next_idx > 0) ? (next_idx - 1) : -1;
    int32_t final_prev = (prev_idx > 0) ? (prev_idx - 1) : -1;
    int32_t final_child = (c_idx > 0) ? (c_idx - 1) : -1;
    int32_t final_parent = (p_idx > 0) ? (p_idx - 1) : -1;

    h_ptr[i * 4 + 0] = final_next;
    h_ptr[i * 4 + 1] = final_prev;
    h_ptr[i * 4 + 2] = final_child;
    h_ptr[i * 4 + 3] = final_parent;
    ++i;
  }

  return py::make_tuple(py_contours, py_hierarchy);
}

PYBIND11_MODULE(filters, m) {
  m.doc() = "Vision filters implemented in C++ via pybind11";
  m.def("gaussian_blur", &gaussian_blur, py::arg("img"),
        py::arg("kernel_size") = 5, py::arg("sigma") = 2.0f,
        "Apply Gaussian blur to grayscale image. Returns image.");
  m.def("get_contours", &find_contours_hierarchy, py::arg("img"),
        "Find contours with hierarchy. Returns (contours,hierarchy).");
}