Skip to content

Commit 2d9605c

Browse files
sfe-SparkFrosfe-SparkFro
committed
Add imwrite()
1 parent da558f4 commit 2d9605c

File tree

3 files changed

+122
-2
lines changed

3 files changed

+122
-2
lines changed

‎src/imgcodecs.cpp

Lines changed: 79 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,84 @@ mp_obj_t cv2_imgcodecs_imread(size_t n_args, const mp_obj_t *pos_args, mp_map_t
7171
mp_raise_msg(&mp_type_Exception, MP_ERROR_TEXT(e.what()));
7272
}
7373

74-
// Return the image
74+
// Return the result
7575
return mat_to_mp_obj(img);
7676
}
77+
78+
mp_obj_t cv2_imgcodecs_imwrite(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
79+
// Define the arguments
80+
enum { ARG_filename, ARG_img, ARG_params };
81+
static const mp_arg_t allowed_args[] = {
82+
{ MP_QSTR_filename, MP_ARG_REQUIRED | MP_ARG_OBJ, { .u_obj = MP_OBJ_NULL } },
83+
{ MP_QSTR_img, MP_ARG_REQUIRED | MP_ARG_OBJ, { .u_obj = MP_OBJ_NULL } },
84+
{ MP_QSTR_params, MP_ARG_OBJ, { .u_obj = mp_const_none } },
85+
};
86+
87+
// Parse the arguments
88+
mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
89+
mp_arg_parse_all(n_args, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
90+
91+
// Convert arguments to required types
92+
mp_obj_t filename = args[ARG_filename].u_obj;
93+
Mat img = mp_obj_to_mat(args[ARG_img].u_obj);
94+
ndarray_obj_t *params;
95+
if (args[ARG_params].u_obj == mp_const_none) {
96+
// If no parameters are provided, use an empty ndarray
97+
params = ndarray_new_linear_array(0, NDARRAY_INT16);
98+
} else {
99+
params = ndarray_from_mp_obj(args[ARG_params].u_obj, 0);
100+
}
101+
102+
// Convert the filename to a std::string
103+
size_t filename_len;
104+
const char *filename_chr = mp_obj_str_get_data(filename, &filename_len);
105+
std::string filename_str(filename_chr, filename_len);
106+
107+
// Create vector of uint8_t for encoding
108+
std::vector<uint8_t> buf;
109+
110+
// Convert the parameters to a vector of int
111+
std::vector<int> params_vec;
112+
if (params->len > 0) {
113+
params_vec.reserve(params->len);
114+
for (size_t i = 0; i < params->len; ++i) {
115+
mp_obj_t val = (mp_obj_t*) mp_binary_get_val_array(params->dtype, params->array, i);
116+
// ndarrays default to float, and mp_obj_get_int() does not support
117+
// float values, so we need to do the type conversion ourselves
118+
mp_float_t val_float = mp_obj_get_float(val);
119+
params_vec.push_back((int) val_float);
120+
}
121+
}
122+
123+
// Encode the image from the buffer
124+
bool retval;
125+
try {
126+
retval = imencode(filename_str, img, buf, params_vec);
127+
} catch(Exception& e) {
128+
mp_raise_msg(&mp_type_Exception, MP_ERROR_TEXT(e.what()));
129+
}
130+
131+
// Convert the vector of uint8_t to a bytes object
132+
mp_obj_t buf_obj = mp_obj_new_bytes((const byte *)buf.data(), buf.size());
133+
134+
// Call MicroPython's `open()` function to write the image file
135+
mp_obj_t open_args[2];
136+
open_args[0] = filename;
137+
open_args[1] = mp_obj_new_str("wb", 2); // Open in binary write mode
138+
mp_map_t open_kw_args = MP_EMPTY_MAP(); // No keyword arguments
139+
mp_obj_t file_obj = mp_builtin_open(2, open_args, &open_kw_args);
140+
141+
// Call the `write()` method on the file object to write the image data
142+
mp_obj_t write_method[3];
143+
mp_load_method(file_obj, MP_QSTR_write, write_method);
144+
write_method[2] = buf_obj; // Set the data to write
145+
mp_call_method_n_kw(1, 0, write_method);
146+
147+
// Close the file object
148+
mp_obj_t close_method[2];
149+
mp_load_method(file_obj, MP_QSTR_close, close_method);
150+
mp_call_method_n_kw(0, 0, close_method);
151+
152+
// Return the result
153+
return mp_obj_new_bool(retval);
154+
}

‎src/imgcodecs.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,3 +2,4 @@
22
#include "py/runtime.h"
33

44
extern mp_obj_t cv2_imgcodecs_imread(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args);
5+
extern mp_obj_t cv2_imgcodecs_imwrite(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args);

‎src/opencv_upy.c

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ static MP_DEFINE_CONST_FUN_OBJ_KW(cv2_highgui_waitKey_obj, 0, cv2_highgui_waitKe
1717

1818
// OpenCV imgcodecs module
1919
static MP_DEFINE_CONST_FUN_OBJ_KW(cv2_imgcodecs_imread_obj, 1, cv2_imgcodecs_imread);
20+
static MP_DEFINE_CONST_FUN_OBJ_KW(cv2_imgcodecs_imwrite_obj, 2, cv2_imgcodecs_imwrite);
2021

2122
// OpenCV imgproc module
2223
static MP_DEFINE_CONST_FUN_OBJ_KW(cv2_imgproc_adaptiveThreshold_obj, 6, cv2_imgproc_adaptiveThreshold);
@@ -92,7 +93,7 @@ static const mp_rom_map_elem_t cv2_module_globals_table[] = {
9293
{ MP_ROM_QSTR(MP_QSTR_BORDER_DEFAULT), MP_ROM_INT(4) },
9394
{ MP_ROM_QSTR(MP_QSTR_BORDER_ISOLATED), MP_ROM_INT(16) },
9495

95-
// Image read mode flags, from opencv2/imgcodecs.hpp
96+
// Image read flags, from opencv2/imgcodecs.hpp
9697
{ MP_ROM_QSTR(MP_QSTR_IMREAD_UNCHANGED), MP_ROM_INT(-1) },
9798
{ MP_ROM_QSTR(MP_QSTR_IMREAD_GRAYSCALE), MP_ROM_INT(0) },
9899
{ MP_ROM_QSTR(MP_QSTR_IMREAD_COLOR_BGR), MP_ROM_INT(1) },
@@ -109,6 +110,45 @@ static const mp_rom_map_elem_t cv2_module_globals_table[] = {
109110
{ MP_ROM_QSTR(MP_QSTR_IMREAD_IGNORE_ORIENTATION), MP_ROM_INT(128) },
110111
{ MP_ROM_QSTR(MP_QSTR_IMREAD_COLOR_RGB), MP_ROM_INT(256) },
111112

113+
// Image write flags, from opencv2/imgcodecs.hpp
114+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEG_QUALITY), MP_ROM_INT(1) },
115+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEG_PROGRESSIVE), MP_ROM_INT(2) },
116+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEG_OPTIMIZE), MP_ROM_INT(3) },
117+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEG_RST_INTERVAL), MP_ROM_INT(4) },
118+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEG_LUMA_QUALITY), MP_ROM_INT(5) },
119+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEG_CHROMA_QUALITY), MP_ROM_INT(6) },
120+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEG_SAMPLING_FACTOR), MP_ROM_INT(7) },
121+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_PNG_COMPRESSION), MP_ROM_INT(16) },
122+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_PNG_STRATEGY), MP_ROM_INT(17) },
123+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_PNG_BILEVEL), MP_ROM_INT(18) },
124+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_PXM_BINARY), MP_ROM_INT(32) },
125+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_EXR_TYPE), MP_ROM_INT((3 << 4) + 0) },
126+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_EXR_COMPRESSION), MP_ROM_INT((3 << 4) + 1) },
127+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_EXR_DWA_COMPRESSION_LEVEL), MP_ROM_INT((3 << 4) + 2) },
128+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_WEBP_QUALITY), MP_ROM_INT(64) },
129+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_HDR_COMPRESSION), MP_ROM_INT((5 << 4) + 0) },
130+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_PAM_TUPLETYPE), MP_ROM_INT(128) },
131+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_TIFF_RESUNIT), MP_ROM_INT(256) },
132+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_TIFF_XDPI), MP_ROM_INT(257) },
133+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_TIFF_YDPI), MP_ROM_INT(258) },
134+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_TIFF_COMPRESSION), MP_ROM_INT(259) },
135+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_TIFF_ROWSPERSTRIP), MP_ROM_INT(278) },
136+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_TIFF_PREDICTOR), MP_ROM_INT(317) },
137+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEG2000_COMPRESSION_X1000), MP_ROM_INT(272) },
138+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_AVIF_QUALITY), MP_ROM_INT(512) },
139+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_AVIF_DEPTH), MP_ROM_INT(513) },
140+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_AVIF_SPEED), MP_ROM_INT(514) },
141+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEGXL_QUALITY), MP_ROM_INT(640) },
142+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEGXL_EFFORT), MP_ROM_INT(641) },
143+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEGXL_DISTANCE), MP_ROM_INT(642) },
144+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_JPEGXL_DECODING_SPEED), MP_ROM_INT(643) },
145+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_GIF_LOOP), MP_ROM_INT(1024) },
146+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_GIF_SPEED), MP_ROM_INT(1025) },
147+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_GIF_QUALITY), MP_ROM_INT(1026) },
148+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_GIF_DITHER), MP_ROM_INT(1027) },
149+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_GIF_TRANSPARENCY), MP_ROM_INT(1028) },
150+
{ MP_ROM_QSTR(MP_QSTR_IMWRITE_GIF_COLORTABLE), MP_ROM_INT(1029) },
151+
112152
// Morphology operation types, from opencv2/imgproc.hpp
113153
{ MP_ROM_QSTR(MP_QSTR_MORPH_ERODE), MP_ROM_INT(0) },
114154
{ MP_ROM_QSTR(MP_QSTR_MORPH_DILATE), MP_ROM_INT(1) },
@@ -250,6 +290,7 @@ static const mp_rom_map_elem_t cv2_module_globals_table[] = {
250290
////////////////////////////////////////////////////////////////////////////
251291

252292
{ MP_ROM_QSTR(MP_QSTR_imread), MP_ROM_PTR(&cv2_imgcodecs_imread_obj) },
293+
{ MP_ROM_QSTR(MP_QSTR_imwrite), MP_ROM_PTR(&cv2_imgcodecs_imwrite_obj) },
253294

254295
////////////////////////////////////////////////////////////////////////////
255296
// OpenCV imgproc functions

0 commit comments

Comments
 (0)