Add chroma format support to kvz_picture

Add picture_alloc_csp to libkvz api to allocated pictures with chroma
format different from 420.

src/encmain.c

@@ -98,14 +98,15 @@ static unsigned get_padding(unsigned width_or_height){
  */
 static void compute_psnr(const kvz_picture *const src,
                          const kvz_picture *const rec,
-                         double psnr[NUM_COLORS])
+                         double psnr[3])
 {
   assert(src->width  == rec->width);
   assert(src->height == rec->height);
 
   int32_t pixels = src->width * src->height;
+  int colors = rec->chroma_format == KVZ_CSP_400 ? 1 : 3;
 
-  for (int32_t c = 0; c < NUM_COLORS; ++c) {
+  for (int32_t c = 0; c < colors; ++c) {
     int32_t num_pixels = pixels;
     if (c != COLOR_Y) {
       num_pixels >>= 2;
@@ -178,8 +179,10 @@ static void* input_read_thread(void* in_args)
       goto done;
     }
 
-    frame_in = args->api->picture_alloc(args->opts->config->width  + args->padding_x,
-                                        args->opts->config->height + args->padding_y);
+    enum kvz_chroma_format csp = KVZ_FORMAT2CSP(args->opts->config->input_format);
+    frame_in = args->api->picture_alloc_csp(csp,
+                                            args->opts->config->width  + args->padding_x,
+                                            args->opts->config->height + args->padding_y);
 
     if (!frame_in) {
       fprintf(stderr, "Failed to allocate image.\n");

src/encoderstate.c

@@ -791,7 +791,7 @@ static void encoder_set_source_picture(encoder_state_t * const state, kvz_pictur
     // In lossless mode, the reconstruction is equal to the source frame.
     state->tile->frame->rec = kvz_image_copy_ref(frame);
   } else {
-    state->tile->frame->rec = kvz_image_alloc(frame->width, frame->height);
+    state->tile->frame->rec = kvz_image_alloc(state->encoder_control->chroma_format, frame->width, frame->height);
   }
 
   kvz_videoframe_set_poc(state->tile->frame, state->frame->poc);

src/image.c

@@ -26,11 +26,21 @@
 #include "strategies/strategies-picture.h"
 #include "threads.h"
 
+/**
+* \brief Allocate a new image with 420.
+* This function signature is part of the libkvz API.
+* \return image pointer or NULL on failure
+*/
+kvz_picture * kvz_image_alloc_420(const int32_t width, const int32_t height)
+{
+  return kvz_image_alloc(KVZ_CSP_420, width, height);
+}
+
 /**
  * \brief Allocate a new image.
  * \return image pointer or NULL on failure
  */
-kvz_picture *kvz_image_alloc(const int32_t width, const int32_t height)
+kvz_picture * kvz_image_alloc(enum kvz_chroma_format chroma_format, const int32_t width, const int32_t height)
 {
   //Assert that we have a well defined image
   assert((width % 2) == 0);
@@ -40,7 +50,10 @@ kvz_picture *kvz_image_alloc(const int32_t width, const int32_t height)
   if (!im) return NULL;
 
   unsigned int luma_size = width * height;
-  unsigned int chroma_size = luma_size / 4;
+  unsigned chroma_sizes[] = { 0, luma_size / 4, luma_size / 2, luma_size };
+  unsigned chroma_size = chroma_sizes[chroma_format];
+
+  im->chroma_format = chroma_format;
 
   //Allocate memory
   im->fulldata = MALLOC(kvz_pixel, (luma_size + 2 * chroma_size));
@@ -54,10 +67,17 @@ kvz_picture *kvz_image_alloc(const int32_t width, const int32_t height)
   im->width = width;
   im->height = height;
   im->stride = width;
+  im->chroma_format = chroma_format;
 
   im->y = im->data[COLOR_Y] = &im->fulldata[0];
-  im->u = im->data[COLOR_U] = &im->fulldata[luma_size];
-  im->v = im->data[COLOR_V] = &im->fulldata[luma_size + chroma_size];
+
+  if (chroma_format == KVZ_CSP_400) {
+    im->u = im->data[COLOR_U] = NULL;
+    im->v = im->data[COLOR_V] = NULL;
+  } else {
+    im->u = im->data[COLOR_U] = &im->fulldata[luma_size];
+    im->v = im->data[COLOR_V] = &im->fulldata[luma_size + chroma_size];
+  }
 
   im->pts = 0;
   im->dts = 0;
@@ -138,10 +158,13 @@ kvz_picture *kvz_image_make_subimage(kvz_picture *const orig_image,
   im->width = width;
   im->height = height;
   im->stride = orig_image->stride;
+  im->chroma_format = orig_image->chroma_format;
 
   im->y = im->data[COLOR_Y] = &orig_image->y[x_offset + y_offset * orig_image->stride];
-  im->u = im->data[COLOR_U] = &orig_image->u[x_offset/2 + y_offset/2 * orig_image->stride/2];
-  im->v = im->data[COLOR_V] = &orig_image->v[x_offset/2 + y_offset/2 * orig_image->stride/2];
+  if (orig_image->chroma_format != KVZ_CSP_400) {
+    im->u = im->data[COLOR_U] = &orig_image->u[x_offset / 2 + y_offset / 2 * orig_image->stride / 2];
+    im->v = im->data[COLOR_V] = &orig_image->v[x_offset / 2 + y_offset / 2 * orig_image->stride / 2];
+  }
 
   im->pts = 0;
   im->dts = 0;

src/image.h

@@ -52,7 +52,8 @@ typedef struct {
 } yuv_t;
 
 
-kvz_picture *kvz_image_alloc(const int32_t width, const int32_t height);
+kvz_picture *kvz_image_alloc_420(const int32_t width, const int32_t height);
+kvz_picture *kvz_image_alloc(enum kvz_chroma_format chroma_format, const int32_t width, const int32_t height);
 
 void kvz_image_free(kvz_picture *im);
 

src/kvazaar.c

@@ -287,11 +287,11 @@ static int kvazaar_field_encoding_adapter(kvz_encoder *enc,
   } first = { 0, 0 }, second = { 0, 0 };
 
   if (pic_in != NULL) {
-    first_field = kvz_image_alloc(state->encoder_control->in.width, state->encoder_control->in.height);
+    first_field = kvz_image_alloc(state->encoder_control->chroma_format, state->encoder_control->in.width, state->encoder_control->in.height);
     if (first_field == NULL) {
       goto kvazaar_field_encoding_adapter_failure;
     }
-    second_field = kvz_image_alloc(state->encoder_control->in.width, state->encoder_control->in.height);
+    second_field = kvz_image_alloc(state->encoder_control->chroma_format, state->encoder_control->in.width, state->encoder_control->in.height);
     if (second_field == NULL) {
       goto kvazaar_field_encoding_adapter_failure;
     }
@@ -355,7 +355,7 @@ static const kvz_api kvz_8bit_api = {
   .config_destroy = kvz_config_destroy,
   .config_parse = kvz_config_parse,
 
-  .picture_alloc = kvz_image_alloc,
+  .picture_alloc = kvz_image_alloc_420,
   .picture_free = kvz_image_free,
 
   .chunk_free = kvz_bitstream_free_chunks,
@@ -364,6 +364,8 @@ static const kvz_api kvz_8bit_api = {
   .encoder_close = kvazaar_close,
   .encoder_headers = kvazaar_headers,
   .encoder_encode = kvazaar_field_encoding_adapter,
+
+  .picture_alloc_csp = kvz_image_alloc,
 };
 
 

src/kvazaar.h

@@ -339,6 +339,7 @@ typedef struct kvz_picture {
   int64_t dts;             //!< \brief Decompression timestamp.
 
   enum kvz_interlacing interlacing; //!< \since 3.2.0 \brief Field order for interlaced pictures.
+  enum kvz_chroma_format chroma_format;
 } kvz_picture;
 
 /**
@@ -616,6 +617,19 @@ typedef struct kvz_api {
                                   kvz_picture **pic_out,
                                   kvz_picture **src_out,
                                   kvz_frame_info *info_out);
+
+  /**
+   * \brief Allocate a kvz_picture.
+   *
+   * The returned kvz_picture should be deallocated by calling picture_free.
+   *
+   * \since 3.12.0
+   * \param chroma_fomat  Chroma subsampling to use.
+   * \param width   width of luma pixel array to allocate
+   * \param height  height of luma pixel array to allocate
+   * \return        allocated picture, or NULL if allocation failed.
+   */
+  kvz_picture * (*picture_alloc_csp)(enum kvz_chroma_format chroma_fomat, int32_t width, int32_t height);
 } kvz_api;
 
 

tests/sad_tests.c

@@ -64,12 +64,12 @@ static kvz_picture *g_ref = 0;
 // SETUP, TEARDOWN AND HELPER FUNCTIONS
 static void setup_tests()
 {
-  g_pic = kvz_image_alloc(8, 8);
+  g_pic = kvz_image_alloc(KVZ_CSP_420, 8, 8);
   for (int i = 0; i < 64; ++i) {
     g_pic->y[i] = pic_data[i] + 48;
   }
 
-  g_ref = kvz_image_alloc(8, 8);
+  g_ref = kvz_image_alloc(KVZ_CSP_420, 8, 8);
   for (int i = 0; i < 64; ++i) {
     g_ref->y[i] = ref_data[i] + 48;
   }