How to pass C++ itk::Image to Python

.def(
        "to_pyarray",
        [](const TImagePointer &img,
           const std::string &contiguous) {
            const auto size =
                    img->GetLargestPossibleRegion().GetSize();
            const auto shape =
                    (contiguous == "F")
                            ? std::vector<size_t>{size[2], size[1],
                                                  size[0]}
                            : std::vector<size_t>{size[0], size[1],
                                                  size[2]};
            return py::array(
                    py::dtype::of<typename TImagePointer::
                                          ObjectType::PixelType>(),
                    shape, img->GetBufferPointer());
        },
        py::arg("contig") = "F")
// TODO: Create a view (non-copy) of the data
// Problems will arise with the contig differences between
// numpy(fortran) and c.
.def(
        "as_pyarray",
        [](const TImagePointer & /* img */,
           const std::string & /* contiguous */) {
            throw std::runtime_error(
                    "not implemented, use to_pyarray");
        },
        py::arg("contig") = "F")

.def(
        "from_pyarray",
        [](TImagePointer &img,
           py::array_t<
                   typename TImagePointer::ObjectType::PixelType>
                   np_array,
           const std::string &contiguous) {
            using PixelType =
                    typename TImagePointer::ObjectType::PixelType;
            using Image = typename TImagePointer::ObjectType;
            using ImporterType =
                    itk::ImportImageFilter<PixelType,
                                           Image::ImageDimension>;
            auto info = np_array.request();

            auto importer = ImporterType::New();
            auto region = img->GetLargestPossibleRegion();
            auto size = region.GetSize();

            if (contiguous == "F") {
                std::copy(info.shape.rbegin(), info.shape.rend(),
                          size.begin());
            } else if (contiguous == "C") {
                std::copy(info.shape.begin(), info.shape.end(),
                          size.begin());
            } else {
                throw std::runtime_error(
                        "Unknown parameter contig: " + contiguous +
                        ". Valid: F or C.");
            }
            region.SetSize(size);
            // Note that region index is kept from the staring img.
            importer->SetRegion(region);
            // Metadata is ignored (defaulted)
            // importer->SetOrigin(img->GetOrigin());
            // importer->SetSpacing(img->GetSpacing());
            // importer->SetDirection(img->GetDirection());
            // img owns the buffer, not the import filter
            const bool importImageFilterWillOwnTheBuffer = false;
            const auto data =
                    static_cast<typename TImagePointer::ObjectType::
                                        PixelType *>(info.ptr);
            const auto numberOfPixels = np_array.size();
            importer->SetImportPointer(
                    data, numberOfPixels,
                    importImageFilterWillOwnTheBuffer);
            importer->Update();
            img = importer->GetOutput();
        },
        py::arg("input"), py::arg("contig") = "F")

1 Like