YUV420P grabber

YUV420PGrabber.cpp

#include <assert.h>
#include <roxlu/core/Utils.h>
#include <roxlu/core/Log.h>
#include <image/Image.h>
#include <video/YUV420PGrabber.h>

YUV420PGrabber::YUV420PGrabber() 
  :y_prog(0)
  ,y_vert(0)
  ,y_frag(0)
  ,uv_prog(0)
  ,uv_frag(0)
  ,pt_prog(0)
  ,pt_frag(0)
  ,scene_fbo(0)
  ,scene_depth(0)
  ,scene_tex(0)
  ,y_tex(0)
  ,u_tex(0)
  ,v_tex(0)
  ,window_w(0)
  ,window_h(0)
  ,video_w(0)
  ,video_h(0)
  ,uv_w(0)
  ,uv_h(0)
  ,vbo(0)
  ,vao(0)
  ,image(NULL)
  ,cb_frame(NULL)
  ,cb_user(NULL)
  ,millis_per_frame(0)
  ,frame_timeout(0)
  ,state(YUV420_STATE_NONE)

#if YUV420_MODE == YUV420_MODE_PBO
  ,read_dx(0)
  ,write_dx(0)
#elif YUV420_MODE == YUV420_MODE_GLREAD
  ,y_buf(NULL)
  ,u_buf(NULL)
  ,v_buf(NULL)
#endif
{
}

YUV420PGrabber::~YUV420PGrabber() {
  if(y_vert) {
    glDeleteShader(y_vert);
  }

  if(y_frag) {
    glDeleteShader(y_frag);
  }

  if(uv_frag) {
    glDeleteShader(uv_frag);
  }

  if(y_prog) {
    glDeleteProgram(y_prog);
  }

  if(uv_prog) {
    glDeleteProgram(uv_prog);
  }

  if(scene_fbo) {
    glDeleteFramebuffers(1, &scene_fbo);
  }

  if(scene_depth) {
    glDeleteRenderbuffers(1, &scene_depth);
  }

  if(scene_tex) {
    glDeleteTextures(1, &scene_tex);
  }

  if(y_tex) {
    glDeleteTextures(1, &y_tex);
  }

  if(u_tex) {
    glDeleteTextures(1, &u_tex);
  }

  if(v_tex) {
    glDeleteTextures(1, &v_tex);
  }

  if(vbo) {
    glDeleteBuffers(1, &vbo);
  }

  if(vao) {
    glDeleteVertexArrays(1, &vao);
  }

  if(pt_prog) {
    glDeleteProgram(pt_prog);
  }

  if(pt_frag) {
    glDeleteShader(pt_frag);
  }

#if YUV420_MODE == YUV420_MODE_PBO
  RX_ERROR("DELETE PBOS");
#elif YUV420_MODE == YUV420_MODE_GLREAD
  if(image) {
    delete[] image;
  }
#endif
  
  window_w = 0;
  window_h = 0;
  video_w = 0;
  video_h = 0;
  uv_w = 0;
  uv_h = 0;
  y_prog = 0;
  y_vert = 0;
  y_frag = 0;
  uv_prog = 0;
  uv_frag = 0;
  pt_prog = 0;
  pt_frag = 0;
  scene_fbo = 0;
  scene_depth = 0;
  scene_tex = 0;
  y_tex = 0;
  u_tex = 0;
  v_tex = 0;
  vbo = 0;
  vao = 0;
  cb_frame = NULL;
  cb_user = NULL;
  image = NULL;
  millis_per_frame = 0;
  frame_timeout = 0;

#if YUV420_MODE == YUV420_MODE_PBO
  read_dx = 0;
  write_dx = 0;
#elif YUV420_MODE == YUV420_MODE_GLREAD
  y_buf = NULL;
  u_buf = NULL;
  v_buf = NULL;
#endif
}

bool YUV420PGrabber::setup(int winW, int winH, int vidW, int vidH, int fps,
                           yuv420p_on_frame frameCB, void* frameUser) 
{
  assert(winW);
  assert(winH);
  assert(vidW);
  assert(vidH);
  assert(frameCB);

  cb_frame = frameCB;
  cb_user = frameUser;
  window_w = winW;
  window_h = winH;
  video_w = vidW;
  video_h = vidH;
  uv_w = video_w * 0.5;
  uv_h = video_h * 0.5;

  millis_per_frame = (1.0 / double(fps)) * 1000;
  
  if(!setupFBO()) {
    return false;
  }

  if(!setupShaders()) {
    return false;
  }

  if(!setupBuffers()) {
    return false;
  }

#if YUV420_MODE == YUV420_MODE_PBO
  if(!setupPBOs()) {
    return false;
  }
#elif YUV420_MODE == YUV420_MODE_GLREAD
  int num_bytes = video_w * video_h + ((uv_w * uv_h) * 2);
  image = new unsigned char[num_bytes];
  y_buf = image;
  u_buf = image + video_w * video_h;
  v_buf = u_buf + uv_w * uv_h;
#endif

  return true;

}

bool YUV420PGrabber::setupFBO() {
  assert(window_h && window_w);
  assert(video_w && video_h);
  
  glGenFramebuffers(1, &scene_fbo);
  glBindFramebuffer(GL_FRAMEBUFFER, scene_fbo);
  
  glGenRenderbuffers(1, &scene_depth);
  glBindRenderbuffer(GL_RENDERBUFFER, scene_depth);
  glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH_COMPONENT32, window_w, window_h);
  glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT,  GL_RENDERBUFFER, scene_depth);

  glGenTextures(1, &scene_tex);
  glBindTexture(GL_TEXTURE_2D, scene_tex);
  glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA8, window_w, window_h, 0, GL_RGB, GL_UNSIGNED_BYTE, NULL);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
  glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, scene_tex, 0);

  glGenTextures(1, &y_tex);
  glBindTexture(GL_TEXTURE_2D, y_tex);
  glTexImage2D(GL_TEXTURE_2D, 0, GL_R8, video_w, video_h, 0, GL_RED, GL_UNSIGNED_BYTE, NULL);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
  glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT1, GL_TEXTURE_2D, y_tex, 0);

  glGenTextures(1, &u_tex);
  glBindTexture(GL_TEXTURE_2D, u_tex);
  glTexImage2D(GL_TEXTURE_2D, 0, GL_R8, uv_w, uv_h, 0, GL_RED, GL_UNSIGNED_BYTE, NULL);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
  glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT2, GL_TEXTURE_2D, u_tex, 0);

  glGenTextures(1, &v_tex);
  glBindTexture(GL_TEXTURE_2D, v_tex);
  glTexImage2D(GL_TEXTURE_2D, 0, GL_R8, uv_w, uv_h, 0, GL_RED, GL_UNSIGNED_BYTE, NULL);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
  glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT3, GL_TEXTURE_2D, v_tex, 0);

  glBindFramebuffer(GL_FRAMEBUFFER, 0);

  GLenum status = glCheckFramebufferStatus(GL_FRAMEBUFFER);
  if(status != GL_FRAMEBUFFER_COMPLETE) {       
    RX_ERROR("Framebuffer is not yet complete");
    return false;
  }

  return true;
}

bool YUV420PGrabber::setupShaders() {
  assert(window_w && window_h);
  assert(video_w && video_h);

  // y - shader
  y_vert = glCreateShader(GL_VERTEX_SHADER);
  glShaderSource(y_vert, 1, &YUV420P_Y_VS, NULL);
  glCompileShader(y_vert); eglGetShaderInfoLog(y_vert);

  y_frag = glCreateShader(GL_FRAGMENT_SHADER);
  glShaderSource(y_frag, 1, &YUV420P_Y_FS, NULL);
  glCompileShader(y_frag); eglGetShaderInfoLog(y_frag);

  y_prog = glCreateProgram();
  glAttachShader(y_prog, y_vert);
  glAttachShader(y_prog, y_frag);
  glBindAttribLocation(y_prog, 0, "a_pos");
  glBindAttribLocation(y_prog, 1, "a_tex");
  glLinkProgram(y_prog); eglGetShaderLinkLog(y_prog);
  
  glUseProgram(y_prog);

  GLint u_tex = glGetUniformLocation(y_prog, "u_tex");
  if(u_tex < 0) {
    RX_ERROR("Error while trying to get u_tex.");
    return false;
  }
  glUniform1i(u_tex, 0);

  // uv - shader
  uv_frag = glCreateShader(GL_FRAGMENT_SHADER);
  glShaderSource(uv_frag, 1, &YUV420P_UV_FS, NULL);
  glCompileShader(uv_frag); eglGetShaderInfoLog(uv_frag);

  uv_prog = glCreateProgram();
  glAttachShader(uv_prog, y_vert);
  glAttachShader(uv_prog, uv_frag);
  glBindAttribLocation(uv_prog, 0, "a_pos");
  glBindAttribLocation(uv_prog, 1, "a_tex");
  glBindFragDataLocation(uv_prog, 0, "u_col");
  glBindFragDataLocation(uv_prog, 1, "v_col");
  glLinkProgram(uv_prog); eglGetShaderLinkLog(uv_prog);

  glUseProgram(uv_prog);

  u_tex = glGetUniformLocation(uv_prog, "u_tex");
  if(u_tex < 0) {
    RX_ERROR("Cannot get u_tex for uv program");
    return false;
  }
  glUniform1i(u_tex, 0);

  // pass through shader
  pt_frag = glCreateShader(GL_FRAGMENT_SHADER);
  glShaderSource(pt_frag, 1, &YUV420_PASSTHROUGH_FS, NULL);
  glCompileShader(pt_frag); eglGetShaderInfoLog(pt_frag);
  
  pt_prog = glCreateProgram();
  glAttachShader(pt_prog, y_vert);
  glAttachShader(pt_prog, pt_frag);
  glBindAttribLocation(pt_prog, 0, "a_pos");
  glBindAttribLocation(pt_prog, 1, "a_tex");
  glLinkProgram(pt_prog); eglGetShaderLinkLog(pt_prog);
  glUseProgram(pt_prog);

  u_tex = glGetUniformLocation(pt_prog, "u_tex");
  if(u_tex < 0) {
    RX_ERROR("Cannot get u_tex for pass though program");
    return false;
  }
  
  glUniform1i(u_tex, 0);

  return true;
}

bool YUV420PGrabber::setupBuffers() {
  assert(window_w && window_h);
  assert(video_w && video_h);
  assert(y_prog && y_vert && y_frag);

  GLfloat vertices[] = {
    -1.0, -1.0, 0.0, 0.0, // a
    1.0, -1.0, 1.0, 0.0,  // b 
    1.0, 1.0, 1.0, 1.0,  // c
    -1.0, -1.0, 0.0, 0.0,// a
    1.0, 1.0, 1.0, 1.0,  // c
    -1.0, 1.0, 0.0, 1.0  // d
  };

  glGenBuffers(1, &vbo);
  glBindBuffer(GL_ARRAY_BUFFER, vbo);
  glBufferData(GL_ARRAY_BUFFER, sizeof(GLfloat) * 6 * 4, vertices, GL_STATIC_DRAW);

  glGenVertexArrays(1, &vao);
  glBindVertexArray(vao);

  glBindBuffer(GL_ARRAY_BUFFER, vbo);
  glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, sizeof(GLfloat) * 4, (GLvoid*)NULL);
  glEnableVertexAttribArray(0);

  glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, sizeof(GLfloat) * 4, (GLvoid*)8);
  glEnableVertexAttribArray(1);

  return true;
}

#if YUV420_MODE == YUV420_MODE_PBO
bool YUV420PGrabber::setupPBOs() {
  createPBOs(y_pbos, video_w, video_h);
  createPBOs(u_pbos, uv_w, uv_h);
  createPBOs(v_pbos, uv_w, uv_h);

  return true;
}

// we assume dest has 2 elements
bool YUV420PGrabber::createPBOs(GLuint* dest, int w, int h) {
  glGenBuffers(2, dest);
  for(int i = 0; i < 2; ++i) {
    glBindBuffer(GL_PIXEL_PACK_BUFFER, dest[i]);
    glBufferData(GL_PIXEL_PACK_BUFFER, w * h, NULL, GL_STREAM_COPY);
  }
  return true;
}
#endif // YUV420_MODE_PBO


void YUV420PGrabber::start() {
  assert(millis_per_frame);

  if(state == YUV420_STATE_STARTED) {
    RX_ERROR("Already started the yuv420 grabber.");
    return;
  }

  state = YUV420_STATE_STARTED;
  frame_timeout = (uv_hrtime() / 1000000) + millis_per_frame;
}

void YUV420PGrabber::stop(){
  assert(millis_per_frame);
  if(state == YUV420_STATE_NONE) {
    RX_ERROR("Not started");
    return;
  }
  state = YUV420_STATE_NONE;
}

void YUV420PGrabber::beginGrab() {
  assert(scene_fbo);
  assert(window_w);
  assert(window_h);

  GLenum bufs[] = { GL_COLOR_ATTACHMENT0 };
  glBindFramebuffer(GL_DRAW_FRAMEBUFFER, scene_fbo);
  glDrawBuffers(1, bufs);
  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
  glViewport(0, 0, window_w, window_h);
}

void YUV420PGrabber::endGrab() {
  assert(scene_fbo);
  assert(cb_frame);

  // render Y
  GLenum draw_buf_y[] = { GL_COLOR_ATTACHMENT1 };
  glDrawBuffers(1, draw_buf_y);
  glClear(GL_COLOR_BUFFER_BIT);
  glViewport(0,0,video_w, video_h);

  glActiveTexture(GL_TEXTURE0);
  glBindTexture(GL_TEXTURE_2D, scene_tex);

  glBindVertexArray(vao);
  glUseProgram(y_prog);
  glDrawArrays(GL_TRIANGLES, 0, 6);

  // render U + V
  GLenum draw_buf_uv[] = { GL_COLOR_ATTACHMENT2, GL_COLOR_ATTACHMENT3 };
  glDrawBuffers(2, draw_buf_uv);
  glClear(GL_COLOR_BUFFER_BIT);
  glViewport(0,0,uv_w, uv_h);
  
  glUseProgram(uv_prog);
  glDrawArrays(GL_TRIANGLES, 0, 6);

  glBindTexture(GL_TEXTURE_2D, 0);
  glBindFramebuffer(GL_FRAMEBUFFER, 0);
  glBindFramebuffer(GL_READ_FRAMEBUFFER, scene_fbo);

  if(state == YUV420_STATE_STARTED) {
    uint64_t now = uv_hrtime() / 1000000;
    if(now >= frame_timeout) {
      frame_timeout = now + millis_per_frame;

#if YUV420_MODE == YUV420_MODE_PBO
#elif YUV420_MODE == YUV420_MODE_GLREAD

      // important note: on Mac 10.8.4 the order of glReadPixels() is important 
      // the last one, from GL_COLOR_ATTACHMENT1, must be read last.. somehow
      // reading this one first will result in a buffer which only contains 1/4th
      // of the fbo.

      glReadBuffer(GL_COLOR_ATTACHMENT3);
      glReadPixels(0,0,uv_w, uv_h, GL_RED, GL_UNSIGNED_BYTE, v_buf);
      
      glReadBuffer(GL_COLOR_ATTACHMENT2);
      glReadPixels(0,0,uv_w, uv_h, GL_RED, GL_UNSIGNED_BYTE, u_buf);

      glReadBuffer(GL_COLOR_ATTACHMENT1);
      glReadPixels(0,0,video_w, video_h, GL_RED, GL_UNSIGNED_BYTE, y_buf);

      cb_frame(image, cb_user);


#endif  
    
    }
  } // if(state == YUV420_STATE_STARTED)

  glBindFramebuffer(GL_FRAMEBUFFER, 0);
  glDrawBuffer(GL_BACK_LEFT);
  glViewport(0,0,window_w, window_h);
}

void YUV420PGrabber::draw() {
  glBindVertexArray(vao);
  glUseProgram(pt_prog);
  glActiveTexture(GL_TEXTURE0);
  glBindTexture(GL_TEXTURE_2D, scene_tex);
  glDrawArrays(GL_TRIANGLES, 0, 6);
}

YUV420PGrabber.h

#ifndef ROXLU_YUV420P_GRABBER_H
#define ROXLU_YUV420P_GRABBER_H

extern "C" {
#  include <uv.h>
}

#include <roxlu/opengl/GL.h>

#define YUV420_MODE_GLREAD 1             /* use simple glRead() do retrieve the y,u,v data */
#define YUV420_MODE_PBO 2                /* use PBOs to retrieve/download the yuv data from gpu */
#define YUV420_MODE YUV420_MODE_GLREAD

#define YUV420_STATE_NONE 0
#define YUV420_STATE_STARTED 1

typedef void(*yuv420p_on_frame)(unsigned char* image, void* user);  /* gets called when we read a new frame. image is an 3 element array, with pointers to the 3 planes */

static const char* YUV420P_Y_VS = ""
  "#version 150\n"
  "in vec4 a_pos;"
  "in vec2 a_tex;"
  "out vec2 v_tex;"
  "void main() {"
  "   gl_Position = a_pos; "
  "   v_tex = a_tex;"
  "}";

static const char* YUV420P_Y_FS = ""
  "#version 150\n"
  "uniform sampler2D u_tex;"
  "in vec2 v_tex; "
  "out float fragcol;"
 
  "float Y(vec3 c)  {"
  "  float result = (0.257 * c.r) + (0.504 * c.g) + (0.098 * c.b) + 0.0625;"
  "  return result;"
  "}"

  "void main() {"
  "  fragcol = Y(texture(u_tex, v_tex).rgb);"
  "}";

static const char* YUV420P_UV_FS = ""
  "#version 150\n"
  "uniform sampler2D u_tex;"
  "in vec2 v_tex;"
  "out float u_col;"
  "out float v_col;"

  "float V(vec3 c) {"
  "  float result = (0.439 * c.r) - (0.368 * c.g) - (0.071 * c.b) + 0.5; "
  "  return result;" 
  "}"

  "float U(vec3 c) {"
  "  float result = -(0.148 * c.r) - (0.291 * c.g) + (0.439 * c.b) + 0.5; "
  "  return result; "
  "}"

  "void main() {"
  "  vec3 col = texture(u_tex, v_tex).rgb;"
  "  u_col = U(col); "
  "  v_col = V(col); "
  "}";

static const char* YUV420_PASSTHROUGH_FS = ""
  "#version 150\n"
  "uniform sampler2D u_tex;"
  "in vec2 v_tex; "
  "out vec4 fragcolor;"
  
  "void main() {"
  "  fragcolor.rgb = texture(u_tex, v_tex).rgb;"
  "  fragcolor.a = 1.0;"
  "}";

class YUV420PGrabber {
 public:
  YUV420PGrabber();
  ~YUV420PGrabber();
  bool setup(int windowW, int windowH, int videoW, int videoH, int fps, yuv420p_on_frame frameCB, void* frameUser);
  void beginGrab();
  void endGrab();
  void draw();  
  void start();                /* start calling the callback function at millis_per_frame rate  */
  void stop();                 /* stop calling the callback function */

 private:
  bool setupFBO();
  bool setupShaders();
  bool setupBuffers();

#if YUV420_MODE == YUV420_MODE_PBO
  bool setupPBOs();
  bool createBPOs(GLuint* dest, int w, int h);
#endif

 public:
  int window_w;
  int window_h;
  int video_w;                 /* the resulting video width, the y plane has the same width */
  int video_h;                 /* the resulting video height, the y plane has the same height */
  int uv_w;                    /* the width of the u and v planes */
  int uv_h;                    /* the height of the u and v planes */

  GLuint y_prog;
  GLuint y_vert;
  GLuint y_frag;
  GLuint uv_prog;
  GLuint uv_frag;              /* fragment shader for uv pass, we reuse y_vert */
  GLuint pt_prog;              /* pass through program */
  GLuint pt_frag;              /* pass through fragment shader, used in draw() */

  GLuint scene_fbo;
  GLuint scene_depth;
  GLuint scene_tex;

  GLuint y_tex;
  GLuint u_tex;
  GLuint v_tex;



  GLuint vbo;
  GLuint vao;

  unsigned char* image;     /* pointer to the 3 YUV420P planes */
  yuv420p_on_frame cb_frame;
  void* cb_user; 
  uint64_t millis_per_frame;
  uint64_t frame_timeout;
  int state;

#if YUV420_MODE == YUV420_MODE_PBO
  GLuint y_pbos[2];
  GLuint u_pbos[2];
  GLuint v_pbos[2];
  int read_dx;
  int write_dx;
#elif YUV420_MODE == YUV420_MODE_GLREAD
  unsigned char* y_buf;
  unsigned char* u_buf;
  unsigned char* v_buf;

#endif

};


#endif