ffmpeg-mux.c 21 KB

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