X-Git-Url: https://git.sesse.net/?a=blobdiff_plain;f=futatabi%2Fmain.cpp;h=19cd5fb3e4162439e22e7cac2cd8208b3a84686e;hb=02ea864dc5a6dde7450c497581ff18d784ab832c;hp=a4b1e422937887652b5821000f824cc446537f9b;hpb=18199cc2f756bd8eb17a475882dbc738194cbee8;p=nageru

diff --git a/futatabi/main.cpp b/futatabi/main.cpp
index a4b1e42..19cd5fb 100644
--- a/futatabi/main.cpp
+++ b/futatabi/main.cpp
@@ -75,7 +75,7 @@ Summary metric_received_frame_size_bytes;
 
 namespace {
 
-FrameOnDisk write_frame(int stream_idx, int64_t pts, const uint8_t *data, size_t size, DB *db)
+FrameOnDisk write_frame(int stream_idx, int64_t pts, const uint8_t *data, size_t size, vector<uint32_t> audio, DB *db)
 {
 	if (open_frame_files.count(stream_idx) == 0) {
 		char filename[256];
@@ -105,6 +105,7 @@ FrameOnDisk write_frame(int stream_idx, int64_t pts, const uint8_t *data, size_t
 	hdr.set_stream_idx(stream_idx);
 	hdr.set_pts(pts);
 	hdr.set_file_size(size);
+	hdr.set_audio_size(audio.size() * sizeof(audio[0]));
 
 	string serialized;
 	if (!hdr.SerializeToString(&serialized)) {
@@ -130,6 +131,12 @@ FrameOnDisk write_frame(int stream_idx, int64_t pts, const uint8_t *data, size_t
 		perror("fwrite");
 		abort();
 	}
+	if (audio.size() > 0) {
+		if (fwrite(audio.data(), hdr.audio_size(), 1, file.fp) != 1) {
+			perror("fwrite");
+			exit(1);
+		}
+	}
 	fflush(file.fp);  // No fsync(), though. We can accept losing a few frames.
 	global_disk_space_estimator->report_write(filename, 8 + sizeof(len) + serialized.size() + size, pts);
 
@@ -138,6 +145,7 @@ FrameOnDisk write_frame(int stream_idx, int64_t pts, const uint8_t *data, size_t
 	frame.filename_idx = filename_idx;
 	frame.offset = offset;
 	frame.size = size;
+	frame.audio_size = audio.size() * sizeof(audio[0]);
 
 	{
 		lock_guard<mutex> lock(frame_mu);
@@ -371,9 +379,10 @@ void load_frame_file(const char *filename, const string &basename, unsigned file
 		}
 		frame.filename_idx = filename_idx;
 		frame.size = hdr.file_size();
+		frame.audio_size = hdr.audio_size();
 
-		if (frame.offset + frame.size > file_len ||
-		    fseek(fp, frame.offset + frame.size, SEEK_SET) == -1) {
+		if (frame.offset + frame.size + frame.audio_size > file_len ||
+		    fseek(fp, frame.offset + frame.size + frame.audio_size, SEEK_SET) == -1) {
 			fprintf(stderr, "WARNING: %s: Could not seek past frame (probably truncated).\n", filename);
 			break;
 		}
@@ -500,8 +509,22 @@ void record_thread_func()
 			continue;
 		}
 
-		int64_t last_pts = -1;
+		// Match any audio streams to video streams, sequentially.
+		vector<int> video_stream_idx, audio_stream_idx;
+		for (unsigned i = 0; i < format_ctx->nb_streams; ++i) {
+			if (format_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
+				video_stream_idx.push_back(i);
+			} else if (format_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
+				audio_stream_idx.push_back(i);
+			}
+		}
+		unordered_map<int, int> audio_stream_to_video_stream_idx;
+		for (size_t i = 0; i < min(video_stream_idx.size(), audio_stream_idx.size()); ++i) {
+			audio_stream_to_video_stream_idx[audio_stream_idx[i]] = video_stream_idx[i];
+		}
 
+		vector<uint32_t> pending_audio[MAX_STREAMS];
+		int64_t last_pts = -1;
 		while (!should_quit.load()) {
 			AVPacket pkt;
 			unique_ptr<AVPacket, decltype(av_packet_unref) *> pkt_cleanup(
@@ -515,8 +538,23 @@ void record_thread_func()
 			if (av_read_frame(format_ctx.get(), &pkt) != 0) {
 				break;
 			}
+
+			AVStream *stream = format_ctx->streams[pkt.stream_index];
+			if (stream->codecpar->codec_type == AVMEDIA_TYPE_AUDIO &&
+			    audio_stream_to_video_stream_idx.count(pkt.stream_index)) {
+				if ((pkt.size % (sizeof(uint32_t) * 2)) != 0) {
+					fprintf(stderr, "Audio stream %u had a packet of strange length %d, ignoring.\n",
+						pkt.stream_index, pkt.size);
+				} else {
+					// TODO: Endianness?
+					const uint32_t *begin = (const uint32_t *)pkt.data;
+					const uint32_t *end = (const uint32_t *)(pkt.data + pkt.size);
+					pending_audio[audio_stream_to_video_stream_idx[pkt.stream_index]].assign(begin, end);
+				}
+			}
+
 			if (pkt.stream_index >= MAX_STREAMS ||
-			    format_ctx->streams[pkt.stream_index]->codecpar->codec_type != AVMEDIA_TYPE_VIDEO) {
+			    stream->codecpar->codec_type != AVMEDIA_TYPE_VIDEO) {
 				continue;
 			}
 
@@ -524,7 +562,7 @@ void record_thread_func()
 			metric_received_frame_size_bytes.count_event(pkt.size);
 
 			// Convert pts to our own timebase.
-			AVRational stream_timebase = format_ctx->streams[pkt.stream_index]->time_base;
+			AVRational stream_timebase = stream->time_base;
 			int64_t pts = av_rescale_q(pkt.pts, stream_timebase, AVRational{ 1, TIMEBASE });
 
 			// Translate offset into our stream.
@@ -535,7 +573,7 @@ void record_thread_func()
 
 			//fprintf(stderr, "Got a frame from camera %d, pts = %ld, size = %d\n",
 			//      pkt.stream_index, pts, pkt.size);
-			FrameOnDisk frame = write_frame(pkt.stream_index, pts, pkt.data, pkt.size, &db);
+			FrameOnDisk frame = write_frame(pkt.stream_index, pts, pkt.data, pkt.size, move(pending_audio[pkt.stream_index]), &db);
 
 			post_to_main_thread([pkt, frame] {
 				global_mainwindow->display_frame(pkt.stream_index, frame);