Save image to buffer

library/public/image.h

@@ -5,6 +5,7 @@
 #include "export.h"
 
 #include <string>
+#include <vector>
 
 namespace f3d
 {
@@ -153,6 +154,12 @@ class F3D_EXPORT image
    */
   void save(const std::string& path, SaveFormat format = SaveFormat::PNG) const;
 
+  /**
+   * Save an image to a memory buffer in the specified format.
+   * Default format is PNG if not specified.
+   */
+  std::vector<unsigned char> saveBuffer(SaveFormat format = SaveFormat::PNG) const;
+
   /**
    * An exception that can be thrown by the image when there.
    * is an error on write.

library/src/image.cxx

@@ -4,6 +4,7 @@
 #include "init.h"
 
 #include <vtkBMPWriter.h>
+#include <vtkDataArrayRange.h>
 #include <vtkImageData.h>
 #include <vtkImageDifference.h>
 #include <vtkImageReader2.h>
@@ -13,17 +14,38 @@
 #include <vtkPointData.h>
 #include <vtkSmartPointer.h>
 #include <vtkTIFFWriter.h>
+#include <vtkUnsignedCharArray.h>
 #include <vtksys/SystemTools.hxx>
 
 #include <cassert>
-#include <vector>
 
 namespace f3d
 {
 class image::internals
 {
 public:
   vtkSmartPointer<vtkImageData> Image;
+
+  template<typename WriterType>
+  std::vector<unsigned char> SaveBuffer()
+  {
+    vtkNew<WriterType> writer;
+    writer->WriteToMemoryOn();
+    writer->SetInputData(this->Image);
+    writer->Write();
+
+    if (writer->GetErrorCode() != 0)
+    {
+      throw write_exception("Failed to write in memory");
+    }
+
+    std::vector<unsigned char> result;
+
+    auto valRange = vtk::DataArrayValueRange(writer->GetResult());
+    std::copy(valRange.begin(), valRange.end(), std::back_inserter(result));
+
+    return result;
+  }
 };
 
 //----------------------------------------------------------------------------
@@ -290,6 +312,25 @@ void image::save(const std::string& path, SaveFormat format) const
   }
 }
 
+//----------------------------------------------------------------------------
+std::vector<unsigned char> image::saveBuffer(SaveFormat format) const
+{
+  switch (format)
+  {
+    case SaveFormat::PNG:
+      return this->Internals->SaveBuffer<vtkPNGWriter>();
+    case SaveFormat::JPG:
+      return this->Internals->SaveBuffer<vtkJPEGWriter>();
+    case SaveFormat::TIF:
+      throw write_exception("Cannot save to buffer in TIF format");
+      break;
+    case SaveFormat::BMP:
+      return this->Internals->SaveBuffer<vtkBMPWriter>();
+  }
+
+  return {};
+}
+
 //----------------------------------------------------------------------------
 image::write_exception::write_exception(const std::string& what)
   : exception(what)

library/testing/TestSDKImage.cxx

@@ -28,6 +28,38 @@ int TestSDKImage(int argc, char* argv[])
   generated.save(std::string(argv[2]) + "TestSDKImage.tif", f3d::image::SaveFormat::TIF);
   generated.save(std::string(argv[2]) + "TestSDKImage.bmp", f3d::image::SaveFormat::BMP);
 
+  // test saveBuffer in different formats
+  std::vector<unsigned char> bufferPNG = generated.saveBuffer();
+  if (bufferPNG.size() == 0)
+  {
+    std::cerr << "PNG buffer empty" << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  std::vector<unsigned char> bufferJPG = generated.saveBuffer(f3d::image::SaveFormat::JPG);
+  if (bufferJPG.size() == 0)
+  {
+    std::cerr << "JPG buffer empty" << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  try
+  {
+    generated.saveBuffer(f3d::image::SaveFormat::TIF);
+    std::cerr << "An exception has not been thrown when saving buffer to TIF format" << std::endl;
+    return EXIT_FAILURE;
+  }
+  catch (const f3d::image::write_exception&)
+  {
+  }
+
+  std::vector<unsigned char> bufferBMP = generated.saveBuffer(f3d::image::SaveFormat::BMP);
+  if (bufferBMP.size() == 0)
+  {
+    std::cerr << "BMP buffer empty" << std::endl;
+    return EXIT_FAILURE;
+  }
+
   // test constructor with different channel sizes
   f3d::image img16(width, height, channels, f3d::image::ChannelType::SHORT);
   f3d::image img32(width, height, channels, f3d::image::ChannelType::FLOAT);

python/F3DPythonBindings.cxx

@@ -102,6 +102,12 @@ PYBIND11_MODULE(pyf3d, module)
     return py::bytes(static_cast<char*>(img.getContent()), expectedSize);
   };
 
+  auto getFileBytes = [](const f3d::image& img, f3d::image::SaveFormat format)
+  {
+    std::vector<unsigned char> result = img.saveBuffer(format);
+    return py::bytes(reinterpret_cast<char*>(result.data()), result.size());
+  };
+
   image //
     .def(py::init<>())
     .def(py::init<const std::string&>())
@@ -116,7 +122,8 @@ PYBIND11_MODULE(pyf3d, module)
     .def_property("content", getImageBytes, setImageBytes)
     .def("compare", &f3d::image::compare)
     .def(
-      "save", &f3d::image::save, py::arg("path"), py::arg("format") = f3d::image::SaveFormat::PNG);
+      "save", &f3d::image::save, py::arg("path"), py::arg("format") = f3d::image::SaveFormat::PNG)
+    .def("save_buffer", getFileBytes, py::arg("format") = f3d::image::SaveFormat::PNG);
 
   // f3d::options
   py::class_<f3d::options> options(module, "Options");

python/__init__.py.in

@@ -95,3 +95,20 @@ def add_deprecation_warnings():
 
 
 add_deprecation_warnings()
+
+
+################################################################################
+# add support for IPython
+
+
+try:
+    from IPython import get_ipython
+
+    def f3d_image_repr_png(f3d_img: Image):
+        return f3d_img.save_buffer()
+
+    if ipython := get_ipython():
+        png_formatter = ipython.display_formatter.formatters["image/png"]
+        png_formatter.for_type(Image, f3d_image_repr_png)
+except ImportError:
+    pass

python/testing/test_image.py

@@ -56,3 +56,9 @@ def test_save(f3d_engine):
 
     img.save(fn, f3d.Image.SaveFormat.BMP)
     assert os.path.isfile(fn)
+
+
+def test_save_buffer(f3d_engine):
+    img = f3d_engine.window.render_to_image(True)
+    buffer = img.save_buffer(f3d.Image.SaveFormat.PNG)
+    assert buffer.startswith(b"\x89PNG")