Cacher is the code snippet organizer for pro developers

/*
 * rwav.c
 *
 *      Author: Arash Shafiei
 *      Email: arash dot shafiei at gmail dot com
 */

#include "rwav.h"

#define MAX_AUDIO_PACKET_SIZE (128 * 1024)

void rwav_register_all() {

	av_register_all();
	avcodec_register_all();
	avdevice_register_all();
	avformat_network_init();

}

int rwav_open_input_audio(InputAudioContext * p_in_audio_ctx) {

	int i;
	AVCodecContext * p_codec_ctx;
	AVCodec * p_codec;

	p_in_audio_ctx->p_fmt_ctx = NULL;

	// Open video
	if (avformat_open_input(&p_in_audio_ctx->p_fmt_ctx,
			p_in_audio_ctx->psz_name, NULL, NULL) != 0) {
		fprintf(stderr, "Cannot open file\n");
		return -1;
	}

	// Retrieve stream information
	if (avformat_find_stream_info(p_in_audio_ctx->p_fmt_ctx, NULL) < 0) {
		fprintf(stderr, "Cannot find stream information\n");
		return -1;
	}

	av_dump_format(p_in_audio_ctx->p_fmt_ctx, 0, p_in_audio_ctx->psz_name, 0);

	// Find the first video stream
	p_in_audio_ctx->i_audio_stream_idx = -1;
	for (i = 0; i < p_in_audio_ctx->p_fmt_ctx->nb_streams; i++) {
		if (p_in_audio_ctx->p_fmt_ctx->streams[i]->codec->codec_type
				== AVMEDIA_TYPE_AUDIO) {
			p_in_audio_ctx->i_audio_stream_idx = i;
		}
	}
	if (p_in_audio_ctx->i_audio_stream_idx == -1) {
		fprintf(stderr, "Cannot find a video stream\n");
		return -1;
	}

	// Get a pointer to the p_codec context for the video stream
	p_codec_ctx =
			p_in_audio_ctx->p_fmt_ctx->streams[p_in_audio_ctx->i_audio_stream_idx]->codec;

	// Find the decoder for the video stream
	p_codec = avcodec_find_decoder(p_codec_ctx->codec_id);
	if (p_codec == NULL) {
		fprintf(stderr, "Input audio p_codec is not supported.\n");
		avformat_close_input(&p_in_audio_ctx->p_fmt_ctx);
		return -1;
	}

	// Open p_codec
	if (avcodec_open2(p_codec_ctx, p_codec, NULL) < 0) {
		fprintf(stderr, "Cannot open input audio p_codec.\n");
		avformat_close_input(&p_in_audio_ctx->p_fmt_ctx);
		return -1;
	}

	// Allocate video frame
	p_in_audio_ctx->p_audio_frame = avcodec_alloc_frame();

	return 0;
}

int rwav_open_input_video(InputVideoContext * p_in_video_ctx) {

	int i;
	AVCodecContext * p_codec_ctx;
	AVCodec * p_codec;

	p_in_video_ctx->p_fmt_ctx = NULL;

	// Open video url
	if (avformat_open_input(&p_in_video_ctx->p_fmt_ctx,
			p_in_video_ctx->psz_name, NULL, NULL) != 0) {
		fprintf(stderr, "Cannot open file\n");
		return -1;
	}

	// Retrieve stream information
	if (avformat_find_stream_info(p_in_video_ctx->p_fmt_ctx, NULL) < 0) {
		fprintf(stderr, "Cannot find stream information\n");
		return -1;
	}

	av_dump_format(p_in_video_ctx->p_fmt_ctx, 0, p_in_video_ctx->psz_name, 0);

	// Find the first video stream
	p_in_video_ctx->i_video_stream_idx = -1;
	for (i = 0; i < p_in_video_ctx->p_fmt_ctx->nb_streams; i++) {
		if (p_in_video_ctx->p_fmt_ctx->streams[i]->codec->codec_type
				== AVMEDIA_TYPE_VIDEO) {
			p_in_video_ctx->i_video_stream_idx = i;
		}
	}
	if (p_in_video_ctx->i_video_stream_idx == -1) {
		fprintf(stderr, "Cannot find a video stream\n");
		return -1;
	}

	// Get a pointer to the codec context for the video stream
	p_codec_ctx =
			p_in_video_ctx->p_fmt_ctx->streams[p_in_video_ctx->i_video_stream_idx]->codec;

	// Find the decoder for the video stream
	p_codec = avcodec_find_decoder(p_codec_ctx->codec_id);
	if (p_codec == NULL) {
		fprintf(stderr, "Codec is not supported.\n");
		avformat_close_input(&p_in_video_ctx->p_fmt_ctx);
		return -1;
	}

	// Open codec
	if (avcodec_open2(p_codec_ctx, p_codec, NULL) < 0) {
		fprintf(stderr, "Cannot open codec.\n");
		avformat_close_input(&p_in_video_ctx->p_fmt_ctx);
		return -1;
	}

	// Allocate video frame
	p_in_video_ctx->p_video_frame = avcodec_alloc_frame();

	return 0;
}

int rwav_open_output_file(OutputFileContext * p_outfile_ctx) {

	p_outfile_ctx->p_fmt_ctx = NULL;
	AVOutputFormat * p_output_fmt;

	//Find output format
	p_output_fmt = av_guess_format(NULL, p_outfile_ctx->psz_name, NULL);
	if (!p_output_fmt) {
		fprintf(stderr, "Cannot find suitable output format\n");
		return -1;
	}

	p_outfile_ctx->p_fmt_ctx = avformat_alloc_context();
	if (!p_outfile_ctx->p_fmt_ctx) {
		fprintf(stderr, "Cannot allocate memory for pOutVideoFormatCtx\n");
		return -1;
	}

	p_outfile_ctx->p_fmt_ctx->oformat = p_output_fmt;
	strcpy(p_outfile_ctx->p_fmt_ctx->filename, p_outfile_ctx->psz_name);

	// Open the output file
	if (!(p_output_fmt->flags & AVFMT_NOFILE)) {
		if (avio_open(&p_outfile_ctx->p_fmt_ctx->pb, p_outfile_ctx->psz_name,
				URL_WRONLY) < 0) {
			fprintf(stderr, "Cannot not open '%s'\n", p_outfile_ctx->psz_name);
			return -1;
		}
	}

	return 0;
}

int rwav_open_output_video_stream(OutputFileContext * p_outfile_ctx,
		OutputVideoContext * p_out_video_ctx) {

	AVCodec * p_video_codec;
	AVStream * p_video_stream;

	int i_video_codec_id = CODEC_ID_H264;

	p_video_codec = avcodec_find_encoder(i_video_codec_id);
	if (p_video_codec == NULL) {
		fprintf(stderr, "Output video codec not found\n");
		return -1;
	}

	//Create new video stream
	p_video_stream = avformat_new_stream(p_outfile_ctx->p_fmt_ctx,
			p_video_codec);
	if (!p_video_stream) {
		fprintf(stderr, "Cannot create output video stream\n");
		return -1;
	}

	p_video_stream->codec->codec_id = p_video_codec->id;
	p_video_stream->codec->codec_type = AVMEDIA_TYPE_VIDEO;
	p_video_stream->codec->bit_rate = p_out_video_ctx->i_vbitrate;
	p_video_stream->codec->width = p_out_video_ctx->i_vwidth;
	p_video_stream->codec->height = p_out_video_ctx->i_vheight;
	p_video_stream->codec->time_base =
			(AVRational) {1 ,p_out_video_ctx->i_vframerate};
	p_video_stream->codec->pix_fmt = PIX_FMT_YUV420P;

	p_video_stream->codec->max_b_frames = 0;
	p_video_stream->codec->thread_count = 1;
	p_video_stream->codec->delay = 0;
	p_video_stream->codec->gop_size = p_out_video_ctx->i_vframerate;
	//videoStream->codec->gop_size = 1;
	p_video_stream->codec->rc_lookahead = 0;

	//videoStream->time_base = (AVRational) {1 , 1000000};
	//videoStream->r_frame_rate = (AVRational) {outVideoCtx->video_framerate, 1};

	//av_opt_set(videoStream->codec->priv_data, "preset", "slow", 0);
	//videoStream->codec->me_range = 16;
	//videoStream->codec->max_qdiff = 4;
	//videoStream->codec->qmin = 10;
	//videoStream->codec->qmax = 51;
	//videoStream->codec->qcompress = 0.6;
	//videoStream->codec->profile = FF_PROFILE_H264_BASELINE;
	//videoStream->codec->level = 10;

	if (p_outfile_ctx->p_fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
		p_video_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;

	p_outfile_ctx->i_video_stream_idx = p_video_stream->index;

	// open the video codec
	if (avcodec_open2(p_video_stream->codec, p_video_codec, NULL) < 0) {
		fprintf(stderr, "Cannot open output video codec\n");
		return -1;
	}

	//Initialize output frame

	p_out_video_ctx->i_video_outbuf_size = 9 * p_out_video_ctx->i_vwidth
			* p_out_video_ctx->i_vheight + 10000;
	p_out_video_ctx->p_video_outbuf = (uint8_t *) av_malloc(
			p_out_video_ctx->i_video_outbuf_size);

	// Allocate an AVFrame structure
	p_out_video_ctx->p_video_frame = avcodec_alloc_frame();
	if (p_out_video_ctx->p_video_frame == NULL) {
		fprintf(stderr, "Cannot allocate YUV frame!\n");
		return -1;
	}

	// Determine required buffer size and allocate buffer
	int i_num_bytes = avpicture_get_size(PIX_FMT_YUV420P,
			p_out_video_ctx->i_vwidth, p_out_video_ctx->i_vheight);
	p_out_video_ctx->p_video_data_buf = (uint8_t *) av_malloc(
			i_num_bytes * sizeof(uint8_t));

	avpicture_fill((AVPicture *) p_out_video_ctx->p_video_frame,
			p_out_video_ctx->p_video_data_buf, PIX_FMT_YUV420P,
			p_out_video_ctx->i_vwidth, p_out_video_ctx->i_vheight);

	return 0;
}

int rwav_open_output_audio_stream(OutputFileContext * p_outfile_ctx,
		OutputAudioContext * p_out_audio_ctx) {

	AVCodec * p_audio_codec;
	AVStream * p_audio_stream;

	int i_audio_codec_id = CODEC_ID_AAC;

	p_audio_codec = avcodec_find_encoder(i_audio_codec_id);
	if (p_audio_codec == NULL) {
		fprintf(stderr, "Output audio codec not found\n");
		return -1;
	}

	//Create new audio stream
	p_audio_stream = avformat_new_stream(p_outfile_ctx->p_fmt_ctx,
			p_audio_codec);
	if (!p_audio_stream) {
		fprintf(stderr, "Cannot create output video stream\n");
		return -1;
	}

	p_audio_stream->codec->codec_id = p_audio_codec->id;
	p_audio_stream->codec->codec_type = AVMEDIA_TYPE_AUDIO;
	p_audio_stream->codec->bit_rate = p_out_audio_ctx->i_abitrate;
	p_audio_stream->codec->sample_rate = p_out_audio_ctx->i_asamplerate;
	p_audio_stream->codec->channels = p_out_audio_ctx->i_achannels;
	p_audio_stream->codec->sample_fmt = AV_SAMPLE_FMT_S16;

	if (p_outfile_ctx->p_fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
		p_audio_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;

	p_outfile_ctx->i_audio_stream_idx = p_audio_stream->index;

	// open the audio codec
	if (avcodec_open2(p_audio_stream->codec, p_audio_codec, NULL) < 0) {
		fprintf(stderr, "Cannot open output audio codec\n");
		return -1;
	}

	p_out_audio_ctx->p_audio_frame = avcodec_alloc_frame();
	p_out_audio_ctx->p_fifo_buf = av_fifo_alloc(2 * MAX_AUDIO_PACKET_SIZE);
	p_out_audio_ctx->p_audio_data_buf = (uint8_t*) av_malloc(
			2 * MAX_AUDIO_PACKET_SIZE);

	return 0;
}

int rwav_open_video_converter(InputVideoContext * p_in_video_ctx,
		OutputVideoContext * p_out_video_ctx) {

	AVCodecContext * p_codec_ctx;

	// Get a pointer to the codec context for the video stream
	p_codec_ctx =
			p_in_video_ctx->p_fmt_ctx->streams[p_in_video_ctx->i_video_stream_idx]->codec;

	p_out_video_ctx->p_sws_ctx = sws_getContext(p_codec_ctx->width,
			p_codec_ctx->height, p_codec_ctx->pix_fmt,
			p_out_video_ctx->i_vwidth, p_out_video_ctx->i_vheight,
			PIX_FMT_YUV420P, SWS_FAST_BILINEAR, NULL, NULL, NULL);

	if (p_out_video_ctx->p_sws_ctx == NULL) {
		fprintf(stderr, "Cannot initialize the conversion context!\n");
		return -1;
	}

	return 0;
}

int rwav_convert_video_frame(InputVideoContext * p_in_video_ctx,
		OutputVideoContext * p_out_video_ctx) {

	AVCodecContext * p_codec_ctx;

	// Get a pointer to the codec context for the video stream
	p_codec_ctx =
			p_in_video_ctx->p_fmt_ctx->streams[p_in_video_ctx->i_video_stream_idx]->codec;

	sws_scale(p_out_video_ctx->p_sws_ctx,
			(const uint8_t * const *) p_in_video_ctx->p_video_frame->data,
			p_in_video_ctx->p_video_frame->linesize, 0, p_codec_ctx->height,
			p_out_video_ctx->p_video_frame->data,
			p_out_video_ctx->p_video_frame->linesize);

	p_out_video_ctx->p_video_frame->pts = p_in_video_ctx->p_video_frame->pts;

	return 0;
}

int rwav_convert_audio_frame(InputAudioContext * p_in_audio_ctx,
		OutputAudioContext * p_out_audio_ctx) {

	p_out_audio_ctx->p_audio_inbuf = p_in_audio_ctx->p_audio_frame->data[0];
	p_out_audio_ctx->i_audio_inbuf_size =
			p_in_audio_ctx->p_audio_frame->linesize[0];

	return 0;
}

int rwav_read_audio(InputAudioContext * p_in_audio_ctx) {

	AVPacket pkt;
	int i_got_frame;
	AVCodecContext * p_codec_ctx;

	// Get a pointer to the codec context for the video stream
	p_codec_ctx =
			p_in_audio_ctx->p_fmt_ctx->streams[p_in_audio_ctx->i_audio_stream_idx]->codec;

	avcodec_get_frame_defaults(p_in_audio_ctx->p_audio_frame);
	// Read frames
	while (av_read_frame(p_in_audio_ctx->p_fmt_ctx, &pkt) >= 0) {

		// Is this a packet from the audio stream?
		if (pkt.stream_index == p_in_audio_ctx->i_audio_stream_idx) {

			// Decode video frame
			if (avcodec_decode_audio4(p_codec_ctx,
					p_in_audio_ctx->p_audio_frame, &i_got_frame, &pkt) < 0) {
				fprintf(stderr, "Error while decoding audio.\n");
				return -1;
			}

			// Did we get a video frame?
			if (i_got_frame) {

				av_free_packet(&pkt);
				return 0;
			}
		}

		// Free the packet that was allocated by av_read_frame
		av_free_packet(&pkt);
	}

	return -1;

}

int rwav_read_video(InputVideoContext * p_in_video_ctx) {

	AVPacket pkt;
	int i_got_frame;

	// Get a pointer to the codec context for the video stream
	AVCodecContext * p_codec_ctx =
			p_in_video_ctx->p_fmt_ctx->streams[p_in_video_ctx->i_video_stream_idx]->codec;

	avcodec_get_frame_defaults(p_in_video_ctx->p_video_frame);

	// Read frames
	while (av_read_frame(p_in_video_ctx->p_fmt_ctx, &pkt) >= 0) {

		// Is this a packet from the video stream?
		if (pkt.stream_index == p_in_video_ctx->i_video_stream_idx) {

			// Decode video frame
			avcodec_decode_video2(p_codec_ctx, p_in_video_ctx->p_video_frame,
					&i_got_frame, &pkt);

			// Did we get a video frame?
			if (i_got_frame) {

				av_free_packet(&pkt);
				return 0;
			}
		}

		// Free the packet that was allocated by av_read_frame
		av_free_packet(&pkt);
	}

	return -1;
}

int rwav_write_video(OutputFileContext * p_outfile_ctx,
		OutputVideoContext * p_out_video_ctx) {

	AVPacket pkt;

	int i_out_size = -1;

	AVStream * p_video_stream =
			p_outfile_ctx->p_fmt_ctx->streams[p_outfile_ctx->i_video_stream_idx];
	AVCodecContext * p_video_codec_ctx = p_video_stream->codec;

	//Set PTS
	p_out_video_ctx->p_video_frame->pts = p_video_codec_ctx->frame_number;

	//Set PTS (alternative method)
	//int64_t now = av_gettime();
	//outVideoCtx->pOutVideoFrame->pts = av_rescale_q(now,AV_TIME_BASE_Q,
	//		videoCodecCtx->time_base);

	//Encoding video
	i_out_size = avcodec_encode_video(p_video_codec_ctx,
			p_out_video_ctx->p_video_outbuf,
			p_out_video_ctx->i_video_outbuf_size,
			p_out_video_ctx->p_video_frame);

	//printf("After pts: %d\n", outVideoCtx->pOutVideoFrame->pts);

	// if zero size, it means the image was buffered
	if (i_out_size > 0) {

		av_init_packet(&pkt);

		if (p_video_codec_ctx->coded_frame->pts != AV_NOPTS_VALUE) {
			pkt.pts = av_rescale_q(p_video_codec_ctx->coded_frame->pts,
					p_video_codec_ctx->time_base, p_video_stream->time_base);
		}

		if (p_video_codec_ctx->coded_frame->key_frame)
			pkt.flags |= AV_PKT_FLAG_KEY;

		pkt.stream_index = p_video_stream->index;
		pkt.data = p_out_video_ctx->p_video_outbuf;
		pkt.size = i_out_size;

		// write the compressed frame in the media file
		if (av_interleaved_write_frame(p_outfile_ctx->p_fmt_ctx, &pkt) != 0)
			fprintf(stderr, "Writing frame is not successful\n");

		av_free_packet(&pkt);

	}

	return i_out_size;
}

int rwav_write_audio(OutputFileContext * p_outfile_ctx,
		OutputAudioContext * p_out_audio_ctx) {

	AVPacket pkt;
	int i_got_pkt;
	int i_frame_bytes;

	AVStream * p_audio_stream =
			p_outfile_ctx->p_fmt_ctx->streams[p_outfile_ctx->i_audio_stream_idx];
	AVCodecContext * p_audio_codec_ctx = p_audio_stream->codec;

	int i_osize = av_get_bytes_per_sample(p_audio_codec_ctx->sample_fmt);

	//Encoding audio
	av_fifo_generic_write(p_out_audio_ctx->p_fifo_buf,
			p_out_audio_ctx->p_audio_inbuf, p_out_audio_ctx->i_audio_inbuf_size,
			NULL);

	i_frame_bytes = p_audio_codec_ctx->frame_size * i_osize
			* p_audio_codec_ctx->channels;

	while (av_fifo_size(p_out_audio_ctx->p_fifo_buf) >= i_frame_bytes) {

		av_fifo_generic_read(p_out_audio_ctx->p_fifo_buf,
				p_out_audio_ctx->p_audio_data_buf, i_frame_bytes, NULL);

		avcodec_get_frame_defaults(p_out_audio_ctx->p_audio_frame);

		p_out_audio_ctx->p_audio_frame->nb_samples =
				i_frame_bytes
						/ (p_audio_codec_ctx->channels
								* av_get_bytes_per_sample(
										p_audio_codec_ctx->sample_fmt));

		if (avcodec_fill_audio_frame(p_out_audio_ctx->p_audio_frame,
				p_audio_codec_ctx->channels, p_audio_codec_ctx->sample_fmt,
				p_out_audio_ctx->p_audio_data_buf, i_frame_bytes, 1) < 0) {
			fprintf(stderr, "Fill audio frame failed\n");
			return -1;
		}

		av_init_packet(&pkt);

		//Set PTS
		//int64_t now = av_gettime();
		//outAudioCtx->pOutAudioFrame->pts = av_rescale_q(now,AV_TIME_BASE_Q,
		//		audioEncCtx->time_base);

		if (avcodec_encode_audio2(p_audio_codec_ctx, &pkt,
				p_out_audio_ctx->p_audio_frame, &i_got_pkt) != 0) {

			fprintf(stderr, "Error while encoding audio.\n");
			return -1;
		}

		if (i_got_pkt) {

			pkt.stream_index = p_audio_stream->index;

			/*
			 if (pkt.pts != AV_NOPTS_VALUE)
			 pkt.pts = av_rescale_q(pkt.pts, audioEncCtx->time_base,
			 audioStream->time_base);
			 */

			pkt.flags |= AV_PKT_FLAG_KEY;

			if (av_interleaved_write_frame(p_outfile_ctx->p_fmt_ctx, &pkt)
					!= 0) {

				fprintf(stderr, "Writing frame is not successful\n");
				av_free_packet(&pkt);
				return -1;
			}

			av_free_packet(&pkt);

		}

	}

	return 0;
}

int rwav_write_header(OutputFileContext * p_outfile_ctx) {

	// write the stream header, if any
	avformat_write_header(p_outfile_ctx->p_fmt_ctx, NULL);

	return 0;
}

int rwav_write_tailer(OutputFileContext * p_outfile_ctx) {

	// write the trailer, if any
	av_write_trailer(p_outfile_ctx->p_fmt_ctx);
	return 0;
}

int rwav_close_input_audio(InputAudioContext * p_in_audio_ctx) {

	// Close the audio format context
	avformat_close_input(&p_in_audio_ctx->p_fmt_ctx);

	// Free the audio frame
	av_free(p_in_audio_ctx->p_audio_frame);

	return 0;
}

int rwav_close_input_video(InputVideoContext * p_in_video_ctx) {

	// Close the video file
	avformat_close_input(&p_in_video_ctx->p_fmt_ctx);

	// Free the YUV frame
	av_free(p_in_video_ctx->p_video_frame);

	return 0;
}

int rwav_close_output_file(OutputFileContext * p_outfile_ctx) {

	int i;

	avio_close(p_outfile_ctx->p_fmt_ctx->pb);

	// free the streams
	for (i = 0; i < p_outfile_ctx->p_fmt_ctx->nb_streams; i++) {
		avcodec_close(p_outfile_ctx->p_fmt_ctx->streams[i]->codec);
		av_freep(&p_outfile_ctx->p_fmt_ctx->streams[i]->info);
	}

	// free the stream
	avformat_free_context(p_outfile_ctx->p_fmt_ctx);

	return 0;
}

int rwav_close_output_video_stream(OutputVideoContext * p_out_video_ctx) {

	av_free(p_out_video_ctx->p_video_outbuf);
	av_free(p_out_video_ctx->p_video_frame);
	av_free(p_out_video_ctx->p_sws_ctx);
	av_free(p_out_video_ctx->p_video_data_buf);

	return 0;
}
int rwav_close_output_audio_stream(OutputAudioContext * p_out_audio_ctx) {

	av_fifo_free(p_out_audio_ctx->p_fifo_buf);
	av_free(p_out_audio_ctx->p_audio_data_buf);
	av_free(p_out_audio_ctx->p_audio_frame);

	return 0;
}