StreamRecoder.cc 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. #include <stdio.h>
  2. #include <stdlib.h>
  3. #include <iostream>
  4. #include <fstream>
  5. #include <sstream>
  6. extern "C" {
  7. #include <libavcodec/avcodec.h>
  8. #include <libavformat/avformat.h>
  9. #include <libavformat/avio.h>
  10. #include <libswscale/swscale.h>
  11. }
  12. int main(int argc, char** argv) {
  13. // Open the initial context variables that are needed
  14. SwsContext *img_convert_ctx;
  15. AVFormatContext* format_ctx = avformat_alloc_context();
  16. AVCodecContext* codec_ctx = NULL;
  17. int video_stream_index;
  18. // Register everything
  19. av_register_all();
  20. avformat_network_init();
  21. //open RTSP
  22. if (avformat_open_input(&format_ctx, "rtsp://134.169.178.187:8554/h264.3gp",
  23. NULL, NULL) != 0) {
  24. return EXIT_FAILURE;
  25. }
  26. if (avformat_find_stream_info(format_ctx, NULL) < 0) {
  27. return EXIT_FAILURE;
  28. }
  29. //search video stream
  30. for (int i = 0; i < format_ctx->nb_streams; i++) {
  31. if (format_ctx->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO)
  32. video_stream_index = i;
  33. }
  34. AVPacket packet;
  35. av_init_packet(&packet);
  36. //open output file
  37. AVFormatContext* output_ctx = avformat_alloc_context();
  38. AVStream* stream = NULL;
  39. int cnt = 0;
  40. //start reading packets from stream and write them to file
  41. av_read_play(format_ctx); //play RTSP
  42. // Get the codec
  43. AVCodec *codec = NULL;
  44. codec = avcodec_find_decoder(AV_CODEC_ID_H264);
  45. if (!codec) {
  46. exit(1);
  47. }
  48. // Add this to allocate the context by codec
  49. codec_ctx = avcodec_alloc_context3(codec);
  50. avcodec_get_context_defaults3(codec_ctx, codec);
  51. avcodec_copy_context(codec_ctx, format_ctx->streams[video_stream_index]->codec);
  52. std::ofstream output_file;
  53. if (avcodec_open2(codec_ctx, codec, NULL) < 0)
  54. exit(1);
  55. img_convert_ctx = sws_getContext(codec_ctx->width, codec_ctx->height,
  56. codec_ctx->pix_fmt, codec_ctx->width, codec_ctx->height, AV_PIX_FMT_RGB24,
  57. SWS_BICUBIC, NULL, NULL, NULL);
  58. int size = avpicture_get_size(AV_PIX_FMT_YUV420P, codec_ctx->width,
  59. codec_ctx->height);
  60. uint8_t* picture_buffer = (uint8_t*) (av_malloc(size));
  61. AVFrame* picture = av_frame_alloc();
  62. AVFrame* picture_rgb = av_frame_alloc();
  63. int size2 = avpicture_get_size(AV_PIX_FMT_RGB24, codec_ctx->width,
  64. codec_ctx->height);
  65. uint8_t* picture_buffer_2 = (uint8_t*) (av_malloc(size2));
  66. avpicture_fill((AVPicture *) picture, picture_buffer, AV_PIX_FMT_YUV420P,
  67. codec_ctx->width, codec_ctx->height);
  68. avpicture_fill((AVPicture *) picture_rgb, picture_buffer_2, AV_PIX_FMT_RGB24,
  69. codec_ctx->width, codec_ctx->height);
  70. while (av_read_frame(format_ctx, &packet) >= 0 && cnt < 1000) { //read ~ 1000 frames
  71. std::cout << "1 Frame: " << cnt << std::endl;
  72. if (packet.stream_index == video_stream_index) { //packet is video
  73. std::cout << "2 Is Video" << std::endl;
  74. if (stream == NULL) { //create stream in file
  75. std::cout << "3 create stream" << std::endl;
  76. stream = avformat_new_stream(output_ctx,
  77. format_ctx->streams[video_stream_index]->codec->codec);
  78. avcodec_copy_context(stream->codec,
  79. format_ctx->streams[video_stream_index]->codec);
  80. stream->sample_aspect_ratio =
  81. format_ctx->streams[video_stream_index]->codec->sample_aspect_ratio;
  82. }
  83. int check = 0;
  84. packet.stream_index = stream->id;
  85. std::cout << "4 decoding" << std::endl;
  86. int result = avcodec_decode_video2(codec_ctx, picture, &check, &packet);
  87. std::cout << "Bytes decoded " << result << " check " << check
  88. << std::endl;
  89. if (cnt > 100) //cnt < 0)
  90. {
  91. sws_scale(img_convert_ctx, picture->data, picture->linesize, 0,
  92. codec_ctx->height, picture_rgb->data, picture_rgb->linesize);
  93. std::stringstream file_name;
  94. file_name << "test" << cnt << ".ppm";
  95. output_file.open(file_name.str().c_str());
  96. output_file << "P3 " << codec_ctx->width << " " << codec_ctx->height
  97. << " 255\n";
  98. for (int y = 0; y < codec_ctx->height; y++) {
  99. for (int x = 0; x < codec_ctx->width * 3; x++)
  100. output_file
  101. << (int) (picture_rgb->data[0]
  102. + y * picture_rgb->linesize[0])[x] << " ";
  103. }
  104. output_file.close();
  105. }
  106. cnt++;
  107. }
  108. av_free_packet(&packet);
  109. av_init_packet(&packet);
  110. }
  111. av_free(picture);
  112. av_free(picture_rgb);
  113. av_free(picture_buffer);
  114. av_free(picture_buffer_2);
  115. av_read_pause(format_ctx);
  116. avio_close(output_ctx->pb);
  117. avformat_free_context(output_ctx);
  118. return (EXIT_SUCCESS);
  119. }