ffmpeg-mux.c 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856
  1. /*
  2. * Copyright (c) 2015 Hugh Bailey <[email protected]>
  3. *
  4. * Permission to use, copy, modify, and distribute this software for any
  5. * purpose with or without fee is hereby granted, provided that the above
  6. * copyright notice and this permission notice appear in all copies.
  7. *
  8. * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
  9. * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
  10. * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
  11. * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
  12. * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
  13. * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
  14. * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
  15. */
  16. #ifdef _WIN32
  17. #include <io.h>
  18. #include <fcntl.h>
  19. #include <windows.h>
  20. #define inline __inline
  21. #endif
  22. #include <stdio.h>
  23. #include <stdlib.h>
  24. #include "ffmpeg-mux.h"
  25. #include <util/dstr.h>
  26. #include <libavformat/avformat.h>
  27. #define ANSI_COLOR_RED "\x1b[0;91m"
  28. #define ANSI_COLOR_MAGENTA "\x1b[0;95m"
  29. #define ANSI_COLOR_RESET "\x1b[0m"
  30. #if LIBAVCODEC_VERSION_MAJOR >= 58
  31. #define CODEC_FLAG_GLOBAL_H AV_CODEC_FLAG_GLOBAL_HEADER
  32. #else
  33. #define CODEC_FLAG_GLOBAL_H CODEC_FLAG_GLOBAL_HEADER
  34. #endif
  35. /* ------------------------------------------------------------------------- */
  36. static char *global_stream_key = "";
  37. struct resize_buf {
  38. uint8_t *buf;
  39. size_t size;
  40. size_t capacity;
  41. };
  42. static inline void resize_buf_resize(struct resize_buf *rb, size_t size)
  43. {
  44. if (!rb->buf) {
  45. rb->buf = malloc(size);
  46. rb->size = size;
  47. rb->capacity = size;
  48. } else {
  49. if (rb->capacity < size) {
  50. size_t capx2 = rb->capacity * 2;
  51. size_t new_cap = capx2 > size ? capx2 : size;
  52. rb->buf = realloc(rb->buf, new_cap);
  53. rb->capacity = new_cap;
  54. }
  55. rb->size = size;
  56. }
  57. }
  58. static inline void resize_buf_free(struct resize_buf *rb)
  59. {
  60. free(rb->buf);
  61. }
  62. /* ------------------------------------------------------------------------- */
  63. struct main_params {
  64. char *file;
  65. /* printable_file is file with any stream key information removed */
  66. struct dstr printable_file;
  67. int has_video;
  68. int tracks;
  69. char *vcodec;
  70. int vbitrate;
  71. int gop;
  72. int width;
  73. int height;
  74. int fps_num;
  75. int fps_den;
  76. int color_primaries;
  77. int color_trc;
  78. int colorspace;
  79. int color_range;
  80. char *acodec;
  81. char *muxer_settings;
  82. };
  83. struct audio_params {
  84. char *name;
  85. int abitrate;
  86. int sample_rate;
  87. int channels;
  88. };
  89. struct header {
  90. uint8_t *data;
  91. int size;
  92. };
  93. struct audio_info {
  94. AVStream *stream;
  95. AVCodecContext *ctx;
  96. };
  97. struct ffmpeg_mux {
  98. AVFormatContext *output;
  99. AVStream *video_stream;
  100. AVCodecContext *video_ctx;
  101. struct audio_info *audio_infos;
  102. struct main_params params;
  103. struct audio_params *audio;
  104. struct header video_header;
  105. struct header *audio_header;
  106. int num_audio_streams;
  107. bool initialized;
  108. char error[4096];
  109. };
  110. static void header_free(struct header *header)
  111. {
  112. free(header->data);
  113. }
  114. static void free_avformat(struct ffmpeg_mux *ffm)
  115. {
  116. if (ffm->output) {
  117. #if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 48, 101)
  118. avcodec_free_context(&ffm->video_ctx);
  119. #endif
  120. if ((ffm->output->oformat->flags & AVFMT_NOFILE) == 0)
  121. avio_close(ffm->output->pb);
  122. avformat_free_context(ffm->output);
  123. ffm->output = NULL;
  124. }
  125. if (ffm->audio_infos) {
  126. #if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 48, 101)
  127. for (int i = 0; i < ffm->num_audio_streams; ++i)
  128. avcodec_free_context(&ffm->audio_infos[i].ctx);
  129. #endif
  130. free(ffm->audio_infos);
  131. }
  132. ffm->video_stream = NULL;
  133. ffm->audio_infos = NULL;
  134. ffm->num_audio_streams = 0;
  135. }
  136. static void ffmpeg_mux_free(struct ffmpeg_mux *ffm)
  137. {
  138. if (ffm->initialized) {
  139. av_write_trailer(ffm->output);
  140. }
  141. free_avformat(ffm);
  142. header_free(&ffm->video_header);
  143. if (ffm->audio_header) {
  144. for (int i = 0; i < ffm->params.tracks; i++) {
  145. header_free(&ffm->audio_header[i]);
  146. }
  147. free(ffm->audio_header);
  148. }
  149. if (ffm->audio) {
  150. free(ffm->audio);
  151. }
  152. dstr_free(&ffm->params.printable_file);
  153. memset(ffm, 0, sizeof(*ffm));
  154. }
  155. static bool get_opt_str(int *p_argc, char ***p_argv, char **str,
  156. const char *opt)
  157. {
  158. int argc = *p_argc;
  159. char **argv = *p_argv;
  160. if (!argc) {
  161. printf("Missing expected option: '%s'\n", opt);
  162. return false;
  163. }
  164. (*p_argc)--;
  165. (*p_argv)++;
  166. *str = argv[0];
  167. return true;
  168. }
  169. static bool get_opt_int(int *p_argc, char ***p_argv, int *i, const char *opt)
  170. {
  171. char *str;
  172. if (!get_opt_str(p_argc, p_argv, &str, opt)) {
  173. return false;
  174. }
  175. *i = atoi(str);
  176. return true;
  177. }
  178. static bool get_audio_params(struct audio_params *audio, int *argc,
  179. char ***argv)
  180. {
  181. if (!get_opt_str(argc, argv, &audio->name, "audio track name"))
  182. return false;
  183. if (!get_opt_int(argc, argv, &audio->abitrate, "audio bitrate"))
  184. return false;
  185. if (!get_opt_int(argc, argv, &audio->sample_rate, "audio sample rate"))
  186. return false;
  187. if (!get_opt_int(argc, argv, &audio->channels, "audio channels"))
  188. return false;
  189. return true;
  190. }
  191. static void ffmpeg_log_callback(void *param, int level, const char *format,
  192. va_list args)
  193. {
  194. char out_buffer[4096];
  195. struct dstr out = {0};
  196. vsnprintf(out_buffer, sizeof(out_buffer), format, args);
  197. dstr_copy(&out, out_buffer);
  198. dstr_replace(&out, global_stream_key, "{stream_key}");
  199. switch (level) {
  200. case AV_LOG_INFO:
  201. fprintf(stdout, "info: [ffmpeg_muxer] %s", out.array);
  202. fflush(stdout);
  203. break;
  204. case AV_LOG_WARNING:
  205. fprintf(stdout, "%swarning: [ffmpeg_muxer] %s%s",
  206. ANSI_COLOR_MAGENTA, out.array, ANSI_COLOR_RESET);
  207. fflush(stdout);
  208. break;
  209. case AV_LOG_ERROR:
  210. fprintf(stderr, "%serror: [ffmpeg_muxer] %s%s", ANSI_COLOR_RED,
  211. out.array, ANSI_COLOR_RESET);
  212. fflush(stderr);
  213. }
  214. dstr_free(&out);
  215. UNUSED_PARAMETER(param);
  216. }
  217. static bool init_params(int *argc, char ***argv, struct main_params *params,
  218. struct audio_params **p_audio)
  219. {
  220. struct audio_params *audio = NULL;
  221. if (!get_opt_str(argc, argv, &params->file, "file name"))
  222. return false;
  223. if (!get_opt_int(argc, argv, &params->has_video, "video track count"))
  224. return false;
  225. if (!get_opt_int(argc, argv, &params->tracks, "audio track count"))
  226. return false;
  227. if (params->has_video > 1 || params->has_video < 0) {
  228. puts("Invalid number of video tracks\n");
  229. return false;
  230. }
  231. if (params->tracks < 0) {
  232. puts("Invalid number of audio tracks\n");
  233. return false;
  234. }
  235. if (params->has_video == 0 && params->tracks == 0) {
  236. puts("Must have at least 1 audio track or 1 video track\n");
  237. return false;
  238. }
  239. if (params->has_video) {
  240. if (!get_opt_str(argc, argv, &params->vcodec, "video codec"))
  241. return false;
  242. if (!get_opt_int(argc, argv, &params->vbitrate,
  243. "video bitrate"))
  244. return false;
  245. if (!get_opt_int(argc, argv, &params->width, "video width"))
  246. return false;
  247. if (!get_opt_int(argc, argv, &params->height, "video height"))
  248. return false;
  249. if (!get_opt_int(argc, argv, &params->color_primaries,
  250. "video color primaries"))
  251. return false;
  252. if (!get_opt_int(argc, argv, &params->color_trc,
  253. "video color trc"))
  254. return false;
  255. if (!get_opt_int(argc, argv, &params->colorspace,
  256. "video colorspace"))
  257. return false;
  258. if (!get_opt_int(argc, argv, &params->color_range,
  259. "video color range"))
  260. return false;
  261. if (!get_opt_int(argc, argv, &params->fps_num, "video fps num"))
  262. return false;
  263. if (!get_opt_int(argc, argv, &params->fps_den, "video fps den"))
  264. return false;
  265. }
  266. if (params->tracks) {
  267. if (!get_opt_str(argc, argv, &params->acodec, "audio codec"))
  268. return false;
  269. audio = calloc(params->tracks, sizeof(*audio));
  270. for (int i = 0; i < params->tracks; i++) {
  271. if (!get_audio_params(&audio[i], argc, argv)) {
  272. free(audio);
  273. return false;
  274. }
  275. }
  276. }
  277. *p_audio = audio;
  278. dstr_copy(&params->printable_file, params->file);
  279. get_opt_str(argc, argv, &global_stream_key, "stream key");
  280. if (strcmp(global_stream_key, "") != 0) {
  281. dstr_replace(&params->printable_file, global_stream_key,
  282. "{stream_key}");
  283. }
  284. #ifdef DEBUG_FFMPEG
  285. av_log_set_callback(ffmpeg_log_callback);
  286. #endif
  287. get_opt_str(argc, argv, &params->muxer_settings, "muxer settings");
  288. return true;
  289. }
  290. static bool new_stream(struct ffmpeg_mux *ffm, AVStream **stream,
  291. const char *name, AVCodec **codec)
  292. {
  293. const AVCodecDescriptor *desc = avcodec_descriptor_get_by_name(name);
  294. if (!desc) {
  295. fprintf(stderr, "Couldn't find encoder '%s'\n", name);
  296. return false;
  297. }
  298. *codec = avcodec_find_encoder(desc->id);
  299. if (!*codec) {
  300. fprintf(stderr, "Couldn't create encoder\n");
  301. return false;
  302. }
  303. *stream = avformat_new_stream(ffm->output, *codec);
  304. if (!*stream) {
  305. fprintf(stderr, "Couldn't create stream for encoder '%s'\n",
  306. name);
  307. return false;
  308. }
  309. (*stream)->id = ffm->output->nb_streams - 1;
  310. return true;
  311. }
  312. static void create_video_stream(struct ffmpeg_mux *ffm)
  313. {
  314. AVCodec *codec;
  315. AVCodecContext *context;
  316. void *extradata = NULL;
  317. if (!new_stream(ffm, &ffm->video_stream, ffm->params.vcodec, &codec))
  318. return;
  319. if (ffm->video_header.size) {
  320. extradata = av_memdup(ffm->video_header.data,
  321. ffm->video_header.size);
  322. }
  323. #if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 48, 101)
  324. context = avcodec_alloc_context3(codec);
  325. #else
  326. context = ffm->video_stream->codec;
  327. #endif
  328. context->bit_rate = (int64_t)ffm->params.vbitrate * 1000;
  329. context->width = ffm->params.width;
  330. context->height = ffm->params.height;
  331. context->coded_width = ffm->params.width;
  332. context->coded_height = ffm->params.height;
  333. context->color_primaries = ffm->params.color_primaries;
  334. context->color_trc = ffm->params.color_trc;
  335. context->colorspace = ffm->params.colorspace;
  336. context->color_range = ffm->params.color_range;
  337. context->extradata = extradata;
  338. context->extradata_size = ffm->video_header.size;
  339. context->time_base =
  340. (AVRational){ffm->params.fps_den, ffm->params.fps_num};
  341. ffm->video_stream->time_base = context->time_base;
  342. #if LIBAVFORMAT_VERSION_MAJOR < 59
  343. // codec->time_base may still be used if LIBAVFORMAT_VERSION_MAJOR < 59
  344. ffm->video_stream->codec->time_base = context->time_base;
  345. #endif
  346. ffm->video_stream->avg_frame_rate = av_inv_q(context->time_base);
  347. if (ffm->output->oformat->flags & AVFMT_GLOBALHEADER)
  348. context->flags |= CODEC_FLAG_GLOBAL_H;
  349. #if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 48, 101)
  350. avcodec_parameters_from_context(ffm->video_stream->codecpar, context);
  351. #endif
  352. ffm->video_ctx = context;
  353. }
  354. static void create_audio_stream(struct ffmpeg_mux *ffm, int idx)
  355. {
  356. AVCodec *codec;
  357. AVCodecContext *context;
  358. AVStream *stream;
  359. void *extradata = NULL;
  360. if (!new_stream(ffm, &stream, ffm->params.acodec, &codec))
  361. return;
  362. av_dict_set(&stream->metadata, "title", ffm->audio[idx].name, 0);
  363. stream->time_base = (AVRational){1, ffm->audio[idx].sample_rate};
  364. if (ffm->audio_header[idx].size) {
  365. extradata = av_memdup(ffm->audio_header[idx].data,
  366. ffm->audio_header[idx].size);
  367. }
  368. #if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 48, 101)
  369. context = avcodec_alloc_context3(codec);
  370. #else
  371. context = stream->codec;
  372. #endif
  373. context->bit_rate = (int64_t)ffm->audio[idx].abitrate * 1000;
  374. context->channels = ffm->audio[idx].channels;
  375. context->sample_rate = ffm->audio[idx].sample_rate;
  376. context->sample_fmt = AV_SAMPLE_FMT_S16;
  377. context->time_base = stream->time_base;
  378. context->extradata = extradata;
  379. context->extradata_size = ffm->audio_header[idx].size;
  380. context->channel_layout =
  381. av_get_default_channel_layout(context->channels);
  382. //AVlib default channel layout for 4 channels is 4.0 ; fix for quad
  383. if (context->channels == 4)
  384. context->channel_layout = av_get_channel_layout("quad");
  385. //AVlib default channel layout for 5 channels is 5.0 ; fix for 4.1
  386. if (context->channels == 5)
  387. context->channel_layout = av_get_channel_layout("4.1");
  388. if (ffm->output->oformat->flags & AVFMT_GLOBALHEADER)
  389. context->flags |= CODEC_FLAG_GLOBAL_H;
  390. #if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 48, 101)
  391. avcodec_parameters_from_context(stream->codecpar, context);
  392. #endif
  393. ffm->audio_infos[ffm->num_audio_streams].stream = stream;
  394. ffm->audio_infos[ffm->num_audio_streams].ctx = context;
  395. ffm->num_audio_streams++;
  396. }
  397. static bool init_streams(struct ffmpeg_mux *ffm)
  398. {
  399. if (ffm->params.has_video)
  400. create_video_stream(ffm);
  401. if (ffm->params.tracks) {
  402. ffm->audio_infos =
  403. calloc(ffm->params.tracks, sizeof(*ffm->audio_infos));
  404. for (int i = 0; i < ffm->params.tracks; i++)
  405. create_audio_stream(ffm, i);
  406. }
  407. if (!ffm->video_stream && !ffm->num_audio_streams)
  408. return false;
  409. return true;
  410. }
  411. static void set_header(struct header *header, uint8_t *data, size_t size)
  412. {
  413. header->size = (int)size;
  414. header->data = malloc(size);
  415. memcpy(header->data, data, size);
  416. }
  417. static void ffmpeg_mux_header(struct ffmpeg_mux *ffm, uint8_t *data,
  418. struct ffm_packet_info *info)
  419. {
  420. if (info->type == FFM_PACKET_VIDEO) {
  421. set_header(&ffm->video_header, data, (size_t)info->size);
  422. } else {
  423. set_header(&ffm->audio_header[info->index], data,
  424. (size_t)info->size);
  425. }
  426. }
  427. static size_t safe_read(void *vdata, size_t size)
  428. {
  429. uint8_t *data = vdata;
  430. size_t total = size;
  431. while (size > 0) {
  432. size_t in_size = fread(data, 1, size, stdin);
  433. if (in_size == 0)
  434. return 0;
  435. size -= in_size;
  436. data += in_size;
  437. }
  438. return total;
  439. }
  440. static bool ffmpeg_mux_get_header(struct ffmpeg_mux *ffm)
  441. {
  442. struct ffm_packet_info info = {0};
  443. bool success = safe_read(&info, sizeof(info)) == sizeof(info);
  444. if (success) {
  445. uint8_t *data = malloc(info.size);
  446. if (safe_read(data, info.size) == info.size) {
  447. ffmpeg_mux_header(ffm, data, &info);
  448. } else {
  449. success = false;
  450. }
  451. free(data);
  452. }
  453. return success;
  454. }
  455. static inline bool ffmpeg_mux_get_extra_data(struct ffmpeg_mux *ffm)
  456. {
  457. if (ffm->params.has_video) {
  458. if (!ffmpeg_mux_get_header(ffm)) {
  459. return false;
  460. }
  461. }
  462. for (int i = 0; i < ffm->params.tracks; i++) {
  463. if (!ffmpeg_mux_get_header(ffm)) {
  464. return false;
  465. }
  466. }
  467. return true;
  468. }
  469. #ifdef _MSC_VER
  470. #pragma warning(disable : 4996)
  471. #endif
  472. static inline int open_output_file(struct ffmpeg_mux *ffm)
  473. {
  474. AVOutputFormat *format = ffm->output->oformat;
  475. int ret;
  476. if ((format->flags & AVFMT_NOFILE) == 0) {
  477. ret = avio_open(&ffm->output->pb, ffm->params.file,
  478. AVIO_FLAG_WRITE);
  479. if (ret < 0) {
  480. fprintf(stderr, "Couldn't open '%s', %s\n",
  481. ffm->params.printable_file.array,
  482. av_err2str(ret));
  483. return FFM_ERROR;
  484. }
  485. }
  486. AVDictionary *dict = NULL;
  487. if ((ret = av_dict_parse_string(&dict, ffm->params.muxer_settings, "=",
  488. " ", 0))) {
  489. fprintf(stderr, "Failed to parse muxer settings: %s\n%s\n",
  490. av_err2str(ret), ffm->params.muxer_settings);
  491. av_dict_free(&dict);
  492. }
  493. if (av_dict_count(dict) > 0) {
  494. printf("Using muxer settings:");
  495. AVDictionaryEntry *entry = NULL;
  496. while ((entry = av_dict_get(dict, "", entry,
  497. AV_DICT_IGNORE_SUFFIX)))
  498. printf("\n\t%s=%s", entry->key, entry->value);
  499. printf("\n");
  500. }
  501. ret = avformat_write_header(ffm->output, &dict);
  502. if (ret < 0) {
  503. fprintf(stderr, "Error opening '%s': %s",
  504. ffm->params.printable_file.array, av_err2str(ret));
  505. av_dict_free(&dict);
  506. return ret == -22 ? FFM_UNSUPPORTED : FFM_ERROR;
  507. }
  508. av_dict_free(&dict);
  509. return FFM_SUCCESS;
  510. }
  511. #define SRT_PROTO "srt"
  512. #define UDP_PROTO "udp"
  513. #define TCP_PROTO "tcp"
  514. static int ffmpeg_mux_init_context(struct ffmpeg_mux *ffm)
  515. {
  516. AVOutputFormat *output_format;
  517. int ret;
  518. bool is_network = false;
  519. if (strncmp(ffm->params.file, SRT_PROTO, sizeof(SRT_PROTO) - 1) == 0 ||
  520. strncmp(ffm->params.file, UDP_PROTO, sizeof(UDP_PROTO) - 1) == 0 ||
  521. strncmp(ffm->params.file, TCP_PROTO, sizeof(TCP_PROTO) - 1) == 0)
  522. is_network = true;
  523. if (is_network) {
  524. avformat_network_init();
  525. output_format = av_guess_format("mpegts", NULL, "video/M2PT");
  526. } else {
  527. output_format = av_guess_format(NULL, ffm->params.file, NULL);
  528. }
  529. if (output_format == NULL) {
  530. fprintf(stderr, "Couldn't find an appropriate muxer for '%s'\n",
  531. ffm->params.printable_file.array);
  532. return FFM_ERROR;
  533. }
  534. ret = avformat_alloc_output_context2(&ffm->output, output_format, NULL,
  535. ffm->params.file);
  536. if (ret < 0) {
  537. fprintf(stderr, "Couldn't initialize output context: %s\n",
  538. av_err2str(ret));
  539. return FFM_ERROR;
  540. }
  541. ffm->output->oformat->video_codec = AV_CODEC_ID_NONE;
  542. ffm->output->oformat->audio_codec = AV_CODEC_ID_NONE;
  543. if (!init_streams(ffm)) {
  544. free_avformat(ffm);
  545. return FFM_ERROR;
  546. }
  547. ret = open_output_file(ffm);
  548. if (ret != FFM_SUCCESS) {
  549. free_avformat(ffm);
  550. return ret;
  551. }
  552. return FFM_SUCCESS;
  553. }
  554. static int ffmpeg_mux_init_internal(struct ffmpeg_mux *ffm, int argc,
  555. char *argv[])
  556. {
  557. argc--;
  558. argv++;
  559. if (!init_params(&argc, &argv, &ffm->params, &ffm->audio))
  560. return FFM_ERROR;
  561. if (ffm->params.tracks) {
  562. ffm->audio_header =
  563. calloc(ffm->params.tracks, sizeof(*ffm->audio_header));
  564. }
  565. #if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(58, 9, 100)
  566. av_register_all();
  567. #endif
  568. if (!ffmpeg_mux_get_extra_data(ffm))
  569. return FFM_ERROR;
  570. /* ffmpeg does not have a way of telling what's supported
  571. * for a given output format, so we try each possibility */
  572. return ffmpeg_mux_init_context(ffm);
  573. }
  574. static int ffmpeg_mux_init(struct ffmpeg_mux *ffm, int argc, char *argv[])
  575. {
  576. int ret = ffmpeg_mux_init_internal(ffm, argc, argv);
  577. if (ret != FFM_SUCCESS) {
  578. ffmpeg_mux_free(ffm);
  579. return ret;
  580. }
  581. ffm->initialized = true;
  582. return ret;
  583. }
  584. static inline int get_index(struct ffmpeg_mux *ffm,
  585. struct ffm_packet_info *info)
  586. {
  587. if (info->type == FFM_PACKET_VIDEO) {
  588. if (ffm->video_stream) {
  589. return ffm->video_stream->id;
  590. }
  591. } else {
  592. if ((int)info->index < ffm->num_audio_streams) {
  593. return ffm->audio_infos[info->index].stream->id;
  594. }
  595. }
  596. return -1;
  597. }
  598. static AVCodecContext *get_codec_context(struct ffmpeg_mux *ffm,
  599. struct ffm_packet_info *info)
  600. {
  601. if (info->type == FFM_PACKET_VIDEO) {
  602. if (ffm->video_stream) {
  603. return ffm->video_ctx;
  604. }
  605. } else {
  606. if ((int)info->index < ffm->num_audio_streams) {
  607. return ffm->audio_infos[info->index].ctx;
  608. }
  609. }
  610. return NULL;
  611. }
  612. static inline AVStream *get_stream(struct ffmpeg_mux *ffm, int idx)
  613. {
  614. return ffm->output->streams[idx];
  615. }
  616. static inline int64_t rescale_ts(struct ffmpeg_mux *ffm,
  617. AVRational codec_time_base, int64_t val,
  618. int idx)
  619. {
  620. AVStream *stream = get_stream(ffm, idx);
  621. return av_rescale_q_rnd(val / codec_time_base.num, codec_time_base,
  622. stream->time_base,
  623. AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX);
  624. }
  625. static inline bool ffmpeg_mux_packet(struct ffmpeg_mux *ffm, uint8_t *buf,
  626. struct ffm_packet_info *info)
  627. {
  628. int idx = get_index(ffm, info);
  629. AVPacket packet = {0};
  630. /* The muxer might not support video/audio, or multiple audio tracks */
  631. if (idx == -1) {
  632. return true;
  633. }
  634. const AVRational codec_time_base =
  635. get_codec_context(ffm, info)->time_base;
  636. av_init_packet(&packet);
  637. packet.data = buf;
  638. packet.size = (int)info->size;
  639. packet.stream_index = idx;
  640. packet.pts = rescale_ts(ffm, codec_time_base, info->pts, idx);
  641. packet.dts = rescale_ts(ffm, codec_time_base, info->dts, idx);
  642. if (info->keyframe)
  643. packet.flags = AV_PKT_FLAG_KEY;
  644. int ret = av_interleaved_write_frame(ffm->output, &packet);
  645. if (ret < 0) {
  646. fprintf(stderr, "av_interleaved_write_frame failed: %s\n",
  647. av_err2str(ret));
  648. }
  649. return ret >= 0;
  650. }
  651. /* ------------------------------------------------------------------------- */
  652. #ifdef _WIN32
  653. int wmain(int argc, wchar_t *argv_w[])
  654. #else
  655. int main(int argc, char *argv[])
  656. #endif
  657. {
  658. struct ffm_packet_info info = {0};
  659. struct ffmpeg_mux ffm = {0};
  660. struct resize_buf rb = {0};
  661. bool fail = false;
  662. int ret;
  663. #ifdef _WIN32
  664. char **argv;
  665. SetErrorMode(SEM_FAILCRITICALERRORS);
  666. argv = malloc(argc * sizeof(char *));
  667. for (int i = 0; i < argc; i++) {
  668. size_t len = wcslen(argv_w[i]);
  669. int size;
  670. size = WideCharToMultiByte(CP_UTF8, 0, argv_w[i], (int)len,
  671. NULL, 0, NULL, NULL);
  672. argv[i] = malloc(size + 1);
  673. WideCharToMultiByte(CP_UTF8, 0, argv_w[i], (int)len, argv[i],
  674. size + 1, NULL, NULL);
  675. argv[i][size] = 0;
  676. }
  677. _setmode(_fileno(stdin), O_BINARY);
  678. #endif
  679. setvbuf(stderr, NULL, _IONBF, 0);
  680. ret = ffmpeg_mux_init(&ffm, argc, argv);
  681. if (ret != FFM_SUCCESS) {
  682. fprintf(stderr, "Couldn't initialize muxer\n");
  683. return ret;
  684. }
  685. while (!fail && safe_read(&info, sizeof(info)) == sizeof(info)) {
  686. resize_buf_resize(&rb, info.size);
  687. if (safe_read(rb.buf, info.size) == info.size) {
  688. fail = !ffmpeg_mux_packet(&ffm, rb.buf, &info);
  689. } else {
  690. fail = true;
  691. }
  692. }
  693. ffmpeg_mux_free(&ffm);
  694. resize_buf_free(&rb);
  695. #ifdef _WIN32
  696. for (int i = 0; i < argc; i++)
  697. free(argv[i]);
  698. free(argv);
  699. #endif
  700. return 0;
  701. }