-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathpython_bindings.cpp
More file actions
84 lines (75 loc) · 3.07 KB
/
python_bindings.cpp
File metadata and controls
84 lines (75 loc) · 3.07 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
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
#include "buffer.hpp"
#include "video.hpp"
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
namespace py = pybind11;
namespace {
std::vector<int64_t> slice_to_indexes(const py::slice& s, int64_t max_size)
{
ssize_t start = 0, stop = 0, step = 0, slicelength = 0;
if (!s.compute(max_size, &start, &stop, &step, &slicelength)) {
throw py::error_already_set();
}
std::vector<int64_t> indices;
indices.reserve(slicelength);
for (ssize_t i = start; i < stop; i += step) {
indices.push_back(i);
}
return indices;
}
std::vector<int64_t>
compute_indexes(const std::variant<int64_t, pybind11::slice, pybind11::list, pybind11::tuple>& index, int64_t max_size)
{
return std::visit(
[max_size]<typename T>(const T& val) -> std::vector<int64_t> {
if constexpr (std::is_same_v<T, int64_t>) {
return {val};
} else if constexpr (std::is_same_v<T, pybind11::slice>) {
return slice_to_indexes(val, max_size);
} else {
std::vector<int64_t> indices;
indices.reserve(val.size());
for (auto& i : val) {
indices.push_back(pybind11::cast<int64_t>(i));
}
return indices;
}
},
index);
}
} // namespace
PYBIND11_MODULE(deepframe, m)
{
py::class_<core::buffer, std::shared_ptr<core::buffer>>(m, "Buffer", py::buffer_protocol())
.def_buffer([](core::buffer& buf) -> py::buffer_info {
auto dims = buf.dims();
auto stride = buf.stride();
return py::buffer_info(buf.data(), /* Pointer to buffer */
sizeof(uint8_t), /* Size of one scalar */
py::format_descriptor<uint8_t>::format(), /* Python struct-style
format descriptor */
dims.size(), /* Number of dimensions */
dims, /* Buffer dimensions */
stride, /* Strides (in bytes) for each index */
true /* Buffer is read-only */
);
});
m.def(
"extract_frames",
[](const std::string& url,
std::variant<int64_t, py::slice, py::list, py::tuple>& index,
int64_t max_size = std::numeric_limits<int64_t>::max()) {
auto indices = compute_indexes(index, max_size);
return codecs::extract_video_frames_from_video_at_url(url, indices);
},
py::arg("url"),
py::arg("index"),
py::arg("max_size") = std::numeric_limits<int64_t>::max());
m.def(
"get_info",
[](const std::string& url) {
return codecs::get_video_info(url);
},
py::arg("url"));
}