recorder.cpp 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166
  1. #include "recorder.h"
  2. #if defined ENABLE_THEORA_RECORDER
  3. #include <random>
  4. #include "simple/file.hpp"
  5. #include "simple/geom/algorithm.hpp"
  6. using simple::graphical::color_vector;
  7. using simple::graphical::rgb_vector;
  8. using simple::graphical::rgb_pixel;
  9. using simple::geom::vector;
  10. using simple::geom::gauss_jordan_elimination;
  11. // color spec from https://theora.org/doc/Theora.pdf
  12. constexpr auto offset = vector(16.f,128.f,128.f);
  13. constexpr auto excursion = vector(219.f,224.f,224.f);
  14. constexpr float Kr = 0.299f;
  15. constexpr float Kb = 0.114f;
  16. constexpr float iKr = 1.f - Kr;
  17. constexpr float iKb = 1.f - Kb;
  18. constexpr float iK = 1.f - Kb - Kr;
  19. // cue 3blue1brown explaining this with projections of
  20. // visualizations in 6 dimensions...
  21. constexpr auto transform_rgb_yuv = gauss_jordan_elimination(
  22. vector(
  23. vector(1.f, 0.f, 2*iKr, 1.f, 0.f, 0.f),
  24. vector(1.f, -2*iKb*Kb/iK, -2*iKr*Kr/iK, 0.f, 1.f, 0.f),
  25. vector(1.f, 2*iKb, 0.f, 0.f, 0.f, 1.f)
  26. )
  27. ).transformed(&vector<float,6>::last<3>);
  28. unsigned ceil_16(unsigned n)
  29. {
  30. return ((n + 0b1111) >> 4 ) << 4;
  31. }
  32. void write(std::ostream& os, const ogg_page& page)
  33. {
  34. os.write(reinterpret_cast<const char*>(page.header), page.header_len);
  35. os.write(reinterpret_cast<const char*>(page.body), page.body_len);
  36. }
  37. recorder::recorder(std::string filename, int2 size) :
  38. file(simple::file::bwopex(filename)),
  39. ycbcr_buffer(size.x()*size.y()*3/2)
  40. {
  41. ogg_stream_init(&video_stream, std::random_device{}());
  42. th_info_init(&video_info);
  43. video_info.frame_width = ceil_16(size.x());
  44. video_info.frame_height = ceil_16(size.y());
  45. video_info.pic_width = size.x();
  46. video_info.pic_height = size.y();
  47. video_info.pic_x = 0;
  48. video_info.pic_y = 0;
  49. video_info.fps_numerator = 60;
  50. video_info.fps_denominator = 1;
  51. video_info.colorspace = TH_CS_ITU_REC_470M;
  52. video_info.pixel_fmt = TH_PF_420;
  53. video_info.quality = 32;
  54. video_info.keyframe_granule_shift = 6;
  55. video_encoder = th_encode_alloc(&video_info);
  56. th_comment video_comments;
  57. th_comment_init(&video_comments);
  58. while(th_encode_flushheader(video_encoder, &video_comments, &packet) > 0)
  59. {
  60. ogg_stream_packetin(&video_stream, &packet);
  61. if(ogg_stream_pageout(&video_stream, &page)) // or while?
  62. write(file, page);
  63. }
  64. th_comment_clear(&video_comments);
  65. // apparently need to flush here
  66. if(ogg_stream_flush(&video_stream, &page))
  67. write(file, page);
  68. ycbcr_view[0].width = video_info.frame_width;
  69. ycbcr_view[0].height = video_info.frame_height;
  70. ycbcr_view[0].stride = video_info.frame_width;
  71. ycbcr_view[0].data = ycbcr_buffer.data();
  72. for(int i = 1; i < 3; ++i)
  73. {
  74. ycbcr_view[i].width = video_info.frame_width/2;
  75. ycbcr_view[i].height = video_info.frame_height/2;
  76. ycbcr_view[i].stride = video_info.frame_width/2;
  77. ycbcr_view[i].data = ycbcr_buffer.data()
  78. + video_info.frame_width*video_info.frame_height
  79. + video_info.frame_width*video_info.frame_height/4 * (i-1);
  80. }
  81. ogg_uint32_t freq = 64;
  82. th_encode_ctl(video_encoder,
  83. TH_ENCCTL_SET_KEYFRAME_FREQUENCY_FORCE, &freq, sizeof(freq));
  84. }
  85. void recorder::record(const simple::graphical::surface& frame, int cell_size, bool is_last)
  86. {
  87. using namespace simple::graphical;
  88. auto palette = *frame.format().palette();
  89. auto pixels = std::get<pixel_writer<pixel_byte>>(frame.pixels());
  90. simple::geom::loop(int2::zero(), pixels.size()/2, int2::one(),
  91. [&](auto i)
  92. {
  93. float2 uv{};
  94. auto flat_i = i.x()+i.y()*video_info.frame_width/2;
  95. simple::geom::loop(int2::zero(), int2::one(2), int2::one(),
  96. [&](auto j)
  97. {
  98. auto sample = 2*i+j;
  99. rgb_vector pixel(palette[pixels[sample/cell_size]]);
  100. auto yuv = round(pixel(transform_rgb_yuv) * excursion + offset);
  101. auto flat_sample = sample.x()+sample.y()*video_info.frame_width;
  102. ycbcr_view[0].data[flat_sample] = yuv[0];
  103. uv += yuv.last<2>();
  104. });
  105. uv /= 4;
  106. uv.floor();
  107. ycbcr_view[1].data[flat_i] = uv[0];
  108. ycbcr_view[2].data[flat_i] = uv[1];
  109. });
  110. th_encode_ycbcr_in(video_encoder, ycbcr_view);
  111. while(th_encode_packetout(video_encoder, is_last, &packet))
  112. {
  113. ogg_stream_packetin(&video_stream, &packet);
  114. while(ogg_stream_pageout(&video_stream, &page))
  115. write(file, page);
  116. }
  117. if(is_last)
  118. if(ogg_stream_flush(&video_stream, &page))
  119. write(file, page);
  120. }
  121. recorder::~recorder()
  122. {
  123. th_info_clear(&video_info);
  124. th_encode_free(video_encoder);
  125. ogg_stream_clear(&video_stream);
  126. }
  127. #else
  128. recorder::recorder(std::string, int2)
  129. {
  130. }
  131. void recorder::record(const simple::graphical::surface&, int, bool)
  132. {
  133. }
  134. recorder::~recorder()
  135. {
  136. }
  137. #endif