enc_recon_frame_test.c 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393
  1. /*
  2. * copyright (c) 2022 Anton Khirnov <anton@khirnov.net>
  3. *
  4. * This file is part of FFmpeg.
  5. *
  6. * FFmpeg is free software; you can redistribute it and/or
  7. * modify it under the terms of the GNU Lesser General Public
  8. * License as published by the Free Software Foundation; either
  9. * version 2.1 of the License, or (at your option) any later version.
  10. *
  11. * FFmpeg is distributed in the hope that it will be useful,
  12. * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13. * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  14. * Lesser General Public License for more details.
  15. *
  16. * You should have received a copy of the GNU Lesser General Public
  17. * License along with FFmpeg; if not, write to the Free Software
  18. * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  19. */
  20. /* A test for AV_CODEC_FLAG_RECON_FRAME
  21. * TODO: dump reconstructed frames to disk */
  22. #include <stdio.h>
  23. #include <stdint.h>
  24. #include <stdlib.h>
  25. #include "decode_simple.h"
  26. #include "libavutil/adler32.h"
  27. #include "libavutil/common.h"
  28. #include "libavutil/error.h"
  29. #include "libavutil/frame.h"
  30. #include "libavutil/imgutils.h"
  31. #include "libavutil/opt.h"
  32. #include "libavformat/avformat.h"
  33. #include "libavcodec/avcodec.h"
  34. #include "libavcodec/codec.h"
  35. #include "libswscale/swscale.h"
  36. typedef struct FrameChecksum {
  37. int64_t ts;
  38. uint32_t checksum[4];
  39. } FrameChecksum;
  40. typedef struct PrivData {
  41. AVCodecContext *enc;
  42. AVCodecContext *dec;
  43. int64_t pts_in;
  44. AVPacket *pkt;
  45. AVFrame *frame, *frame_recon;
  46. struct SwsContext *scaler;
  47. FrameChecksum *checksums_decoded;
  48. size_t nb_checksums_decoded;
  49. FrameChecksum *checksums_recon;
  50. size_t nb_checksums_recon;
  51. } PrivData;
  52. static int frame_hash(FrameChecksum **pc, size_t *nb_c, int64_t ts,
  53. const AVFrame *frame)
  54. {
  55. FrameChecksum *c;
  56. int shift_h[4] = { 0 }, shift_v[4] = { 0 };
  57. c = av_realloc_array(*pc, *nb_c + 1, sizeof(*c));
  58. if (!c)
  59. return AVERROR(ENOMEM);
  60. *pc = c;
  61. (*nb_c)++;
  62. c += *nb_c - 1;
  63. memset(c, 0, sizeof(*c));
  64. av_pix_fmt_get_chroma_sub_sample(frame->format, &shift_h[1], &shift_v[1]);
  65. shift_h[2] = shift_h[1];
  66. shift_v[2] = shift_v[1];
  67. c->ts = ts;
  68. for (int p = 0; frame->data[p]; p++) {
  69. const uint8_t *data = frame->data[p];
  70. int linesize = av_image_get_linesize(frame->format, frame->width, p);
  71. uint32_t checksum = 0;
  72. for (int j = 0; j < frame->height >> shift_v[p]; j++) {
  73. checksum = av_adler32_update(checksum, data, linesize);
  74. data += frame->linesize[p];
  75. }
  76. c->checksum[p] = checksum;
  77. }
  78. return 0;
  79. }
  80. static int recon_frame_process(PrivData *pd, const AVPacket *pkt)
  81. {
  82. AVFrame *f = pd->frame_recon;
  83. int ret;
  84. ret = avcodec_receive_frame(pd->enc, f);
  85. if (ret < 0) {
  86. fprintf(stderr, "Error retrieving a reconstructed frame\n");
  87. return ret;
  88. }
  89. // the encoder's internal format (in which the reconsturcted frames are
  90. // exported) may be different from the user-facing pixel format
  91. if (f->format != pd->enc->pix_fmt) {
  92. if (!pd->scaler) {
  93. pd->scaler = sws_getContext(f->width, f->height, f->format,
  94. f->width, f->height, pd->enc->pix_fmt,
  95. SWS_BITEXACT, NULL, NULL, NULL);
  96. if (!pd->scaler)
  97. return AVERROR(ENOMEM);
  98. }
  99. ret = sws_scale_frame(pd->scaler, pd->frame, f);
  100. if (ret < 0) {
  101. fprintf(stderr, "Error converting pixel formats\n");
  102. return ret;
  103. }
  104. av_frame_unref(f);
  105. f = pd->frame;
  106. }
  107. ret = frame_hash(&pd->checksums_recon, &pd->nb_checksums_recon,
  108. pkt->pts, f);
  109. av_frame_unref(f);
  110. return 0;
  111. }
  112. static int process_frame(DecodeContext *dc, AVFrame *frame)
  113. {
  114. PrivData *pd = dc->opaque;
  115. int ret;
  116. if (!avcodec_is_open(pd->enc)) {
  117. if (!frame) {
  118. fprintf(stderr, "No input frames were decoded\n");
  119. return AVERROR_INVALIDDATA;
  120. }
  121. pd->enc->width = frame->width;
  122. pd->enc->height = frame->height;
  123. pd->enc->pix_fmt = frame->format;
  124. pd->enc->thread_count = dc->decoder->thread_count;
  125. pd->enc->thread_type = dc->decoder->thread_type;
  126. // real timestamps do not matter for this test, so we just
  127. // pretend the input is 25fps CFR to avoid any timestamp issues
  128. pd->enc->time_base = (AVRational){ 1, 25 };
  129. ret = avcodec_open2(pd->enc, NULL, NULL);
  130. if (ret < 0) {
  131. fprintf(stderr, "Error opening the encoder\n");
  132. return ret;
  133. }
  134. }
  135. if (frame) {
  136. frame->pts = pd->pts_in++;
  137. // avoid forcing coded frame type
  138. frame->pict_type = AV_PICTURE_TYPE_NONE;
  139. }
  140. ret = avcodec_send_frame(pd->enc, frame);
  141. if (ret < 0) {
  142. fprintf(stderr, "Error submitting a frame for encoding\n");
  143. return ret;
  144. }
  145. while (1) {
  146. AVPacket *pkt = pd->pkt;
  147. ret = avcodec_receive_packet(pd->enc, pkt);
  148. if (ret == AVERROR(EAGAIN))
  149. break;
  150. else if (ret == AVERROR_EOF)
  151. pkt = NULL;
  152. else if (ret < 0) {
  153. fprintf(stderr, "Error receiving a frame from the encoder\n");
  154. return ret;
  155. }
  156. if (pkt) {
  157. ret = recon_frame_process(pd, pkt);
  158. if (ret < 0)
  159. return ret;
  160. }
  161. if (!avcodec_is_open(pd->dec)) {
  162. if (!pkt) {
  163. fprintf(stderr, "No packets were received from the encoder\n");
  164. return AVERROR(EINVAL);
  165. }
  166. pd->dec->width = pd->enc->width;
  167. pd->dec->height = pd->enc->height;
  168. pd->dec->pix_fmt = pd->enc->pix_fmt;
  169. pd->dec->thread_count = dc->decoder->thread_count;
  170. pd->dec->thread_type = dc->decoder->thread_type;
  171. if (pd->enc->extradata_size) {
  172. pd->dec->extradata = av_memdup(pd->enc->extradata,
  173. pd->enc->extradata_size + AV_INPUT_BUFFER_PADDING_SIZE);
  174. if (!pd->dec->extradata)
  175. return AVERROR(ENOMEM);
  176. }
  177. ret = avcodec_open2(pd->dec, NULL, NULL);
  178. if (ret < 0) {
  179. fprintf(stderr, "Error opening the decoder\n");
  180. return ret;
  181. }
  182. }
  183. ret = avcodec_send_packet(pd->dec, pkt);
  184. if (ret < 0) {
  185. fprintf(stderr, "Error sending a packet to decoder\n");
  186. return ret;
  187. }
  188. while (1) {
  189. ret = avcodec_receive_frame(pd->dec, pd->frame);
  190. if (ret == AVERROR(EAGAIN))
  191. break;
  192. else if (ret == AVERROR_EOF)
  193. return 0;
  194. else if (ret < 0) {
  195. fprintf(stderr, "Error receving a frame from decoder\n");
  196. return ret;
  197. }
  198. ret = frame_hash(&pd->checksums_decoded, &pd->nb_checksums_decoded,
  199. pd->frame->pts, pd->frame);
  200. av_frame_unref(pd->frame);
  201. if (ret < 0)
  202. return ret;
  203. }
  204. }
  205. return 0;
  206. }
  207. static int frame_checksum_compare(const void *a, const void *b)
  208. {
  209. const FrameChecksum *ca = a;
  210. const FrameChecksum *cb = b;
  211. if (ca->ts == cb->ts)
  212. return 0;
  213. return FFSIGN(ca->ts - cb->ts);
  214. }
  215. int main(int argc, char **argv)
  216. {
  217. PrivData pd;
  218. DecodeContext dc;
  219. const char *filename, *enc_name, *enc_opts, *thread_type = NULL, *nb_threads = NULL;
  220. const AVCodec *enc, *dec;
  221. int ret = 0, max_frames = 0;
  222. if (argc < 4) {
  223. fprintf(stderr,
  224. "Usage: %s <input file> <encoder> <encoder options> "
  225. "[<max frame count> [<thread count> <thread type>]\n",
  226. argv[0]);
  227. return 0;
  228. }
  229. filename = argv[1];
  230. enc_name = argv[2];
  231. enc_opts = argv[3];
  232. if (argc >= 5)
  233. max_frames = strtol(argv[4], NULL, 0);
  234. if (argc >= 6)
  235. nb_threads = argv[5];
  236. if (argc >= 7)
  237. thread_type = argv[6];
  238. memset(&dc, 0, sizeof(dc));
  239. memset(&pd, 0, sizeof(pd));
  240. enc = avcodec_find_encoder_by_name(enc_name);
  241. if (!enc) {
  242. fprintf(stderr, "No such encoder: %s\n", enc_name);
  243. return 1;
  244. }
  245. if (!(enc->capabilities & AV_CODEC_CAP_ENCODER_RECON_FRAME)) {
  246. fprintf(stderr, "Encoder '%s' cannot ouput reconstructed frames\n",
  247. enc->name);
  248. return 1;
  249. }
  250. dec = avcodec_find_decoder(enc->id);
  251. if (!dec) {
  252. fprintf(stderr, "No decoder for: %s\n", avcodec_get_name(enc->id));
  253. return 1;
  254. }
  255. pd.enc = avcodec_alloc_context3(enc);
  256. if (!pd.enc) {
  257. fprintf(stderr, "Error allocating encoder\n");
  258. return 1;
  259. }
  260. ret = av_set_options_string(pd.enc, enc_opts, "=", ",");
  261. if (ret < 0) {
  262. fprintf(stderr, "Error setting encoder options\n");
  263. goto fail;
  264. }
  265. pd.enc->flags |= AV_CODEC_FLAG_RECON_FRAME | AV_CODEC_FLAG_BITEXACT;
  266. pd.dec = avcodec_alloc_context3(dec);
  267. if (!pd.dec) {
  268. fprintf(stderr, "Error allocating decoder\n");
  269. goto fail;
  270. }
  271. pd.dec->flags |= AV_CODEC_FLAG_BITEXACT;
  272. pd.dec->err_recognition |= AV_EF_CRCCHECK;
  273. pd.frame = av_frame_alloc();
  274. pd.frame_recon = av_frame_alloc();
  275. pd.pkt = av_packet_alloc();
  276. if (!pd.frame ||!pd.frame_recon || !pd.pkt) {
  277. ret = 1;
  278. goto fail;
  279. }
  280. ret = ds_open(&dc, filename, 0);
  281. if (ret < 0) {
  282. fprintf(stderr, "Error opening the file\n");
  283. goto fail;
  284. }
  285. dc.process_frame = process_frame;
  286. dc.opaque = &pd;
  287. dc.max_frames = max_frames;
  288. ret = av_dict_set(&dc.decoder_opts, "threads", nb_threads, 0);
  289. ret |= av_dict_set(&dc.decoder_opts, "thread_type", thread_type, 0);
  290. ret = ds_run(&dc);
  291. if (ret < 0)
  292. goto fail;
  293. if (pd.nb_checksums_decoded != pd.nb_checksums_recon) {
  294. fprintf(stderr, "Mismatching frame counts: recon=%zu decoded=%zu\n",
  295. pd.nb_checksums_recon, pd.nb_checksums_decoded);
  296. ret = 1;
  297. goto fail;
  298. }
  299. // reconstructed frames are in coded order, sort them by pts into presentation order
  300. qsort(pd.checksums_recon, pd.nb_checksums_recon, sizeof(*pd.checksums_recon),
  301. frame_checksum_compare);
  302. for (size_t i = 0; i < pd.nb_checksums_decoded; i++) {
  303. const FrameChecksum *d = &pd.checksums_decoded[i];
  304. const FrameChecksum *r = &pd.checksums_recon[i];
  305. for (int p = 0; p < FF_ARRAY_ELEMS(d->checksum); p++)
  306. if (d->checksum[p] != r->checksum[p]) {
  307. fprintf(stderr, "Checksum mismatch in frame ts=%"PRId64", plane %d\n",
  308. d->ts, p);
  309. ret = 1;
  310. goto fail;
  311. }
  312. }
  313. fprintf(stderr, "All %zu encoded frames match\n", pd.nb_checksums_decoded);
  314. fail:
  315. avcodec_free_context(&pd.enc);
  316. avcodec_free_context(&pd.dec);
  317. av_freep(&pd.checksums_decoded);
  318. av_freep(&pd.checksums_recon);
  319. av_frame_free(&pd.frame);
  320. av_frame_free(&pd.frame_recon);
  321. av_packet_free(&pd.pkt);
  322. ds_free(&dc);
  323. return !!ret;
  324. }