codekingpro/portable-devtools

1/*
2 * This copyright notice applies to this header file only:
3 *
4 * Copyright (c) 2010-2024 NVIDIA Corporation
5 *
6 * Permission is hereby granted, free of charge, to any person
7 * obtaining a copy of this software and associated documentation
8 * files (the "Software"), to deal in the Software without
9 * restriction, including without limitation the rights to use,
10 * copy, modify, merge, publish, distribute, sublicense, and/or sell
11 * copies of the software, and to permit persons to whom the
12 * software is furnished to do so, subject to the following
13 * conditions:
14 *
15 * The above copyright notice and this permission notice shall be
16 * included in all copies or substantial portions of the Software.
17 *
18 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
19 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
20 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
21 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
22 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
23 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
24 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
25 * OTHER DEALINGS IN THE SOFTWARE.
26 */
27#pragma once
28
29#include <thread>
30#include <mutex>
31extern "C" {
32#include <libavformat/avformat.h>
33#include <libavutil/opt.h>
34#include <libswresample/swresample.h>
35};
36#include "Logger.h"
37
38using namespace std;
39
40extern simplelogger::Logger *logger;
41
42static string AvErrorToString(int av_error_code) {
43    const auto buf_size = 1024U;
44    char *err_string = (char *)calloc(buf_size, sizeof(*err_string));
45    if (!err_string) {
46        return string();
47    }
48
49    if (0 != av_strerror(av_error_code, err_string, buf_size - 1)) {
50        free(err_string);
51        stringstream ss;
52        ss << "Unknown error with code " << av_error_code;
53        return ss.str();
54    }
55
56    string str(err_string);
57    free(err_string);
58    return str;
59}
60
61class FFmpegStreamer {
62private:
63    AVFormatContext *oc = NULL;
64    AVStream *vs = NULL;
65    int nFps = 0;
66
67public:
68    FFmpegStreamer(AVCodecID eCodecId, int nWidth, int nHeight, int nFps, const char *szInFilePath) : nFps(nFps) {
69        avformat_network_init();
70
71        int ret = 0;
72
73        if ((eCodecId == AV_CODEC_ID_H264) || (eCodecId == AV_CODEC_ID_HEVC))
74            ret = avformat_alloc_output_context2(&oc, NULL, "mpegts", NULL);
75        else if (eCodecId == AV_CODEC_ID_AV1)
76            ret = avformat_alloc_output_context2(&oc, NULL, "ivf", NULL);
77
78        if (ret < 0) {
79            LOG(ERROR) << "FFmpeg: failed to allocate an AVFormatContext. Error message: "
80                       << AvErrorToString(ret);
81            return;
82        }
83
84        oc->url = av_strdup(szInFilePath);
85        LOG(INFO) << "Streaming destination: " << oc->url;
86
87        // Add video stream to oc
88        vs = avformat_new_stream(oc, NULL);
89        if (!vs) {
90            LOG(ERROR) << "FFMPEG: Could not alloc video stream";
91            return;
92        }
93        vs->id = 0;
94
95        // Set video parameters
96        AVCodecParameters *vpar = vs->codecpar;
97        vpar->codec_id = eCodecId;
98        vpar->codec_type = AVMEDIA_TYPE_VIDEO;
99        vpar->width = nWidth;
100        vpar->height = nHeight;
101
102        // Everything is ready. Now open the output stream.
103        if (avio_open(&oc->pb, oc->url, AVIO_FLAG_WRITE) < 0) {
104            LOG(ERROR) << "FFMPEG: Could not open " << oc->url;
105            return ;
106        }
107
108        // Write the container header
109        if (avformat_write_header(oc, NULL)) {
110            LOG(ERROR) << "FFMPEG: avformat_write_header error!";
111            return;
112        }
113    }
114    ~FFmpegStreamer() {
115        if (oc) {
116            av_write_trailer(oc);
117            avio_close(oc->pb);
118            avformat_free_context(oc);
119        }
120    }
121
122    bool Stream(uint8_t *pData, int nBytes, int nPts) {
123        AVPacket *pkt = av_packet_alloc();
124        if (!pkt) {
125            LOG(ERROR) << "AVPacket allocation failed !";
126            return false;
127        }
128        pkt->pts = av_rescale_q(nPts++, AVRational {1, nFps}, vs->time_base);
129        // No B-frames
130        pkt->dts = pkt->pts;
131        pkt->stream_index = vs->index;
132        pkt->data = pData;
133        pkt->size = nBytes;
134
135        if(!memcmp(pData, "\x00\x00\x00\x01\x67", 5)) {
136            pkt->flags |= AV_PKT_FLAG_KEY;
137        }
138
139        // Write the compressed frame into the output
140        int ret = av_write_frame(oc, pkt);
141        av_write_frame(oc, NULL);
142        if (ret < 0) {
143            LOG(ERROR) << "FFMPEG: Error while writing video frame";
144        }
145
146        av_packet_free(&pkt);
147        return true;
148    }
149};
150