bluray_remux_invalid_dts.c

#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <stdbool.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <libgen.h>
#include <unistd.h>
#include <signal.h>
#include "libbluray/bluray.h"
#include "libbluray/meta_data.h"
#include <libavutil/dict.h>
#include <libavutil/avutil.h>
#include <libavutil/file.h>
#include <libavutil/opt.h>
#include <libavutil/timestamp.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>

// To compile:
// $ gcc -o bluray_remux bluray_remux.c `pkg-config --libs --cflags libavutil libavcodec libavformat libbluray`
// To run:
// ./bluray_remux <device name>

// Blu-ray block size is 192 bytes (I think)
// bdsplice uses 192 * 1024 (196608)
// ffmpeg chokes on too high a buffer, so for fun (and safety),
// set it to a DVD block length of 2048
#define BLURAY_BUFFER_SIZE 2048

static BLURAY *bd;
static int bluray_packet_offset;

// Quit writing when ctl c happens
static bool interrupted;

void interrupt_me(int i) {
	interrupted = true;
}

int bluray_read_packet(void *opaque, uint8_t *bluray_buffer, int buffer_size) {

	memset(bluray_buffer, '\0', BLURAY_BUFFER_SIZE);

	printf("[bluray_remux] bd_tell(): %" PRIu64 "\r", bd_tell(bd));
	fflush(stdout);
	int size_read = bd_read(bd, bluray_buffer, BLURAY_BUFFER_SIZE);

	if(size_read == 0)
		return AVERROR_EOF;

	bluray_packet_offset++;

	return size_read;

}

struct buffer_data {
	uint8_t *ptr;
	size_t size;
};

int main(int argc, char *argv[]) {

	av_log_set_level(AV_LOG_DEBUG);

	bluray_packet_offset = 0;

	int retval = 0;

	const char *device_filename = NULL;
	if (argv[optind])
		device_filename = argv[optind];
	else
		device_filename = "/dev/sr0";

	// Open the device
	bd = bd_open(device_filename, NULL);
	assert(bd != NULL);

	// Select a bluray playlist (HARD CODED during development)
	retval = bd_select_playlist(bd, 800);
	assert(retval == 1);

	char filename[] = "bluray_remux_00.mkv";

	// New input context
	AVFormatContext *input = avformat_alloc_context();

	// Map bluray_read_packet() function
	struct buffer_data av_buffer_data;
	av_buffer_data.ptr = NULL;
	av_buffer_data.size = BLURAY_BUFFER_SIZE;

	uint8_t *avio_buffer = av_malloc(BLURAY_BUFFER_SIZE);

	// Set bluray copy function as input
	AVIOContext *avio = avio_alloc_context(avio_buffer, BLURAY_BUFFER_SIZE, 0, &av_buffer_data, &bluray_read_packet, NULL, NULL);
	input->pb = avio;

	// Let ffmpeg find the format demuxer (should be mpegts)
	retval = avformat_open_input(&input, NULL, NULL, NULL);
		assert(retval == 0);

	// This will run bluray_read_packet() for the first time
	retval = avformat_find_stream_info(input, NULL);
		assert(retval >= 0);

	// Display progress so far of probed stream
	av_dump_format(input, 0, NULL, 0);

	// Get the stream indexes
	int ix;
	int video_ix = -1;
	unsigned int input_streams = 0, video_streams = 0, audio_streams = 0, pgs_streams = 0;
	for(ix = 0; (unsigned int)ix < input->nb_streams; ix++) {

		input_streams++;

		if(input->streams[ix]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO && video_ix == -1) {
			video_ix = ix;
			video_streams++;
		} else if(input->streams[ix]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
			audio_streams++;
		} else if(input->streams[ix]->codecpar->codec_type == AVMEDIA_TYPE_SUBTITLE) {
			pgs_streams++;
		}

	}

	// Finished with the INPUT
	// Time to move on to the OUPUT

	AVFormatContext *output = NULL;
	retval = avformat_alloc_output_context2(&output, NULL, NULL, filename);
		assert(retval >= 0);
		assert(output != NULL);

	// New video stream
	AVStream *mkv_video = NULL;
	mkv_video = avformat_new_stream(output, NULL);
		assert(mkv_video != NULL);
	retval = avcodec_parameters_copy(mkv_video->codecpar, input->streams[video_ix]->codecpar);
		assert(retval >= 0);
	mkv_video->time_base.num = 1;
	mkv_video->time_base.den = 90000;
	mkv_video->codecpar->codec_tag = 0; // I need to figure out this one

	// Create an audio stream
	AVStream *mkv_audio[audio_streams];
	unsigned int input_stream_ix = video_streams;
	for(ix = 0; ix < audio_streams; ix++, input_stream_ix++) {

		mkv_audio[ix] = avformat_new_stream(output, NULL);

		retval = avcodec_parameters_copy(mkv_audio[ix]->codecpar, input->streams[input_stream_ix]->codecpar);
			assert(retval >= 0);

		mkv_audio[ix]->time_base.num = 1;
		mkv_audio[ix]->time_base.den = 90000;
		mkv_audio[ix]->codecpar->codec_tag = 0;

		av_dict_set(&mkv_audio[ix]->metadata, "language", "eng", 0);

	}

	// Add subtitles
	AVStream *mkv_pgs[pgs_streams];
	input_stream_ix = video_streams + audio_streams;
	for(ix = 0; ix < pgs_streams; ix++, input_stream_ix++) {

		mkv_pgs[ix] = avformat_new_stream(output, NULL);

		retval = avcodec_parameters_copy(mkv_pgs[ix]->codecpar, input->streams[input_stream_ix]->codecpar);
			assert(retval >= 0);

		mkv_pgs[ix]->time_base.num = 1;
		mkv_pgs[ix]->time_base.den = 90000;
		mkv_pgs[ix]->codecpar->codec_tag = 0;

		// Add metadata to audio stream
		retval = av_dict_set(&mkv_pgs[ix]->metadata, "language", "eng", 0);
			assert(retval >= 0);

	}

	// See progress of output streams so far
	av_dump_format(output, 0, filename, 1);
		// return 0;

	// Set muxer name in container :D
	AVDictionary *mkv_opts = NULL;
	retval = av_dict_set(&output->metadata, "encoding_tool", "bluray_remux", 0);
		assert(retval >= 0);

	// Create output file
	avio_open(&output->pb, filename, AVIO_FLAG_WRITE);
		assert(retval >= 0);
	retval = avformat_write_header(output, &mkv_opts);
		assert(retval >= 0);

	// Start reading through packets
	AVPacket bluray_packet;

	// Catch interrupt signal
	interrupted = false;
	signal(SIGINT, interrupt_me);

	while(av_read_frame(input, &bluray_packet) == 0) {

		if(interrupted)
			break;

		// Streams on a Blu-ray should show up at initial probe, so ignore any that show up during remux
		if(bluray_packet.stream_index + 1 > input_streams) {
			av_packet_unref(&bluray_packet);
			continue;
		}

		// Skip corrupt packets (rare)
		if(bluray_packet.flags & AV_PKT_FLAG_CORRUPT) {
			av_packet_unref(&bluray_packet);
			continue;
		}

		if(input->streams[bluray_packet.stream_index]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO && bluray_packet.stream_index == 0) {

			bluray_packet.pts = av_rescale_q_rnd(bluray_packet.pts, input->streams[bluray_packet.stream_index]->time_base, mkv_video->time_base, AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX);
			bluray_packet.dts = av_rescale_q_rnd(bluray_packet.dts, input->streams[bluray_packet.stream_index]->time_base, mkv_video->time_base, AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX);
			bluray_packet.duration = av_rescale_q(bluray_packet.duration, input->streams[bluray_packet.stream_index]->time_base, mkv_video->time_base);

		} else if(input->streams[bluray_packet.stream_index]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {

			bluray_packet.pts = av_rescale_q_rnd(bluray_packet.pts, input->streams[bluray_packet.stream_index]->time_base, mkv_audio[0]->time_base, AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX);
			bluray_packet.dts = av_rescale_q_rnd(bluray_packet.dts, input->streams[bluray_packet.stream_index]->time_base, mkv_audio[0]->time_base, AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX);
			bluray_packet.duration = av_rescale_q(bluray_packet.duration, input->streams[bluray_packet.stream_index]->time_base, mkv_audio[0]->time_base);

		} else if(input->streams[bluray_packet.stream_index]->codecpar->codec_type == AVMEDIA_TYPE_SUBTITLE) {

			bluray_packet.pts = av_rescale_q_rnd(bluray_packet.pts, input->streams[bluray_packet.stream_index]->time_base, mkv_pgs[0]->time_base, AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX);
			bluray_packet.dts = av_rescale_q_rnd(bluray_packet.dts, input->streams[bluray_packet.stream_index]->time_base, mkv_pgs[0]->time_base, AV_ROUND_NEAR_INF|AV_ROUND_PASS_MINMAX);
			bluray_packet.duration = av_rescale_q(bluray_packet.duration, input->streams[bluray_packet.stream_index]->time_base, mkv_pgs[0]->time_base);

		} else {
			av_packet_unref(&bluray_packet);
			continue;
		}

		// Write the packet to the output filename
		retval = av_interleaved_write_frame(output, &bluray_packet);
		if(retval < 0) {
			av_packet_unref(&bluray_packet);
			goto end;
		}

	}

	end:

	printf("\n");

	// Finish writing to file
	av_write_trailer(output);
	avformat_close_input(&input);
	avio_closep(&output->pb);

	return 0;

}