google-deepmind / spiral

We provide a pre-trained model for unconditional 19-step generation of CelebA-HQ images
Apache License 2.0
326 stars 45 forks source link

how to edit pylibmypaint.cc to add pickle support for libmypaint? #4

Closed urw7rs closed 4 years ago

urw7rs commented 4 years ago

I want to add pickle support to libmypaint so I can run it in parallel using multiprocessing. I tried to edit pylibmypaint.cc following pybind11's documentations.

This is how I edited pylibmypaint.cc:

PYBIND11_MODULE(pylibmypaint, m) {
  py::class_<BrushWrapper> brush_wrapper(m, "BrushWrapper");
  brush_wrapper
      .def(py::init())
      .def("LoadFromFile", &BrushWrapper::LoadFromFile)
      .def("SetBaseValue", &BrushWrapper::SetBaseValue)
      .def("SetSurface", &BrushWrapper::SetSurface)
      .def("Reset", &BrushWrapper::Reset)
      .def("NewStroke", &BrushWrapper::NewStroke)
      .def("StrokeTo", &BrushWrapper::StrokeTo)
      .def("SetSurface", &BrushWrapper::SetSurface)
      .def("GetSurface", &BrushWrapper::GetSurface).        // added for pickling
      .def("GetState", &BrushWrapper::GetState)
      .def("SetState", &BrushWrapper::SetState)
      .def(py::pickle(
          [](BrushWrapper &b) {
              /* Return a tuple that fully encodes the state of the object */
          std::vector<float> state = b.GetState();
              auto surface = b.GetSurface();
              return py::make_tuple(state, surface);
          },
          [](py::tuple t) { 
              if (t.size() != 2)
                   throw std::runtime_error("Invalid state!");
              /* Invoke the in-place constructor. Note that this is needed even
                 when the object just has a trivial default constructor */
              auto b = BrushWrapper();

              /* Assign any additional state */
          b.SetState(t[0].cast<std::vector<float>>());
          return b;
      }
      ));

  py::class_<SurfaceWrapper> surface_wrapper(m, "SurfaceWrapper");
  surface_wrapper
      .def(py::init<int, int, SurfaceWrapper::Background>())
      .def("BeginAtomic", &SurfaceWrapper::BeginAtomic)
      .def("EndAtomic", &SurfaceWrapper::EndAtomic)
      .def("Clear", &SurfaceWrapper::Clear)
      .def("BufferAsNumpy", [](SurfaceWrapper &sw) {
           uint16_t* buffer = sw.GetBuffer();
           auto buffer_dims = sw.GetBufferDims();
           return py::array_t<uint16_t>(buffer_dims, buffer);
      })
      .def(py::pickle(
          [](SurfaceWrapper sw) {
              /* Return a tuple that fully encodes the state of the object */
              uint16_t* buffer_ = sw.GetBuffer();
              int n = sizeof (buffer_) / sizeof (buffer_[0]);
              std::vector<uint16_t> buffer(buffer_, buffer_ + n);
              return py::make_tuple(sw.w, sw.h, sw.c, buffer);
          },
          [](py::tuple t) { 
              if (t.size() != 4)
                  throw std::runtime_error("Invalid state!");
              /* Invoke the in-place constructor. Note that this is needed even
                 when the object just has a trivial default constructor */
              auto sw = SurfaceWrapper(t[0].cast<int>(), 
                           t[1].cast<int>(), 
                           t[2].cast<SurfaceWrapper::Background>());
          uint16_t* buffer_ = sw.GetBuffer();
              std::vector<uint16_t> buffer = t[3].cast<std::vector<uint16_t>>();
          std::copy(buffer.begin(), buffer.end(), buffer_);
          return sw;
      }
      ));

I did get it to build without errors on ubuntu 18.04 using __getstate__ and __setstate__ but whenever I try to pickle SurfaceWrapper, I got a segmentation fault. And using the recommended py::pickle method, I get errors like this: error: use of deleted function ‘spiral::libmypaint::SurfaceWrapper::SurfaceWrapper(const spiral::libmypaint::SurfaceWrapper&)’. I think it's because of

  SurfaceWrapper(const SurfaceWrapper&) = delete;
  void operator=(const SurfaceWrapper&) = delete;

from libmypaint_wrapper/surface.h at line 58, 59.

I'm not familiar with c++ and pybind. I would really appreciate your help.