jpegio: add bitmaptools.blit like clipping to decode

Author: jepler

shared-bindings/jpegio/JpegDecoder.c

@@ -28,6 +28,7 @@
 #include "py/builtin.h"
 #include "py/runtime.h"
 
+#include "shared-bindings/bitmaptools/__init__.h"
 #include "shared-bindings/displayio/Bitmap.h"
 #include "shared-bindings/jpegio/JpegDecoder.h"
 #include "shared-module/jpegio/JpegDecoder.h"
@@ -106,6 +107,15 @@ MP_DEFINE_CONST_FUN_OBJ_2(jpegio_jpegdecoder_open_obj, jpegio_jpegdecoder_open);
 //|         self,
 //|         bitmap: displayio.Bitmap,
 //|         scale: int = 0,
+//|         x: int = 0,
+//|         y: int = 0,
+//|         *,
+//|         x1: int,
+//|         y1: int,
+//|         x2: int,
+//|         y2: int,
+//|         skip_source_index: int,
+//|         skip_dest_index: int,
 //|     ) -> None:
 //|         """Decode JPEG data
 //|
@@ -115,21 +125,45 @@ MP_DEFINE_CONST_FUN_OBJ_2(jpegio_jpegdecoder_open_obj, jpegio_jpegdecoder_open);
 //|         The image is optionally downscaled by a factor of ``2**scale``.
 //|         Scaling by a factor of 8 (scale=3) is particularly efficient in terms of decoding time.
 //|
+//|         The remaining parameters are as for `bitmaptools.blit`.
+//|         Because JPEG is a lossy data format, chroma keying based on the "source
+//|         index" is not reliable, because the same original RGB value might end
+//|         up being decompressed as a similar but not equal color value. Using a
+//|         higher JPEG encoding quality can help, but ultimately it will not be
+//|         perfect.
+//|
 //|         After a call to ``decode``, you must ``open`` a new JPEG. It is not
 //|         possible to repeatedly ``decode`` the same jpeg data, even if it is to
 //|         select different scales or crop regions from it.
 //|
 //|         :param Bitmap bitmap: Optional output buffer
 //|         :param int scale: Scale factor from 0 to 3, inclusive.
+//|         :param int x: Horizontal pixel location in bitmap where source_bitmap upper-left
+//|                       corner will be placed
+//|         :param int y: Vertical pixel location in bitmap where source_bitmap upper-left
+//|                       corner will be placed
+//|         :param int x1: Minimum x-value for rectangular bounding box to be copied from the source bitmap
+//|         :param int y1: Minimum y-value for rectangular bounding box to be copied from the source bitmap
+//|         :param int x2: Maximum x-value (exclusive) for rectangular bounding box to be copied from the source bitmap
+//|         :param int y2: Maximum y-value (exclusive) for rectangular bounding box to be copied from the source bitmap
+//|         :param int skip_source_index: bitmap palette index in the source that will not be copied,
+//|                                set to None to copy all pixels
+//|         :param int skip_dest_index: bitmap palette index in the destination bitmap that will not get overwritten
+//|                                 by the pixels from the source
 //|         """
 //|
 STATIC mp_obj_t jpegio_jpegdecoder_decode(mp_uint_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
     jpegio_jpegdecoder_obj_t *self = MP_OBJ_TO_PTR(pos_args[0]);
 
-    enum { ARG_bitmap, ARG_scale };
+    enum { ARG_bitmap, ARG_scale, ARG_x, ARG_y, ARGS_X1_Y1_X2_Y2, ARG_skip_source_index, ARG_skip_dest_index };
     static const mp_arg_t allowed_args[] = {
         { MP_QSTR_bitmap, MP_ARG_OBJ | MP_ARG_REQUIRED, {.u_obj = mp_const_none } },
         { MP_QSTR_scale, MP_ARG_INT, {.u_int = 0 } },
+        { MP_QSTR_x, MP_ARG_INT, {.u_int = 0 } },
+        { MP_QSTR_y, MP_ARG_INT, {.u_int = 0 } },
+        ALLOWED_ARGS_X1_Y1_X2_Y2(0, 0),
+        {MP_QSTR_skip_source_index, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_obj = mp_const_none} },
+        {MP_QSTR_skip_dest_index, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_obj = mp_const_none} },
     };
     mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
     mp_arg_parse_all(n_args - 1, pos_args + 1, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
@@ -141,7 +175,33 @@ STATIC mp_obj_t jpegio_jpegdecoder_decode(mp_uint_t n_args, const mp_obj_t *pos_
     int scale = args[ARG_scale].u_int;
     mp_arg_validate_int_range(scale, 0, 3, MP_QSTR_scale);
 
-    common_hal_jpegio_jpegdecoder_decode_into(self, bitmap, scale);
+    int x = mp_arg_validate_int_range(args[ARG_x].u_int, 0, bitmap->width, MP_QSTR_x);
+    int y = mp_arg_validate_int_range(args[ARG_x].u_int, 0, bitmap->height, MP_QSTR_y);
+    bitmaptools_rect_t lim = bitmaptools_validate_coord_range_pair(&args[ARG_x1], bitmap->width, bitmap->height);
+
+    uint32_t skip_source_index;
+    bool skip_source_index_none; // flag whether skip_value was None
+
+    if (args[ARG_skip_source_index].u_obj == mp_const_none) {
+        skip_source_index = 0;
+        skip_source_index_none = true;
+    } else {
+        skip_source_index = mp_obj_get_int(args[ARG_skip_source_index].u_obj);
+        skip_source_index_none = false;
+    }
+
+    uint32_t skip_dest_index;
+    bool skip_dest_index_none; // flag whether skip_self_value was None
+
+    if (args[ARG_skip_dest_index].u_obj == mp_const_none) {
+        skip_dest_index = 0;
+        skip_dest_index_none = true;
+    } else {
+        skip_dest_index = mp_obj_get_int(args[ARG_skip_dest_index].u_obj);
+        skip_dest_index_none = false;
+    }
+    common_hal_jpegio_jpegdecoder_decode_into(self, bitmap, scale, x, y, &lim, skip_source_index, skip_source_index_none, skip_dest_index, skip_dest_index_none);
+
     return mp_const_none;
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_KW(jpegio_jpegdecoder_decode_obj, 1, jpegio_jpegdecoder_decode);

shared-bindings/jpegio/JpegDecoder.h

@@ -27,6 +27,7 @@
 #include "py/obj.h"
 #include "py/stream.h"
 #include "shared-module/displayio/Bitmap.h"
+#include "shared-bindings/bitmaptools/__init__.h"
 
 extern const mp_obj_type_t jpegio_jpegdecoder_type;
 
@@ -36,4 +37,9 @@ void common_hal_jpegio_jpegdecoder_construct(jpegio_jpegdecoder_obj_t *self);
 void common_hal_jpegio_jpegdecoder_close(jpegio_jpegdecoder_obj_t *self);
 mp_obj_t common_hal_jpegio_jpegdecoder_set_source_buffer(jpegio_jpegdecoder_obj_t *self, mp_obj_t jpeg_data);
 mp_obj_t common_hal_jpegio_jpegdecoder_set_source_file(jpegio_jpegdecoder_obj_t *self, mp_obj_t file_obj);
-void common_hal_jpegio_jpegdecoder_decode_into(jpegio_jpegdecoder_obj_t *self, displayio_bitmap_t *bitmap, int scale);
+void common_hal_jpegio_jpegdecoder_decode_into(
+    jpegio_jpegdecoder_obj_t *self,
+    displayio_bitmap_t *bitmap, int scale, int16_t x, int16_t y,
+    bitmaptools_rect_t *lim,
+    uint32_t skip_source_index, bool skip_source_index_none,
+    uint32_t skip_dest_index, bool skip_dest_index_none);

shared-module/jpegio/JpegDecoder.c

@@ -143,23 +143,90 @@ mp_obj_t common_hal_jpegio_jpegdecoder_set_source_buffer(jpegio_jpegdecoder_obj_
     return common_hal_jpegio_jpegdecoder_decode_common(self, buffer_input);
 }
 
+#define DECODER_CONTINUE (1)
+#define DECODER_INTERRUPT (0)
 static int bitmap_output(JDEC *jd, void *data, JRECT *rect) {
     jpegio_jpegdecoder_obj_t *self = CONTAINER_OF(jd, jpegio_jpegdecoder_obj_t, decoder);
-    common_hal_bitmaptools_arrayblit(self->dest, data, 2, rect->left, rect->top, rect->right + 1, rect->bottom + 1, false, 0);
+    int src_width = rect->right - rect->left + 1, src_pixel_stride = src_width /* in units of pixels! */, src_height = rect->bottom - rect->top + 1;
+
+    displayio_bitmap_t src = {
+        .width = src_width,
+        .height = src_height,
+        .data = data,
+        .stride = src_pixel_stride / 2, /* in units of uint32_t */
+        .bits_per_value = 16,
+        .x_shift = 1,
+        .x_mask = 1,
+        .bitmask = 0xffff,
+    };
+
+    int x = self->x;
+    int y = self->y;
+    int x1 = self->lim.x1 - rect->left;
+    int x2 = self->lim.x2 - rect->left;
+    int y1 = self->lim.y1 - rect->top;
+    int y2 = self->lim.y2 - rect->top;
+
+    if (y2 < y1) {
+        // The last row in the source image to copy FROM is above of this, so
+        // no more pixels on any rows
+        return DECODER_INTERRUPT;
+    }
+
+    y2 = MIN(y2, src_height);
+    if (x2 < x1) {
+        // The last column in the source image to copy FROM is left of this, so
+        // no more pixels on this row but could be on subsequent rows
+        return DECODER_CONTINUE;
+    }
+    x2 = MIN(x2, src_width);
+
+    // The first column in the source image to copy FROM is left of this, so copy starting with
+    // *local source* x1 equal to 0, and *target* x adjusted right by the same amount
+    if (x1 < 0) {
+        x += -x1;
+        x1 = 0;
+    }
+
+    // the same, but for Y coordinates
+    if (y1 < 0) {
+        y += -y1;
+        y1 = 0;
+    }
+
+    // blit takes care of x, y out of range
+    assert(x1 >= 0);
+    assert(y1 >= 0);
+    // blit takes care of x1 >= x2 and y1 >= y2 cases
+    assert(x2 <= src_width);
+    assert(y2 <= src_height);
+
+    common_hal_bitmaptools_blit(self->dest, &src, x, y, x1, y1, x2, y2, self->skip_source_index, self->skip_source_index_none, self->skip_dest_index, self->skip_dest_index_none);
     return 1;
 }
 
-void common_hal_jpegio_jpegdecoder_decode_into(jpegio_jpegdecoder_obj_t *self, displayio_bitmap_t *bitmap, int scale) {
+void common_hal_jpegio_jpegdecoder_decode_into(
+    jpegio_jpegdecoder_obj_t *self,
+    displayio_bitmap_t *bitmap, int scale, int16_t x, int16_t y,
+    bitmaptools_rect_t *lim,
+    uint32_t skip_source_index, bool skip_source_index_none,
+    uint32_t skip_dest_index, bool skip_dest_index_none) {
     if (self->data_obj == MP_OBJ_NULL) {
         mp_raise_RuntimeError_varg(MP_ERROR_TEXT("%q() without %q()"), MP_QSTR_decode, MP_QSTR_open);
     }
-    int dst_height = self->decoder.height >> scale;
-    int dst_width = self->decoder.width >> scale;
-    if (dst_width > bitmap->width || dst_height > bitmap->height) {
-        mp_raise_ValueError(MP_ERROR_TEXT("Destination bitmap too small to contain image"));
-    }
+
+    self->x = x;
+    self->y = y;
+    self->lim = *lim;
+    self->skip_source_index = skip_source_index;
+    self->skip_source_index_none = skip_source_index_none;
+    self->skip_dest_index = skip_dest_index;
+    self->skip_dest_index_none = skip_dest_index_none;
+
     self->dest = bitmap;
     JRESULT result = jd_decomp(&self->decoder, bitmap_output, scale);
     common_hal_jpegio_jpegdecoder_close(self);
-    check_jresult(result);
+    if (result != JDR_INTR) {
+        check_jresult(result);
+    }
 }

shared-module/jpegio/JpegDecoder.h

@@ -37,5 +37,9 @@ typedef struct jpegio_jpegdecoder_obj {
     mp_obj_t data_obj;
     mp_buffer_info_t bufinfo;
     displayio_bitmap_t *dest;
-    int scale;
+    uint16_t x, y;
+    bitmaptools_rect_t lim;
+    uint32_t skip_source_index, skip_dest_index;
+    bool skip_source_index_none, skip_dest_index_none;
+    uint8_t scale;
 } jpegio_jpegdecoder_obj_t;

tests/circuitpython/jpegio_decompress.py

@@ -3,6 +3,7 @@
 from displayio import Bitmap
 import binascii
 import jpegio
+import bitmaptools
 
 content = binascii.a2b_base64(
     b"""
@@ -54,21 +55,40 @@
 decoder = jpegio.JpegDecoder()
 
 
-def test(jpeg_input, scale):
+def dump_bitmap(b):
+    w = b.width
+    h = b.height
+
+    for y in range(min(30, h)):
+        for x in range(min(72, w)):
+            print(end="#" if (b[x, y] & 0x7E0) < 0x400 else " ")
+        print("..." if w > 72 else "")
+    print("...\n" if h > 24 else "")
+
+
+def test(jpeg_input, scale, fill=0xFFFF, **position_and_crop):
     w, h = decoder.open(jpeg_input)
     w >>= scale
     h >>= scale
     print(f"{w}x{h}")
 
     b = Bitmap(w, h, 65535)
+    b.fill(fill)
 
-    decoder.decode(b, scale=scale)
+    decoder.decode(b, scale=scale, **position_and_crop)
 
-    for y in range(min(30, h)):
-        for x in range(min(72, w)):
-            print(end="#" if (b[x, y] & 0x7E0) < 0x400 else " ")
-        print("..." if w > 72 else "")
-    print("...\n" if h > 24 else "")
+    dump_bitmap(b)
+
+    if position_and_crop:
+        position_and_crop.setdefault("x", 0)
+        position_and_crop.setdefault("y", 0)
+        full = Bitmap(w, h, 65535)
+        decoder.open(jpeg_input)
+        decoder.decode(full, scale=scale)
+        refb = Bitmap(w, h, 65535)
+        refb.fill(fill)
+        bitmaptools.blit(refb, full, **position_and_crop)
+        print(f"{memoryview(refb) == memoryview(b)=}")
 
 
 print("bytes")
@@ -84,3 +104,12 @@ def test(jpeg_input, scale):
 test(io.BytesIO(content), scale=1)
 test(io.BytesIO(content), scale=2)
 test(io.BytesIO(content), scale=3)
+
+print("crop & move")
+test(content, scale=3, x=8, y=8)
+test(content, scale=3, x1=8, y1=8)
+test(content, scale=3, x2=16, y2=16)
+test(content, scale=3, x=12, y=12, x1=8, y1=12, x2=16, y2=16)
+
+print("color key")
+test(content, scale=0, skip_source_index=0x4529, fill=0)