diff --git a/ojph/_imread.py b/ojph/_imread.py index 927a90b..e40ce9a 100644 --- a/ojph/_imread.py +++ b/ojph/_imread.py @@ -303,19 +303,13 @@ def read_image( max_val = iinfo.max if self._num_components == 1: - # Single component - always HW format - for h in range(height): - self._codestream_pull(0, out=image[h], min_val=min_val, max_val=max_val) + self._codestream_pull_lines_optimized(0, height, image, min_val, max_val) elif self._channel_order == 'CHW': - # Non-RGB multi-component - use planar flag for format detection for c in range(self._num_components): - for h in range(height): - self._codestream_pull(c, out=image[c, h, :], min_val=min_val, max_val=max_val) + self._codestream_pull_lines_optimized(c, height, image[c, :, :], min_val, max_val) else: - # Non-planar mode was used for writing - return HWC format for c in range(self._num_components): - for h in range(height): - self._codestream_pull(c, out=image[h, :, c], min_val=min_val, max_val=max_val) + self._codestream_pull_lines_optimized(c, height, image[:, :, c], min_val, max_val) self._close_codestream_and_file() return image @@ -340,6 +334,25 @@ def _codestream_pull(self, component, out, min_val=None, max_val=None): else: out[:] = line_array + def _codestream_pull_lines_optimized(self, component, num_lines, out_array, min_val=None, max_val=None): + if out_array.ndim != 2: + raise ValueError("Output array must be 2D for optimized pull") + + height, width = out_array.shape + if num_lines != height: + raise ValueError(f"num_lines ({num_lines}) must match array height ({height})") + + lines_pulled = self._codestream.pull_lines_to_array( + component, + num_lines, + out_array, + min_val if min_val is not None else None, + max_val if max_val is not None else None + ) + + if lines_pulled != num_lines: + raise RuntimeError(f"Expected {num_lines} lines, but only pulled {lines_pulled}") + def _close_codestream_and_file(self): if self._codestream is not None: self._codestream.close() diff --git a/ojph/ojph_bindings.cpp b/ojph/ojph_bindings.cpp index 30a2351..76abcee 100644 --- a/ojph/ojph_bindings.cpp +++ b/ojph/ojph_bindings.cpp @@ -1,7 +1,9 @@ #include #include #include - +#include +#include +#include #include #include @@ -13,7 +15,7 @@ using namespace ojph; PYBIND11_MODULE(ojph_bindings, m) { py::class_(m, "InfileBase") - .def("read", &infile_base::read) + .def("read", &infile_base::read, py::call_guard()) .def("seek", &infile_base::seek) .def("tell", &infile_base::tell) .def("eof", &infile_base::eof) @@ -22,7 +24,7 @@ PYBIND11_MODULE(ojph_bindings, m) { py::class_(m, "J2CInfile") .def(py::init<>()) .def("open", &j2c_infile::open) - .def("read", &j2c_infile::read) + .def("read", &j2c_infile::read, py::call_guard()) .def("seek", [](infile_base& self, si64 offset, int origin) { return self.seek(offset, static_cast(origin)); }) @@ -50,6 +52,7 @@ PYBIND11_MODULE(ojph_bindings, m) { if (buf.size < size) { throw py::value_error("Buffer size is smaller than requested read size"); } + py::gil_scoped_release release; return self.read(static_cast(buf.ptr), size); }, py::arg("buffer"), py::arg("size")) .def("seek", [](mem_infile& self, si64 offset, int origin) { @@ -61,7 +64,7 @@ PYBIND11_MODULE(ojph_bindings, m) { py::class_(m, "outfileBase") - .def("write", &outfile_base::write) + .def("write", &outfile_base::write, py::call_guard()) .def("seek", &outfile_base::seek) .def("tell", &outfile_base::tell) .def("close", &outfile_base::close); @@ -69,14 +72,14 @@ PYBIND11_MODULE(ojph_bindings, m) { py::class_(m, "J2COutfile") .def(py::init<>()) .def("open", &j2c_outfile::open) - .def("write", &j2c_outfile::write) + .def("write", &j2c_outfile::write, py::call_guard()) .def("tell", &j2c_outfile::tell) .def("close", &j2c_outfile::close); py::class_(m, "MemOutfile") .def(py::init<>()) .def("open", &mem_outfile::open, py::arg("initial_size") = 65536, py::arg("clear_mem") = false) - .def("write", &mem_outfile::write) + .def("write", &mem_outfile::write, py::call_guard()) .def("tell", &mem_outfile::tell) .def("get_used_size", &mem_outfile::get_used_size) .def("get_buf_size", &mem_outfile::get_buf_size) @@ -84,7 +87,7 @@ PYBIND11_MODULE(ojph_bindings, m) { return self.seek(offset, static_cast(origin)); }) .def("close", &mem_outfile::close) - .def("write_to_file", &mem_outfile::write_to_file) + .def("write_to_file", &mem_outfile::write_to_file, py::call_guard()) .def("get_data", [](mem_outfile& self) { const ui8* data = self.get_data(); si64 size = self.tell(); @@ -105,6 +108,7 @@ PYBIND11_MODULE(ojph_bindings, m) { [](codestream &self, outfile_base *file, py::object comments, ui32 num_comments) { // Check if the comments argument is None and convert it to nullptr if so const comment_exchange* comments_ptr = comments.is_none() ? nullptr : comments.cast(); + py::gil_scoped_release release; self.write_headers(file, comments_ptr, num_comments); }, py::arg("file"), py::arg("comments") = py::none(), py::arg("num_comments") = 0) @@ -114,15 +118,76 @@ PYBIND11_MODULE(ojph_bindings, m) { if (!line_buf_obj.is_none()) { buf = line_buf_obj.cast(); } + py::gil_scoped_release release; return self.exchange(buf, next_component); }, py::arg("line_buf_obj") = py::none(), py::arg("next_component") = 0) - .def("flush", &codestream::flush) + .def("flush", &codestream::flush, py::call_guard()) .def("enable_resilience", &codestream::enable_resilience) - .def("read_headers", &codestream::read_headers) + .def("read_headers", &codestream::read_headers, py::call_guard()) .def("restrict_input_resolution", &codestream::restrict_input_resolution) - .def("create", &codestream::create) - .def("pull", &codestream::pull) + .def("create", &codestream::create, py::call_guard()) + .def("pull", &codestream::pull, py::call_guard()) + .def("pull_lines_to_array", + [](codestream &self, ui32 component, ui32 num_lines, + py::array_t out_array, py::object min_val_obj, py::object max_val_obj) -> ui32 { + py::buffer_info out_buf = out_array.request(); + if (out_buf.ndim != 2) { + throw py::value_error("Output array must be 2-dimensional"); + } + + size_t line_size = out_buf.shape[1]; + si32* out_ptr = static_cast(out_buf.ptr); + + bool has_clipping = !min_val_obj.is_none() || !max_val_obj.is_none(); + si32 min_val = min_val_obj.is_none() ? std::numeric_limits::min() + : min_val_obj.cast(); + si32 max_val = max_val_obj.is_none() ? std::numeric_limits::max() + : max_val_obj.cast(); + + ui32 lines_pulled = 0; + ui32 current_comp = component; + + { + py::gil_scoped_release release; + + for (ui32 h = 0; h < num_lines; ++h) { + line_buf* line = self.pull(current_comp); + if (line == nullptr || line->i32 == nullptr) { + break; + } + + if (current_comp != component) { + break; + } + + if (line->size != line_size) { + throw py::value_error("Line size mismatch"); + } + + const si32* src_ptr = line->i32; + si32* dst_ptr = out_ptr + h * line_size; + + if (has_clipping) { + for (size_t i = 0; i < line_size; ++i) { + si32 val = src_ptr[i]; + dst_ptr[i] = std::max(min_val, std::min(max_val, val)); + } + } else { + std::memcpy(dst_ptr, src_ptr, line_size * sizeof(si32)); + } + + lines_pulled++; + } + } + + return lines_pulled; + }, + py::arg("component"), + py::arg("num_lines"), + py::arg("out_array"), + py::arg("min_val") = py::none(), + py::arg("max_val") = py::none()) .def("close", &codestream::close) .def("access_siz", &codestream::access_siz) .def("access_cod", &codestream::access_cod)