OpenVulkano/openVulkanoCpp/AR/ArRecorder.cpp

/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#include "ArRecorder.hpp"
#include "ArSession.hpp"
#include "ArFrame.hpp"
#include "IO/Archive/MultiPartArchiveWriter.hpp"
#include "IO/Files/Pfm.hpp"
#include "IO/Files/Pnm.hpp"
#include "IO/AppFolders.hpp"
#include "Base/BlockProfiler.hpp"
#include "Image/YuvUtils.hpp"
#include <fstream>
#include <iostream>
//#if __has_include("turbojpeg.h")
#include <turbojpeg.h>
#define TURBO_JPEG
//#endif

namespace openVulkanoCpp::AR
{
	namespace
	{
		std::filesystem::path GeneratePath(const std::filesystem::path& baseDir, std::string_view name)
		{
			std::stringstream ss;
			auto t = std::time(nullptr);
			auto localTime = *std::localtime(&t);
			ss << std::put_time(&localTime, "%Y-%m-%d_%H_%M_%S");
			return baseDir / name / ss.str();
		}

		std::string GetFileName(size_t frameId, std::string_view fileExtension)
		{
			std::stringstream fnStream;
			fnStream << std::setw(7) << std::setfill('0') << frameId << '.' << fileExtension;
			return fnStream.str();
		}
	}

	ArRecorder::ArRecorder(ArSession* session)
			: m_session(session)
	{
		m_settings.path = GeneratePath(AppFolders::GetAppDataHomeDir(), "ar_recording");
		session->OnNewFrameHighResolution += EventHandler(this, &ArRecorder::SaveHighResolution);
	}

	ArRecorder::~ArRecorder() = default;

	void ArRecorder::WriteColorImage(ArFrame* arFrame, MultiPartArchiveWriter* colorWriter, bool highRes)
	{
		std::string fileName = GetFileName(arFrame->GetFrameId(), "jpg");
#ifndef TURBO_JPEG
		if (arFrame->GetCameraImageAsJpeg([&fileName, this](const char* data, size_t len){ m_colorWriter->AddFile(fileName.c_str(), data, len); }))
			return;
		//TODO stb???
		Logger::AR->error("Failed to create JPEG! Missing turbojpeg.h");
#else
		//TODO handle non nv12 images
		auto img = arFrame->GetCameraImage();
		auto sizeLum = img.luminescenceOrColor.resolution.x * img.luminescenceOrColor.resolution.y;
		auto sizeUV = img.uv.resolution.x * img.uv.resolution.y;
		auto resX = img.luminescenceOrColor.resolution.x;
		auto resY = img.luminescenceOrColor.resolution.y;

		tjhandle handle = tjInitCompress();
		const uint8_t* buffers[3];
		std::unique_ptr<uint8_t[]> dataBuffer;
		if (m_settings.downsampleColor && !highRes)
		{
			dataBuffer = YuvUtils::PlansFromNV12(static_cast<uint8_t*>(img.luminescenceOrColor.data), static_cast<uint8_t*>(img.uv.data),
												 resX, resY, img.uv.resolution.x, img.uv.resolution.y, 2, 2);
			resX /= 2;
			resY /= 2;
			buffers[0] = dataBuffer.get();
			buffers[1] = buffers[0] + sizeLum / 4;
			buffers[2] = buffers[1] + sizeUV / 4;
		}
		else
		{
			dataBuffer = std::unique_ptr<uint8_t[]>(new uint8_t[sizeUV + sizeUV]);
			YuvUtils::ChromaPlanesFromNV12((uint8_t*)img.uv.data, dataBuffer.get(), sizeUV);
			buffers[0] = static_cast<uint8_t*>(img.luminescenceOrColor.data);
			buffers[1] = dataBuffer.get();
			buffers[2] = buffers[1] + sizeUV;
		}

		uint8_t* outBuffer = nullptr;
		unsigned long size = 0;
		if (tjCompressFromYUVPlanes(handle, buffers, resX, nullptr, resY, TJSAMP_420, &outBuffer, &size, 95, TJFLAG_FASTDCT))
			Logger::AR->error("Failed to create JPEG! {}", tjGetErrorStr());
		else
			colorWriter->AddFile(fileName.c_str(), outBuffer, size);
		tjFree(outBuffer);
#endif
	}

	void ArRecorder::WriteDepthImage(ArFrame* arFrame, MultiPartArchiveWriter* depthWriter, MultiPartArchiveWriter* confWriter)
	{
		if (!m_depthWriter || !m_confidenceWriter) return;
		auto depthImg = arFrame->GetDepthImage();
		std::vector<std::pair<const void*, size_t>> buffers(2);
		{ // TODO handle alternative depth formats!!!!
			BlockProfiler profile("Save AR Frame - Depth");
			PfmHeader depthHeader(static_cast<uint32_t>(depthImg.depth.resolution.x), static_cast<uint32_t>(depthImg.depth.resolution.y), 5.0f, false);
			std::string header = depthHeader.ToString();
			buffers[0].first = header.c_str();
			buffers[0].second = header.size();
			buffers[1].first = static_cast<const char *>(depthImg.depth.data);
			buffers[1].second = depthImg.depth.resolution.x * depthImg.depth.resolution.y * sizeof(float);

			const std::string fileName = GetFileName(arFrame->GetFrameId(), "pfm");
			depthWriter->AddFile(fileName.c_str(), buffers);
		}

		{
			BlockProfiler profile("Save AR Frame - Confi");
			PnmHeader confidenceHeader(static_cast<uint32_t>(depthImg.confidence.resolution.x), static_cast<uint32_t>(depthImg.confidence.resolution.y), false, 2);
			std::string header = confidenceHeader.ToString();
			buffers[0].first = header.c_str();
			buffers[0].second = header.size();
			buffers[1].first = static_cast<const char *>(static_cast<void*>(depthImg.confidence.data));
			buffers[1].second = static_cast<size_t>(depthImg.confidence.resolution.x * depthImg.confidence.resolution.y);

			const std::string fileName = GetFileName(arFrame->GetFrameId(), "pgm");
			confWriter->AddFile(fileName.c_str(), buffers);
		}
	}

	void ArRecorder::Write(ArFrame* frame, bool highRes)
	{
		if (frame->IsSaved()) return;
		frame->SetSaved();
		bool useHighResWriter = highRes && m_settings.highResFramesInSeparateArchive;
		BlockProfiler profile("Save AR Frame");
		{
			BlockProfiler profileMeta("Save AR Frame -  Meta");
			std::string metaContent = frame->GetFrameMetadata().ToXML();
			std::string fileName = GetFileName(frame->GetFrameId(), "meta");
			(useHighResWriter ? m_highResWriter : m_metadataWriter)->AddFile(fileName.c_str(), metaContent.c_str(), metaContent.size());
		}
		{
			BlockProfiler profile("Save AR Frame - Image");
			WriteColorImage(frame, useHighResWriter ? m_highResWriter.get() : m_colorWriter.get(), highRes);
		}
		{
			//BlockProfiler profile("Save AR Frame - Depth");
			WriteDepthImage(frame,
							useHighResWriter ? m_highResWriter.get() : m_depthWriter.get(),
							useHighResWriter ? m_highResWriter.get() : m_confidenceWriter.get());
		}
	}

	void ArRecorder::Start()
	{
		if (!m_colorWriter)
		{
			m_colorWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path, "color_{:05d}.tar", ArchiveConfig::TAR, m_settings.archiveSize, true);
			m_depthWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path, "depth_{:05d}.tar", ArchiveConfig::TAR, m_settings.archiveSize, true);
			m_confidenceWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path, "confidence_{:05d}.tar.gz", ArchiveConfig::TAR_GZ, m_settings.archiveSize, true);
			m_metadataWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path, "meta_{:05d}.tar.gz", ArchiveConfig::TAR_GZ, m_settings.archiveSize, true);
			m_highResWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path, "highres_{:05d}.tar", ArchiveConfig::TAR, m_settings.archiveSize, true);

			std::ofstream platformInfoStream(m_settings.path / "ArRecording.xml");
			platformInfoStream << m_session->GetSessionMetadata().ToXML();
			platformInfoStream.close();
		}
		m_recording = true;
	}

	void ArRecorder::Stop()
	{
		if (!m_recording) return;
		m_recording = false;
		for(MultiPartArchiveWriter* writer : { m_colorWriter.get(), m_depthWriter.get(), m_confidenceWriter.get(), m_metadataWriter.get(), m_highResWriter.get() })
		{
			writer->Split();
		}
	}

	void ArRecorder::SetRecordingPath(const std::string& path)
	{
		if (!m_colorWriter)
		{
			for (MultiPartArchiveWriter* writer: { m_colorWriter.get(), m_depthWriter.get(), m_confidenceWriter.get(), m_metadataWriter.get(), m_highResWriter.get() })
			{
				writer->Move(path);
			}
			std::filesystem::rename(m_settings.path / "ArRecording.xml", path + "/ArRecording.xml");
		}

		m_persistent = true;
		m_settings.path = path;
	}

	void ArRecorder::SetRecordingMode(RecordingMode mode)
	{
		if (m_settings.recordingMode == mode) return;
		if (m_settings.recordingMode == RecordingMode::NEW_FRAME && m_newFrameHandler)
		{
			m_session->OnNewFrame -= m_newFrameHandler;
			m_newFrameHandler = nullptr;
		}
		m_settings.recordingMode = mode;
		if (m_settings.recordingMode == RecordingMode::NEW_FRAME)
		{
			m_newFrameHandler = m_session->OnNewFrame += EventHandler(this, &ArRecorder::Save);
		}
	}

	void ArRecorder::Save(const std::shared_ptr<ArFrame>& frame)
	{
		if (!m_recording) return;
		if (m_settings.asyncRecording)
		{

		}
		else Write(frame.get());
	}

	void ArRecorder::SaveHighResolution(const std::shared_ptr<ArFrame>& frame)
	{
		if (!m_recording) return;
		if (m_settings.asyncRecording)
		{

		}
		else Write(frame.get(), true);
	}
}