pybind.cc 3.7 KB
Newer Older
S
superjom 已提交
1 2
#include <ctype.h>
#include <pybind11/pybind11.h>
S
superjom 已提交
3
#include <pybind11/stl.h>
S
superjom 已提交
4

S
superjom 已提交
5
#include "visualdl/logic/sdk.h"
S
superjom 已提交
6 7 8

namespace py = pybind11;
namespace vs = visualdl;
S
superjom 已提交
9
namespace cp = visualdl::components;
S
superjom 已提交
10

Q
Qiao Longfei 已提交
11 12
PYBIND11_PLUGIN(core) {
  py::module m("core", "C++ core of VisualDL");
S
superjom 已提交
13

S
superjom 已提交
14
#define READER_ADD_SCALAR(T)                                            \
S
superjom 已提交
15 16 17 18 19
  .def("get_scalar_" #T, [](vs::Reader& self, const std::string& tag) { \
    auto tablet = self.tablet(tag);                                     \
    return vs::components::ScalarReader<T>(std::move(tablet));          \
  })
  py::class_<vs::Reader>(m, "Reader")
S
superjom 已提交
20 21 22 23 24 25 26
      .def("__init__",
           [](vs::Reader& instance, const std::string& dir) {
             new (&instance) vs::Reader(dir);
           })
      .def("as_mode", &vs::Reader::AsMode)
      .def("modes", [](vs::Reader& self) { return self.storage().modes(); })
      .def("tags", &vs::Reader::tags)
S
superjom 已提交
27
      // clang-format off
S
superjom 已提交
28 29 30 31 32 33 34 35 36
      READER_ADD_SCALAR(float)
      READER_ADD_SCALAR(double)
      READER_ADD_SCALAR(int)
      // clang-format on
      .def("get_image", [](vs::Reader& self, const std::string& tag) {
        auto tablet = self.tablet(tag);
        return vs::components::ImageReader(self.mode(), tablet);
      });
#undef READER_ADD_SCALAR
S
superjom 已提交
37

S
superjom 已提交
38
#define WRITER_ADD_SCALAR(T)                                            \
S
superjom 已提交
39 40 41 42
  .def("new_scalar_" #T, [](vs::Writer& self, const std::string& tag) { \
    auto tablet = self.AddTablet(tag);                                  \
    return cp::Scalar<T>(tablet);                                       \
  })
S
superjom 已提交
43

S
superjom 已提交
44
  py::class_<vs::Writer>(m, "Writer")
S
superjom 已提交
45
      .def("__init__",
S
superjom 已提交
46
           [](vs::Writer& instance, const std::string& dir, int sync_cycle) {
S
superjom 已提交
47
             new (&instance) vs::Writer(dir, sync_cycle);
S
superjom 已提交
48 49
           })
      .def("as_mode", &vs::Writer::AsMode)
S
superjom 已提交
50
      // clang-format off
S
superjom 已提交
51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93
      WRITER_ADD_SCALAR(float)
      WRITER_ADD_SCALAR(double)
      WRITER_ADD_SCALAR(int)
      // clang-format on
      .def("new_image",
           [](vs::Writer& self, const std::string& tag, int num_samples) {
             auto tablet = self.AddTablet(tag);
             return vs::components::Image(tablet, num_samples);
           });

//------------------- components --------------------
#define ADD_SCALAR_READER(T)                               \
  py::class_<cp::ScalarReader<T>>(m, "ScalarReader__" #T)  \
      .def("records", &cp::ScalarReader<T>::records)       \
      .def("timestamps", &cp::ScalarReader<T>::timestamps) \
      .def("ids", &cp::ScalarReader<T>::ids)               \
      .def("caption", &cp::ScalarReader<T>::caption);
  ADD_SCALAR_READER(int);
  ADD_SCALAR_READER(float);
  ADD_SCALAR_READER(double);
  ADD_SCALAR_READER(int64_t);
#undef ADD_SCALAR_READER

#define ADD_SCALAR_WRITER(T)                          \
  py::class_<cp::Scalar<T>>(m, "ScalarWriter__" #T)   \
      .def("set_caption", &cp::Scalar<T>::SetCaption) \
      .def("add_record", &cp::Scalar<T>::AddRecord);
  ADD_SCALAR_WRITER(int);
  ADD_SCALAR_WRITER(float);
  ADD_SCALAR_WRITER(double);
#undef ADD_SCALAR_WRITER

  // clang-format on
  py::class_<cp::Image>(m, "ImageWriter")
      .def("set_caption", &cp::Image::SetCaption)
      .def("start_sampling", &cp::Image::StartSampling)
      .def("is_sample_taken", &cp::Image::IsSampleTaken)
      .def("finish_sampling", &cp::Image::FinishSampling)
      .def("set_sample", &cp::Image::SetSample);

  py::class_<cp::ImageReader>(m, "ImageReader")
      .def("caption", &cp::ImageReader::caption)
      .def("num_records", &cp::ImageReader::num_records)
S
superjom 已提交
94 95
      .def("num_samples", &cp::ImageReader::num_samples)
      .def("timestamp", &cp::ImageReader::timestamp)
S
superjom 已提交
96 97
      .def("data", &cp::ImageReader::data)
      .def("shape", &cp::ImageReader::shape);
S
superjom 已提交
98 99

}  // end pybind