Fix reconstruction snapshots

Reconstruction snapshots are now fixed. The snapshot is taken from the
next complete IFrame, so reconstruction snapshots cannot be taken while
the reconstruction is paused.

applications/reconstruct/src/main.cpp

@@ -45,6 +45,9 @@
 #include <ftl/cuda/normals.hpp>
 #include <ftl/registration.hpp>
 
+#include <ftl/codecs/h264.hpp>
+#include <ftl/codecs/hevc.hpp>
+
 #include <cuda_profiler_api.h>
 
 #ifdef WIN32
@@ -269,25 +272,55 @@ static void run(ftl::Configurable *root) {
 	std::ofstream fileout;
 	ftl::codecs::Writer writer(fileout);
 
+	std::ofstream snapshotout;
+	ftl::codecs::Writer snapshotwriter(snapshotout);
+
 	root->set("record", false);
 
-	int64_t timestamp = 0;
-	std::vector<std::pair<ftl::codecs::StreamPacket, ftl::codecs::Packet>> currentFrame, completeFrame;
+	int64_t timestamp = -1;
+	bool writingSnapshot = false;
+	std::unordered_set<int64_t> precedingFrames, followingFrames;
 	// Add a recording callback to all reconstruction scenes
 	for (size_t i=0; i<sources.size(); ++i) {
-		sources[i]->addRawCallback([&writer,&groups,&timestamp,&currentFrame,&completeFrame,i](ftl::rgbd::Source *src, const ftl::codecs::StreamPacket &spkt, const ftl::codecs::Packet &pkt) {
+		sources[i]->addRawCallback([&writer,&groups,&snapshotout,&snapshotwriter,&timestamp,&writingSnapshot,&precedingFrames,&followingFrames,i](ftl::rgbd::Source *src, const ftl::codecs::StreamPacket &spkt, const ftl::codecs::Packet &pkt) {
 			ftl::codecs::StreamPacket s = spkt;
 
 			// Patch stream ID to match order in group
 			s.streamID = i;
 			writer.write(s, pkt);
 
-			if (s.timestamp > timestamp) {
+			if (snapshotwriter.active()) {
+				// The frame that is captured is the next IFrame, unless that
+				// IFrame is one of the first two frames seen. In this case a
+				// part of the frame might already have been missed, so the
+				// IFrame after that one is captured instead.
+				if (precedingFrames.size() >= 2) {
+					bool isIFrame = false;
+					switch (pkt.codec) {
+						case ftl::codecs::codec_t::H264:
+							isIFrame = ftl::codecs::h264::isIFrame(pkt.data);
+							break;
+						case ftl::codecs::codec_t::HEVC:
+							isIFrame = ftl::codecs::hevc::isIFrame(pkt.data);
+					}
+
+					if (isIFrame && precedingFrames.count(s.timestamp) == 0) {
 						timestamp = s.timestamp;
-				completeFrame = currentFrame;
-				currentFrame = std::vector<std::pair<ftl::codecs::StreamPacket, ftl::codecs::Packet>>({ std::make_pair(s, pkt) });
-			} else if (s.timestamp == timestamp) {
-				currentFrame.push_back(std::make_pair(s, pkt));
+						writingSnapshot = true;
+						snapshotwriter.write(s, pkt);
+					} else if (writingSnapshot && s.timestamp > timestamp) {
+						followingFrames.insert(s.timestamp);
+					}
+
+					// Keep looking for packets of the captured frame until
+					// packets from two following frames have been seen.
+					if (followingFrames.size() >= 2) {
+						snapshotwriter.end();
+						snapshotout.close();
+					}
+				} else {
+					precedingFrames.insert(s.timestamp);
+				}
 			}
 		});
 	}
@@ -315,33 +348,17 @@ static void run(ftl::Configurable *root) {
 		}
 	});
 
-	std::ofstream snapshotout;
-	ftl::codecs::Writer snapshotwriter(snapshotout);
-
-	root->on("3D-snapshot", [&snapshotout,&snapshotwriter,&sources,&currentFrame](const ftl::config::Event &e) {
+	root->on("3D-snapshot", [&snapshotout,&snapshotwriter,&writingSnapshot,&precedingFrames,&followingFrames](const ftl::config::Event &e) {
+		if (!snapshotwriter.active()) {
 			char timestamp[18];
 			std::time_t t=std::time(NULL);
 			std::strftime(timestamp, sizeof(timestamp), "%F-%H%M%S", std::localtime(&t));
 			snapshotout.open(e.entity->value<std::string>("3D-snapshot", std::string(timestamp) + ".ftl"));
-
+			writingSnapshot = false;
+			precedingFrames.clear();
+			followingFrames.clear();
 			snapshotwriter.begin();
-
-		for (size_t i=0; i<sources.size(); ++i) {
-			auto packet_pair = sources[i]->make_packet(Channel::Calibration, sources[i]->parameters(), Channel::Left, sources[i]->getCapabilities());
-			packet_pair.first.streamID = i;
-			snapshotwriter.write(packet_pair.first, packet_pair.second);
-
-			packet_pair = sources[i]->make_packet(sources[i]->getPose());
-			packet_pair.first.streamID = i;
-			snapshotwriter.write(packet_pair.first, packet_pair.second);
-		}
-
-		for (auto &p : currentFrame) {
-			snapshotwriter.write(p.first, p.second);
 		}
-
-		snapshotwriter.end();
-		snapshotout.close();
 	});
 
 	// -------------------------------------------------------------------------

components/codecs/include/ftl/codecs/writer.hpp

@@ -19,6 +19,7 @@ class Writer {
 	bool begin();
 	bool write(const ftl::codecs::StreamPacket &, const ftl::codecs::Packet &);
 	bool end();
+	bool active();
 
 	private:
 	std::ostream *stream_;

components/codecs/src/writer.cpp

@@ -46,3 +46,7 @@ bool Writer::write(const ftl::codecs::StreamPacket &s, const ftl::codecs::Packet
 	(*stream_).write(buffer.data(), buffer.size());
 	return true;
 }
+
+bool Writer::active() {
+	return active_;
+}
\ No newline at end of file