Skip to content

Instantly share code, notes, and snippets.

@asmwarrior
Created May 2, 2023 13:50
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save asmwarrior/701bdd914c53d6853beb6e50dd99c6f5 to your computer and use it in GitHub Desktop.
Save asmwarrior/701bdd914c53d6853beb6e50dd99c6f5 to your computer and use it in GitHub Desktop.
#include <iostream>
#include <vector>
#include <cstring>
#include <fstream>
#include <sstream>
#include <stdexcept>
#include <opencv2/opencv.hpp>
extern "C" {
#include <libavutil/imgutils.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/opt.h>
}
#include<cstdlib> // to generate time stamps
using namespace std;
using namespace cv;
int main()
{
// Set up input frames as BGR byte arrays
vector<Mat> frames;
int width = 640;
int height = 480;
int num_frames = 100;
Scalar black(0, 0, 0);
Scalar white(255, 255, 255);
int font = FONT_HERSHEY_SIMPLEX;
double font_scale = 1.0;
int thickness = 2;
for (int i = 0; i < num_frames; i++) {
Mat frame = Mat::zeros(height, width, CV_8UC3);
putText(frame, std::to_string(i), Point(width / 2 - 50, height / 2), font, font_scale, white, thickness);
frames.push_back(frame);
}
// generate a serial of time stamps which is used to set the PTS value
// suppose they are in ms unit, the time interval is between 30ms to 59ms
vector<int> timestamps;
for (int i = 0; i < num_frames; i++) {
int timestamp;
if (i == 0)
timestamp = 0;
else
{
int random = 30 + (rand() % 30);
timestamp = timestamps[i-1] + random;
}
timestamps.push_back(timestamp);
}
// Populate frames with BGR byte arrays
// Initialize FFmpeg
//av_register_all();
// Set up output file
AVFormatContext* outFormatCtx = nullptr;
//AVCodec* outCodec = nullptr;
AVCodecContext* outCodecCtx = nullptr;
//AVStream* outStream = nullptr;
//AVPacket outPacket;
const char* outFile = "output.mp4";
int outWidth = frames[0].cols;
int outHeight = frames[0].rows;
int fps = 25;
// Open the output file context
avformat_alloc_output_context2(&outFormatCtx, nullptr, nullptr, outFile);
if (!outFormatCtx) {
cerr << "Error: Could not allocate output format context" << endl;
return -1;
}
// Open the output file
if (avio_open(&outFormatCtx->pb, outFile, AVIO_FLAG_WRITE) < 0) {
cerr << "Error opening output file" << std::endl;
return -1;
}
// Set up output codec
const AVCodec* outCodec = avcodec_find_encoder(AV_CODEC_ID_H264);
if (!outCodec) {
cerr << "Error: Could not find H.264 codec" << endl;
return -1;
}
outCodecCtx = avcodec_alloc_context3(outCodec);
if (!outCodecCtx) {
cerr << "Error: Could not allocate output codec context" << endl;
return -1;
}
outCodecCtx->codec_id = AV_CODEC_ID_H264;
outCodecCtx->codec_type = AVMEDIA_TYPE_VIDEO;
outCodecCtx->pix_fmt = AV_PIX_FMT_YUV420P;
outCodecCtx->width = outWidth;
outCodecCtx->height = outHeight;
outCodecCtx->time_base = { 1, fps*1000 }; // 25000
//outCodecCtx->time_base = { 1, fps}; // 25000
outCodecCtx->framerate = {fps, 1}; // 25
outCodecCtx->bit_rate = 4000000;
//https://github.com/leandromoreira/ffmpeg-libav-tutorial
//We set the flag AV_CODEC_FLAG_GLOBAL_HEADER which tells the encoder that it can use the global headers.
if (outFormatCtx->oformat->flags & AVFMT_GLOBALHEADER)
{
outCodecCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER; //
}
// Open output codec
if (avcodec_open2(outCodecCtx, outCodec, nullptr) < 0) {
cerr << "Error: Could not open output codec" << endl;
return -1;
}
// Create output stream
AVStream* outStream = avformat_new_stream(outFormatCtx, outCodec);
if (!outStream) {
cerr << "Error: Could not allocate output stream" << endl;
return -1;
}
// Configure output stream parameters (e.g., time base, codec parameters, etc.)
// ...
// Connect output stream to format context
outStream->codecpar->codec_id = outCodecCtx->codec_id;
outStream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
outStream->codecpar->width = outCodecCtx->width;
outStream->codecpar->height = outCodecCtx->height;
outStream->codecpar->format = outCodecCtx->pix_fmt;
outStream->time_base = outCodecCtx->time_base;
int ret = avcodec_parameters_from_context(outStream->codecpar, outCodecCtx);
if (ret < 0) {
cerr << "Error: Could not copy codec parameters to output stream" << endl;
return -1;
}
outStream->avg_frame_rate = outCodecCtx->framerate;
//outStream->id = outFormatCtx->nb_streams++; <--- We shouldn't modify outStream->id
ret = avformat_write_header(outFormatCtx, nullptr);
if (ret < 0) {
cerr << "Error: Could not write output header" << endl;
return -1;
}
// Convert frames to YUV format and write to output file
int frame_count = -1;
AVPacket* outPacket = av_packet_alloc();
for (const auto& frame : frames) {
frame_count++;
AVFrame* yuvFrame = av_frame_alloc();
if (!yuvFrame) {
cerr << "Error: Could not allocate YUV frame" << endl;
return -1;
}
av_image_alloc(yuvFrame->data, yuvFrame->linesize, outWidth, outHeight, AV_PIX_FMT_YUV420P, 32);
yuvFrame->width = outWidth;
yuvFrame->height = outHeight;
yuvFrame->format = AV_PIX_FMT_YUV420P;
// Convert BGR frame to YUV format
Mat yuvMat;
cvtColor(frame, yuvMat, COLOR_BGR2YUV_I420);
memcpy(yuvFrame->data[0], yuvMat.data, outWidth * outHeight);
memcpy(yuvFrame->data[1], yuvMat.data + outWidth * outHeight, outWidth * outHeight / 4);
memcpy(yuvFrame->data[2], yuvMat.data + outWidth * outHeight * 5 / 4, outWidth * outHeight / 4);
// Set up output packet
//av_init_packet(&outPacket); //error C4996: 'av_init_packet': was declared deprecated
memset(outPacket, 0, sizeof(outPacket)); //Use memset instead of av_init_packet (probably unnecessary).
//outPacket->data = nullptr;
//outPacket->size = 0;
// set the frame pts, do I have to set the package pts?
// yuvFrame->pts = av_rescale_q(i, outCodecCtx->time_base, outStream->time_base); //Set PTS timestamp
yuvFrame->pts = av_rescale_q(timestamps[frame_count]*100, outCodecCtx->time_base, outStream->time_base); //Set PTS timestamp
// if (outPacket->pts != AV_NOPTS_VALUE)
// outPacket->pts = av_rescale_q(outPacket->pts, outCodecCtx->time_base, outStream->time_base);
// if (outPacket->dts != AV_NOPTS_VALUE)
// outPacket->dts = av_rescale_q(outPacket->dts, outCodecCtx->time_base, outStream->time_base);
// Encode frame and write to output file
int ret = avcodec_send_frame(outCodecCtx, yuvFrame);
if (ret < 0) {
cerr << "Error: Could not send frame to output codec" << endl;
return -1;
}
while (ret >= 0)
{
ret = avcodec_receive_packet(outCodecCtx, outPacket);
if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
{
break;
}
else if (ret < 0)
{
cerr << "Error: Could not receive packet from output codec" << endl;
return -1;
}
//av_packet_rescale_ts(outPacket, outCodecCtx->time_base, outStream->time_base);
outPacket->stream_index = outStream->index;
//outPacket->dts = av_rescale_q(frame_count, outCodecCtx->time_base, outStream->time_base); //Set PTS timestamp
outPacket->duration = av_rescale_q(30, {1,25}, outStream->time_base); // Set packet duration
ret = av_interleaved_write_frame(outFormatCtx, outPacket);
static int call_write = 0;
call_write++;
printf("av_interleaved_write_frame %d\n", call_write);
av_packet_unref(outPacket);
if (ret < 0) {
cerr << "Error: Could not write packet to output file" << endl;
return -1;
}
//av_packet_free(&outPacket);
}
av_frame_free(&yuvFrame);
}
// Flush the encoder
ret = avcodec_send_frame(outCodecCtx, nullptr);
if (ret < 0) {
std::cerr << "Error flushing encoder: " << std::endl;
return -1;
}
AVPacket* pkt = av_packet_alloc();
while (ret >= 0) {
if (!pkt) {
std::cerr << "Error allocating packet" << std::endl;
return -1;
}
ret = avcodec_receive_packet(outCodecCtx, pkt);
// Write the packet to the output file
if (ret == 0)
{
pkt->stream_index = outStream->index;
pkt->duration = av_rescale_q(1, outCodecCtx->time_base, outStream->time_base); // <---- Set packet duration
ret = av_interleaved_write_frame(outFormatCtx, pkt);
static int call_write1 = 0;
call_write1++;
printf("av_interleaved_write_frame 1 %d\n", call_write1);
av_packet_unref(pkt);
if (ret < 0) {
std::cerr << "Error writing packet to output file: " << std::endl;
return -1;
}
//av_packet_free(&pkt);
}
}
av_packet_free(&pkt);
av_packet_free(&outPacket);
// Write output trailer
av_write_trailer(outFormatCtx);
// Clean up
avcodec_close(outCodecCtx);
avcodec_free_context(&outCodecCtx);
avformat_free_context(outFormatCtx);
return 0;
}
@asmwarrior
Copy link
Author

This is the way I need to learn to read the video back, and show it depends on the PTS.

For python code: python - Getting timestamp of each frame in a video - Stack Overflow

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment