Merge pull request #3 from BinaryNoggin/jpeg-to-yuv

[API BREAKING] Add feature to extract pixel data from jpegs

Author: adkron

c_src/turbojpeg_native/turbojpeg_native.c

@@ -4,55 +4,199 @@
  * Supported pixel formats: :I420 | :I422 | :I444 
  * Unsupported pixel formats: :RGB | :BGRA | :RGBA | :NV12 | :NV21 | :YV12 | :AYUV
 */
-UNIFEX_TERM create(UnifexEnv *env, int width, int height, int jpegQuality, char *format) {
-  UNIFEX_TERM res;
-  State *state = unifex_alloc_state(env);
-  state->width = width; state->height = height; 
-  state->flags = 0; state->quality = jpegQuality;
+int format_to_tjsamp(char* format) {
   if(strcmp(format, "I420") == 0) {
-    state->format = TJSAMP_420;
+    return TJSAMP_420;
   } else if(strcmp(format, "I422") == 0) {
-    state->format = TJSAMP_422;
+    return TJSAMP_422;
   } else if(strcmp(format, "I444") == 0) {
-    state->format = TJSAMP_444;
+    return TJSAMP_444;
   } else {
-    res = create_result_error(env, "unsupported_format");
-    unifex_release_state(env, state);
-    return(res); 
+    return -1;
   }
-  res = create_result_ok(env, state);
-  unifex_release_state(env, state);
-  return(res);
+}
+
+/** 
+ * Supported pixel formats: :I420 | :I422 | :I444 
+ * Unsupported pixel formats: :RGB | :BGRA | :RGBA | :NV12 | :NV21 | :YV12 | :AYUV
+*/
+const char* tjsamp_to_format(enum TJSAMP tjsamp) {
+  switch(tjsamp) {
+    case(TJSAMP_420):
+      return("I420");
+    case(TJSAMP_422):
+      return("I422");
+    case(TJSAMP_444):
+      return("I444");
+    default:
+      return("unknown_format");
+  }
+}
+
+/**
+ * Returns %{height: int, width: int, format: atom}
+*/
+UNIFEX_TERM get_jpeg_header(UnifexEnv* env, UnifexPayload *payload) {
+  tjhandle tjh;
+  enum TJSAMP tjsamp;
+  enum TJCS cspace;
+  int res, width, height;
+  ERL_NIF_TERM map_out, ret;
+
+  tjh = tjInitDecompress();
+  if(!tjh)
+    return get_jpeg_header_result_error(env, tjGetErrorStr());
+
+  res = tjDecompressHeader3(
+    tjh,
+    payload->data,
+    payload->size, 
+    &width, &height, 
+    (int*)&tjsamp, (int*)&cspace
+  );
+  if(res < 0) {
+    ret = get_jpeg_header_result_error(env, tjGetErrorStr2(tjh));
+    goto cleanup;
+  }
+
+  // unifex does not support maps yes.
+  // See https://github.com/membraneframework/unifex/issues/32
+  if(!enif_make_map_from_arrays(
+    env,
+    (ERL_NIF_TERM []) {
+      enif_make_atom(env, "width"),
+      enif_make_atom(env, "height"),
+      enif_make_atom(env, "format"),
+
+    },
+    (ERL_NIF_TERM []) {
+      enif_make_int(env, width),
+      enif_make_int(env, height),
+      enif_make_atom(env, tjsamp_to_format(tjsamp))
+    },
+    3, &map_out
+  )) {
+    ret = get_jpeg_header_result_error(env, "make_map");
+    goto cleanup;
+  } else {
+    // Generated code does not support maps currently.
+    ret = enif_make_tuple_from_array(env, (ERL_NIF_TERM []) {enif_make_atom(env, "ok"), map_out}, 2);
+    goto cleanup;
+  }
+cleanup:
+  if(tjh) tjDestroy(tjh);
+  return ret;
 }
 
 /**
  * Convert a binary h264 payload into a jpeg encoded payload
 */
-UNIFEX_TERM to_jpeg(UnifexEnv* env, UnifexPayload *payload, State *state) {
-  state->tjh = tjInitCompress();
-  
+UNIFEX_TERM yuv_to_jpeg(UnifexEnv* env, UnifexPayload *payload, int width, int height, int quality, char* format) {
+  tjhandle tjh = NULL;
+  enum TJSAMP tjsamp;
   unsigned char *jpegBuf = NULL;
   unsigned long jpegSize;
+  int res;
+  UnifexPayload *jpegFrame;
+  UNIFEX_TERM ret;
 
-  int res = tjCompressFromYUV(
-    state->tjh, 
+  res = format_to_tjsamp(format);
+  if(res < 0) {
+    return(yuv_to_jpeg_result_error(env, "unsupported_format")); 
+  } else {
+    tjsamp = (enum TJSAMP)res;
+  }
+
+  tjh = tjInitCompress();
+  if(!tjh)
+    return yuv_to_jpeg_result_error(env, tjGetErrorStr());
+  
+  res = tjCompressFromYUV(
+    tjh, 
     payload->data, 
-    state->width, 4, state->height, state->format, 
+    width, 4, height, tjsamp, 
     &jpegBuf, &jpegSize, 
-    state->quality, state->flags
+    quality, 0
   );
 
-  if(res) {
-    return to_jpeg_result_error(env, tjGetErrorStr());
+  if(res < 0) {
+    ret = yuv_to_jpeg_result_error(env, tjGetErrorStr2(tjh));
+    goto cleanup;
   } else {
-    UnifexPayload *jpegFrame = unifex_payload_alloc(env, UNIFEX_PAYLOAD_SHM, jpegSize);
-    memcpy(jpegFrame->data, jpegBuf, jpegSize);
-    tjFree(jpegBuf);
-    return to_jpeg_result_ok(env, jpegFrame);
+    jpegFrame = unifex_payload_alloc(env, UNIFEX_PAYLOAD_SHM, jpegSize);
+    if(!jpegFrame) {
+      ret = yuv_to_jpeg_result_error(env, "payload_alloc");
+    } else {
+      memcpy(jpegFrame->data, jpegBuf, jpegSize);
+      ret = yuv_to_jpeg_result_ok(env, jpegFrame);
+    }
+    goto cleanup;
   }
+
+cleanup:
+  if(jpegBuf) tjFree(jpegBuf);
+  if(tjh) tjDestroy(tjh);
+  return ret;
+}
+
+/**
+ * Convert a binary jpeg payload into a yuv encoded payload
+*/
+UNIFEX_TERM jpeg_to_yuv(UnifexEnv* env, UnifexPayload *payload) {
+  tjhandle tjh;
+  enum TJSAMP tjsamp;
+  enum TJCS cspace;
+  unsigned long yuvBufSize;
+  UnifexPayload *yuvFrame;
+  int res, width, height;
+  UNIFEX_TERM ret;
+
+  tjh = tjInitDecompress();
+  if(!tjh)
+    return jpeg_to_yuv_result_error(env, tjGetErrorStr());
+
+  res = tjDecompressHeader3(
+    tjh,
+    payload->data,
+    payload->size, 
+    &width, &height, 
+    (int*)&tjsamp, (int*)&cspace
+  );
+  if(res < 0) {
+    ret = jpeg_to_yuv_result_error(env, tjGetErrorStr2(tjh));
+    goto cleanup;
+  } 
+  
+  yuvBufSize = tjBufSizeYUV2(width, 4, height, tjsamp);
+  yuvFrame = unifex_payload_alloc(env, UNIFEX_PAYLOAD_SHM, yuvBufSize);
+  if(!yuvFrame) {
+    ret = jpeg_to_yuv_result_error(env, "could not allocate frame");
+    goto cleanup;
+  }
+
+  res = tjDecompressToYUV2(
+    tjh, 
+    payload->data,
+    payload->size,
+    yuvFrame->data,
+    width, 4, height,
+    0
+	);
+
+  if(res < 0) {
+    ret = jpeg_to_yuv_result_error(env, tjGetErrorStr2(tjh));
+    goto cleanup;
+  } else {
+    ret = jpeg_to_yuv_result_ok(env, yuvFrame);
+    goto cleanup;
+  }
+
+cleanup:
+  if(tjh) tjDestroy(tjh);
+  return ret;
 }
 
 void handle_destroy_state(UnifexEnv* env, State* state) {
   UNIFEX_UNUSED(env);
-  if(state->tjh) { tjDestroy(state->tjh); }
+  UNIFEX_UNUSED(state);
 }

c_src/turbojpeg_native/turbojpeg_native.h

@@ -4,15 +4,10 @@
 #include <turbojpeg.h>
 
 /** NIF State */
-typedef struct _h264_parser_state {
-  /** Handle to turbojpeg */
-  tjhandle tjh;
-  enum TJSAMP format;
-  int flags;
-  int width; int height; int quality;
+typedef struct _turbojpeg_native_state {
 } UnifexNifState;
 
-/** NIF State */
+// /** NIF State */
 typedef UnifexNifState State;
 
 #include "_generated/turbojpeg_native.h"

c_src/turbojpeg_native/turbojpeg_native.spec.exs

@@ -1,4 +1,5 @@
 module Turbojpeg.Native
 
-spec create(width::int, height::int, quality::int, format::atom) :: {:ok :: label, state} | {:error :: label, reason :: atom}
-spec to_jpeg(payload, state) :: {:ok :: label, payload} | {:error :: label, reason :: atom}
+spec yuv_to_jpeg(payload, width::int, height::int, quality::int, format::atom) :: {:ok :: label, payload} | {:error :: label, reason :: atom}
+spec jpeg_to_yuv(payload) :: {:ok :: label, payload} | {:error :: label, reason :: atom}
+spec get_jpeg_header(payload) :: {:ok :: label, data::int} | {:error::label, reason::atom}

fixture/ff0000_i444.jpg

@@ -15,8 +15,10 @@ defmodule Turbojpeg.Sink do
   def handle_init(options) do
     state = %{
       quality: options.quality,
+      height: nil,
+      width: nil,
+      format: nil,
       filename: options.filename,
-      jpeg_ref: nil,
       timer_started?: false
     }
 
@@ -37,16 +39,16 @@ defmodule Turbojpeg.Sink do
     %{input: input} = ctx.pads
 
     if !input.caps || caps == input.caps do
-      {:ok, ref} = Native.create(caps.width, caps.height, state.quality, caps.format)
-      {:ok, %{state | jpeg_ref: ref}}
+      {:ok, %{state | width: caps.width, height: caps.height, format: caps.format}}
     else
       raise "Caps have changed while playing. This is not supported."
     end
   end
 
   @impl true
   def handle_write(:input, %Buffer{payload: payload}, _ctx, state) do
-    with {:ok, data} <- Native.to_jpeg(payload, state.jpeg_ref),
+    with {:ok, data} <-
+           Native.yuv_to_jpeg(payload, state.width, state.height, state.quality, state.format),
          :ok <- File.write(state.filename, Shmex.to_binary(data)) do
       {:ok, state}
     else

lib/turbojpeg/sink.ex

@@ -2,12 +2,27 @@ defmodule TurbojpegTest do
   use ExUnit.Case
   @jpeg_header <<255, 216, 255>>
   @i420_fixture "fixture/i420.yuv"
+  @ff0000_fixture "fixture/ff0000_i444.jpg"
 
   test "Converts an i420 frame into a jpeg" do
-    {:ok, native} = Turbojpeg.Native.create(1920, 1080, 75, :I420)
     frame = File.read!(@i420_fixture)
     shmex = Shmex.new(frame)
-    {:ok, jpeg} = Turbojpeg.Native.to_jpeg(shmex, native)
+    {:ok, jpeg} = Turbojpeg.Native.yuv_to_jpeg(shmex, 1920, 1080, 100, :I420)
     assert match?(@jpeg_header <> _, Shmex.to_binary(jpeg))
   end
+
+  test "extracts i444 frame from jpeg" do
+    jpeg = Shmex.new(File.read!(@ff0000_fixture))
+    {:ok, yuv} = Turbojpeg.Native.jpeg_to_yuv(jpeg)
+    {:ok, new_jpeg} = Turbojpeg.Native.yuv_to_jpeg(yuv, 64, 64, 100, :I444)
+    assert Shmex.to_binary(jpeg) == Shmex.to_binary(new_jpeg)
+  end
+
+  test "get jpeg header" do
+    jpeg = Shmex.new(File.read!(@ff0000_fixture))
+    {:ok, result} = Turbojpeg.Native.get_jpeg_header(jpeg)
+    assert result.width == 64
+    assert result.height == 64
+    assert result.format == :I444
+  end
 end