Add colorspace and ranges handling to streamer

8 files changed, 152 insertions, 49 deletions

diff --git a/chroma.glsl b/chroma.glsl
index 9004793..dd103b1 100644
--- a/chroma.glsl
+++ b/chroma.glsl
@@ -16,27 +16,26 @@
  */
 
 uniform sampler2D img_input;
-uniform mediump vec2 chroma_offsets;
+uniform mediump vec2 sample_offsets[4];
+uniform mediump mat3 colorspace;
+uniform mediump vec3 ranges[2];
 
 varying mediump vec2 texcoord;
 
-mediump vec2 rgb2chroma(in mediump vec4 rgb) {
-  // mburakov: This hardcodes BT.709 full-range.
-  mediump float y = rgb.r * 0.2126 + rgb.g * 0.7152 + rgb.b * 0.0722;
-  mediump float u = (rgb.b - y) / (2.0 * (1.0 - 0.0722));
-  mediump float v = (rgb.r - y) / (2.0 * (1.0 - 0.2126));
-  return vec2(u + 0.5, v + 0.5);
+mediump vec4 supersample() {
+  return texture2D(img_input, texcoord + sample_offsets[0]) +
+         texture2D(img_input, texcoord + sample_offsets[1]) +
+         texture2D(img_input, texcoord + sample_offsets[2]) +
+         texture2D(img_input, texcoord + sample_offsets[3]);
+}
+
+mediump vec3 rgb2yuv(in mediump vec3 rgb) {
+  mediump vec3 yuv = colorspace * rgb.rgb + vec3(0.0, 0.5, 0.5);
+  return ranges[0] + yuv * ranges[1];
 }
 
 void main() {
-  mediump vec2 sample_points[4];
-  sample_points[0] = texcoord;
-  sample_points[1] = texcoord + vec2(chroma_offsets.x, 0.0);
-  sample_points[2] = texcoord + vec2(0.0, chroma_offsets.y);
-  sample_points[3] = texcoord + chroma_offsets;
-  mediump vec4 rgb = texture2D(img_input, sample_points[0]) +
-                     texture2D(img_input, sample_points[1]) +
-                     texture2D(img_input, sample_points[2]) +
-                     texture2D(img_input, sample_points[3]);
-  gl_FragColor = vec4(rgb2chroma(rgb / 4.0), 0.0, 1.0);
+  mediump vec4 rgb = supersample() / 4.0;
+  mediump vec3 yuv = rgb2yuv(rgb.rgb);
+  gl_FragColor = vec4(yuv.yz, 0.0, 1.0);
 }
diff --git a/encode.c b/encode.c
index 5c7e886..4fe0af0 100644
--- a/encode.c
+++ b/encode.c
@@ -77,8 +77,33 @@ static bool SetHwFramesContext(struct EncodeContext* encode_context, int width,
   return true;
 }
 
+static enum AVColorSpace ConvertColorspace(enum YuvColorspace colorspace) {
+  switch (colorspace) {
+    case kItuRec601:
+      // TODO(mburakov): No dedicated definition for BT601?
+      return AVCOL_SPC_SMPTE170M;
+    case kItuRec709:
+      return AVCOL_SPC_BT709;
+    default:
+      __builtin_unreachable();
+  }
+}
+
+static enum AVColorRange ConvertRange(enum YuvRange range) {
+  switch (range) {
+    case kNarrowRange:
+      return AVCOL_RANGE_MPEG;
+    case kFullRange:
+      return AVCOL_RANGE_JPEG;
+    default:
+      __builtin_unreachable();
+  }
+}
+
 struct EncodeContext* EncodeContextCreate(struct GpuContext* gpu_context,
-                                          uint32_t width, uint32_t height) {
+                                          uint32_t width, uint32_t height,
+                                          enum YuvColorspace colrospace,
+                                          enum YuvRange range) {
   struct AUTO(EncodeContext)* encode_context =
       malloc(sizeof(struct EncodeContext));
   if (!encode_context) {
@@ -119,8 +144,8 @@ struct EncodeContext* EncodeContextCreate(struct GpuContext* gpu_context,
   encode_context->codec_context->max_b_frames = 0;
   encode_context->codec_context->refs = 1;
   encode_context->codec_context->global_quality = 18;
-  encode_context->codec_context->colorspace = AVCOL_SPC_BT709;
-  encode_context->codec_context->color_range = AVCOL_RANGE_JPEG;
+  encode_context->codec_context->colorspace = ConvertColorspace(colrospace);
+  encode_context->codec_context->color_range = ConvertRange(range);
 
   if (!SetHwFramesContext(encode_context, (int)width, (int)height)) {
     LOG("Failed to set hwframes context");
diff --git a/encode.h b/encode.h
index f68b2e2..085b0df 100644
--- a/encode.h
+++ b/encode.h
@@ -21,12 +21,16 @@
 #include <stdbool.h>
 #include <stdint.h>
 
+#include "colorspace.h"
+
 struct EncodeContext;
 struct GpuContext;
 struct GpuFrame;
 
 struct EncodeContext* EncodeContextCreate(struct GpuContext* gpu_context,
-                                          uint32_t width, uint32_t height);
+                                          uint32_t width, uint32_t height,
+                                          enum YuvColorspace colorspace,
+                                          enum YuvRange range);
 const struct GpuFrame* EncodeContextGetFrame(
     struct EncodeContext* encode_context);
 bool EncodeContextEncodeFrame(struct EncodeContext* encode_context, int fd);
diff --git a/gpu.c b/gpu.c
index 43f61b7..fed3f08 100644
--- a/gpu.c
+++ b/gpu.c
@@ -56,7 +56,7 @@ struct GpuContext {
   PFNGLEGLIMAGETARGETTEXTURE2DOESPROC glEGLImageTargetTexture2DOES;
   GLuint program_luma;
   GLuint program_chroma;
-  GLint chroma_offsets;
+  GLint sample_offsets;
   GLuint framebuffer;
   GLuint vertices;
 };
@@ -179,15 +179,71 @@ bail_out:
   return program;
 }
 
-static bool SetupImgInputUniform(GLuint program) {
-  GLint img_input = glGetUniformLocation(program, "img_input");
-  if (img_input == -1) {
-    LOG("Failed to find img_input uniform (%s)", GlErrorString(glGetError()));
-    return false;
+static const GLfloat* GetColorspaceMatrix(enum YuvColorspace colorspace) {
+  static const GLfloat rec601[] = {
+      _(0.299f, 0.587f, 0.114f),
+      _(-0.168736f, -0.331264f, 0.5f),
+      _(0.5f, -0.418688f, -0.081312f),
+  };
+  static const GLfloat rec709[] = {
+      _(0.2126f, 0.7152f, 0.0722f),
+      _(-0.1146f, -0.3854f, 0.5f),
+      _(0.5f, -0.4542f, -0.0458f),
+  };
+  switch (colorspace) {
+    case kItuRec601:
+      return rec601;
+    case kItuRec709:
+      return rec709;
+    default:
+      __builtin_unreachable();
+  }
+}
+
+static const GLfloat* GetRangeVectors(enum YuvRange range) {
+  static const GLfloat narrow[] = {
+      _(16.f / 255.f, 16.f / 255.f, 16.f / 255.f),
+      _((235.f - 16.f) / 255.f, (240.f - 16.f) / 255.f, (240.f - 16.f) / 255.f),
+  };
+  static const GLfloat full[] = {
+      _(0.f, 0.f, 0.f),
+      _(1.f, 1.f, 1.f),
+  };
+  switch (range) {
+    case kNarrowRange:
+      return narrow;
+    case kFullRange:
+      return full;
+    default:
+      __builtin_unreachable();
+  }
+}
+
+static bool SetupCommonUniforms(GLuint program, enum YuvColorspace colorspace,
+                                enum YuvRange range) {
+  struct {
+    const char* name;
+    GLint location;
+  } uniforms[] = {
+      {.name = "img_input"},
+      {.name = "colorspace"},
+      {.name = "ranges"},
+  };
+
+  for (size_t i = 0; i < LENGTH(uniforms); i++) {
+    uniforms[i].location = glGetUniformLocation(program, uniforms[i].name);
+    if (uniforms[i].location == -1) {
+      LOG("Failed to locate %s uniform (%s)", uniforms[i].name,
+          GlErrorString(glGetError()));
+      return false;
+    }
   }
 
   glUseProgram(program);
-  glUniform1i(img_input, 0);
+  glUniform1i(uniforms[0].location, 0);
+  glUniformMatrix3fv(uniforms[1].location, 1, GL_TRUE,
+                     GetColorspaceMatrix(colorspace));
+  glUniform3fv(uniforms[2].location, 2, GetRangeVectors(range));
   GLenum error = glGetError();
   if (error != GL_NO_ERROR) {
     LOG("Failed to set img_input uniform (%s)", GlErrorString(glGetError()));
@@ -196,7 +252,8 @@ static bool SetupImgInputUniform(GLuint program) {
   return true;
 }
 
-struct GpuContext* GpuContextCreate(void) {
+struct GpuContext* GpuContextCreate(enum YuvColorspace colorspace,
+                                    enum YuvRange range) {
   struct AUTO(GpuContext)* gpu_context = malloc(sizeof(struct GpuContext));
   if (!gpu_context) {
     LOG("Failed to allocate gpu context (%s)", strerror(errno));
@@ -208,7 +265,7 @@ struct GpuContext* GpuContextCreate(void) {
       .glEGLImageTargetTexture2DOES = NULL,
       .program_luma = 0,
       .program_chroma = 0,
-      .chroma_offsets = -1,
+      .sample_offsets = -1,
       .framebuffer = 0,
       .vertices = 0,
   };
@@ -255,11 +312,9 @@ struct GpuContext* GpuContextCreate(void) {
   }
 
   static const EGLint context_attribs[] = {
-#define _(...) __VA_ARGS__
       _(EGL_CONTEXT_MAJOR_VERSION, 3),
       _(EGL_CONTEXT_MINOR_VERSION, 1),
       EGL_NONE,
-#undef _
   };
   gpu_context->context = eglCreateContext(
       gpu_context->display, EGL_NO_CONFIG_KHR, EGL_NO_CONTEXT, context_attribs);
@@ -295,7 +350,7 @@ struct GpuContext* GpuContextCreate(void) {
       CreateGlProgram(_binary_vertex_glsl_start, _binary_vertex_glsl_end,
                       _binary_luma_glsl_start, _binary_luma_glsl_end);
   if (!gpu_context->program_luma ||
-      !SetupImgInputUniform(gpu_context->program_luma)) {
+      !SetupCommonUniforms(gpu_context->program_luma, colorspace, range)) {
     LOG("Failed to create luma program");
     return NULL;
   }
@@ -304,14 +359,14 @@ struct GpuContext* GpuContextCreate(void) {
       CreateGlProgram(_binary_vertex_glsl_start, _binary_vertex_glsl_end,
                       _binary_chroma_glsl_start, _binary_chroma_glsl_end);
   if (!gpu_context->program_chroma ||
-      !SetupImgInputUniform(gpu_context->program_luma)) {
+      !SetupCommonUniforms(gpu_context->program_chroma, colorspace, range)) {
     LOG("Failed to create chroma program");
     return NULL;
   }
-  gpu_context->chroma_offsets =
-      glGetUniformLocation(gpu_context->program_chroma, "chroma_offsets");
-  if (gpu_context->chroma_offsets == -1) {
-    LOG("Failed to find chroma_offsets uniform (%s)",
+  gpu_context->sample_offsets =
+      glGetUniformLocation(gpu_context->program_chroma, "sample_offsets");
+  if (gpu_context->sample_offsets == -1) {
+    LOG("Failed to find sample_offsets uniform (%s)",
         GlErrorString(glGetError()));
     return NULL;
   }
@@ -380,11 +435,9 @@ static EGLImage CreateEglImage(struct GpuContext* gpu_context, uint32_t width,
   };
 
   EGLAttrib attrib_list[7 + LENGTH(attrib_keys) * 2] = {
-#define _(...) __VA_ARGS__
       _(EGL_WIDTH, width),
       _(EGL_HEIGHT, height),
       _(EGL_LINUX_DRM_FOURCC_EXT, fourcc),
-#undef _
   };
 
   EGLAttrib* pairs = &attrib_list[6];
@@ -533,9 +586,15 @@ bool GpuFrameConvert(const struct GpuFrame* from, const struct GpuFrame* to) {
     return false;
   }
 
+  const GLfloat sample_offsets[] = {
+      _(0.f, 0.f),
+      _(1.f / (GLfloat)from->width, 0.f),
+      _(0.f, 1.f / (GLfloat)from->height),
+      _(1.f / (GLfloat)from->width, 1.f / (GLfloat)from->height),
+  };
+
   glUseProgram(from->gpu_context->program_chroma);
-  glUniform2f(from->gpu_context->chroma_offsets, 1.f / (GLfloat)from->width,
-              1.f / (GLfloat)from->height);
+  glUniform2fv(from->gpu_context->sample_offsets, 4, sample_offsets);
   glViewport(0, 0, (GLsizei)to->width / 2, (GLsizei)to->height / 2);
   if (!GpuFrameConvertImpl(from->textures[0], to->textures[1])) {
     LOG("Failed to convert chroma plane");
diff --git a/gpu.h b/gpu.h
index a57ada0..0043aca 100644
--- a/gpu.h
+++ b/gpu.h
@@ -22,6 +22,8 @@
 #include <stddef.h>
 #include <stdint.h>
 
+#include "colorspace.h"
+
 struct GpuFramePlane {
   int dmabuf_fd;
   uint32_t pitch;
@@ -29,7 +31,8 @@ struct GpuFramePlane {
   uint64_t modifier;
 };
 
-struct GpuContext* GpuContextCreate(void);
+struct GpuContext* GpuContextCreate(enum YuvColorspace colorspace,
+                                    enum YuvRange range);
 bool GpuContextSync(struct GpuContext* gpu_context);
 void GpuContextDestroy(struct GpuContext** gpu_context);
 
diff --git a/luma.glsl b/luma.glsl
index cc56377..4dac957 100644
--- a/luma.glsl
+++ b/luma.glsl
@@ -16,15 +16,18 @@
  */
 
 uniform sampler2D img_input;
+uniform mediump mat3 colorspace;
+uniform mediump vec3 ranges[2];
 
 varying mediump vec2 texcoord;
 
-mediump float rgb2luma(in mediump vec4 rgb) {
-  // mburakov: This hardcodes BT.709 full-range.
-  return rgb.r * 0.2126 + rgb.g * 0.7152 + rgb.b * 0.0722;
+mediump vec3 rgb2yuv(in mediump vec3 rgb) {
+  mediump vec3 yuv = colorspace * rgb.rgb + vec3(0.0, 0.5, 0.5);
+  return ranges[0] + yuv * ranges[1];
 }
 
 void main() {
   mediump vec4 rgb = texture2D(img_input, texcoord);
-  gl_FragColor = vec4(rgb2luma(rgb), 0.0, 0.0, 1.0);
+  mediump vec3 yuv = rgb2yuv(rgb.rgb);
+  gl_FragColor = vec4(yuv.x, 0.0, 0.0, 1.0);
 }
diff --git a/main.c b/main.c
index 8e63703..8d7a42c 100644
--- a/main.c
+++ b/main.c
@@ -25,10 +25,18 @@
 #include <unistd.h>
 
 #include "capture.h"
+#include "colorspace.h"
 #include "encode.h"
 #include "gpu.h"
 #include "util.h"
 
+// TODO(mburakov): Currently zwp_linux_dmabuf_v1 has no way to provide
+// colorspace and range information to the compositor. Maybe this would change
+// in the future, i.e keep an eye on color-representation Wayland protocol:
+// https://gitlab.freedesktop.org/wayland/wayland-protocols/-/merge_requests/183
+static const enum YuvColorspace colorspace = kItuRec601;
+static const enum YuvRange range = kNarrowRange;
+
 static volatile sig_atomic_t g_signal;
 static void OnSignal(int status) { g_signal = status; }
 
@@ -69,7 +77,7 @@ int main(int argc, char* argv[]) {
     return EXIT_FAILURE;
   }
 
-  struct AUTO(GpuContext)* gpu_context = GpuContextCreate();
+  struct AUTO(GpuContext)* gpu_context = GpuContextCreate(colorspace, range);
   if (!gpu_context) {
     LOG("Failed to create gpu context");
     return EXIT_FAILURE;
@@ -104,7 +112,8 @@ int main(int argc, char* argv[]) {
     if (!encode_context) {
       uint32_t width, height;
       GpuFrameGetSize(captured_frame, &width, &height);
-      encode_context = EncodeContextCreate(gpu_context, width, height);
+      encode_context =
+          EncodeContextCreate(gpu_context, width, height, colorspace, range);
       if (!encode_context) {
         LOG("Failed to create encode context");
         return EXIT_FAILURE;
diff --git a/util.h b/util.h
index 18e8c5e..47c4008 100644
--- a/util.h
+++ b/util.h
@@ -27,6 +27,7 @@
 #define LENGTH(x) (sizeof(x) / sizeof *(x))
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
+#define _(...) __VA_ARGS__
 
 static inline void* Release(void** x) {
   void* result = *x;