297 lines
10 KiB
C++
297 lines
10 KiB
C++
/*
|
|
* 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 "Base/Logger.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
|
|
{
|
|
constexpr std::string_view RECORDING_METADATA_FILENAME = "ArRecording.xml";
|
|
|
|
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_asyncProcessor(this)
|
|
{
|
|
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) const
|
|
{
|
|
BlockProfiler profile("Save AR Frame - Image");
|
|
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)
|
|
{
|
|
//BlockProfiler profile("Save AR Frame - Depth");
|
|
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::WriteMetadata(ArFrame* frame, MultiPartArchiveWriter* metaWriter)
|
|
{
|
|
BlockProfiler profileMeta("Save AR Frame - Meta");
|
|
std::string metaContent = frame->GetFrameMetadata().ToXML();
|
|
std::string fileName = GetFileName(frame->GetFrameId(), "meta");
|
|
metaWriter->AddFile(fileName.c_str(), metaContent.c_str(), metaContent.size());
|
|
}
|
|
|
|
void ArRecorder::Write(ArFrame* frame, bool highRes)
|
|
{
|
|
if (frame->IsSaved()) return;
|
|
frame->SetSaved();
|
|
bool useHighResWriter = highRes && m_settings.highResFramesInSeparateArchive;
|
|
BlockProfiler profile("Save AR Frame");
|
|
WriteMetadata(frame, useHighResWriter ? m_highResWriter.get() : m_metadataWriter.get());
|
|
WriteColorImage(frame, useHighResWriter ? m_highResWriter.get() : m_colorWriter.get(), highRes);
|
|
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 / RECORDING_METADATA_FILENAME);
|
|
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)
|
|
{
|
|
std::filesystem::path p(path);
|
|
for (MultiPartArchiveWriter* writer: { m_colorWriter.get(), m_depthWriter.get(), m_confidenceWriter.get(), m_metadataWriter.get(), m_highResWriter.get() })
|
|
{
|
|
if (writer) writer->Move(p);
|
|
}
|
|
if (std::filesystem::exists(m_settings.path / RECORDING_METADATA_FILENAME))
|
|
std::filesystem::rename(m_settings.path / RECORDING_METADATA_FILENAME, p / RECORDING_METADATA_FILENAME);
|
|
}
|
|
|
|
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)
|
|
{
|
|
m_asyncProcessor.Queue(frame, false);
|
|
}
|
|
else Write(frame.get());
|
|
}
|
|
|
|
void ArRecorder::SaveHighResolution(const std::shared_ptr<ArFrame>& frame)
|
|
{
|
|
if (!m_recording) return;
|
|
if (m_settings.asyncRecording)
|
|
{
|
|
m_asyncProcessor.Queue(frame, true);
|
|
}
|
|
else Write(frame.get(), true);
|
|
}
|
|
|
|
//region AsyncProcessor
|
|
ArRecorder::AsyncProcessor::AsyncProcessor(ArRecorder* recorder)
|
|
: recorder(recorder), processingThread(&ArRecorder::AsyncProcessor::Handler, this)
|
|
{}
|
|
|
|
ArRecorder::AsyncProcessor::~AsyncProcessor()
|
|
{
|
|
requestExit = true;
|
|
newDataAvailable.notify_one();
|
|
if (processingThread.joinable()) processingThread.join();
|
|
}
|
|
|
|
void ArRecorder::AsyncProcessor::Queue(const std::shared_ptr<ArFrame>& frame, bool highRes)
|
|
{
|
|
if (requestExit) return; // no need to queue up on shutdown
|
|
{
|
|
std::unique_lock lock(queueMutex);
|
|
if (highRes) highResFrameQueue.push(frame);
|
|
else frameQueue.push(frame);
|
|
}
|
|
newDataAvailable.notify_all();
|
|
}
|
|
|
|
void ArRecorder::AsyncProcessor::Handler()
|
|
{
|
|
Utils::SetThreadName("ArRecorder");
|
|
std::unique_lock lock(queueMutex);
|
|
do
|
|
{
|
|
newDataAvailable.wait(lock, [this]{ return !frameQueue.empty() || !highResFrameQueue.empty() || requestExit; });
|
|
while(!highResFrameQueue.empty())
|
|
{
|
|
auto frame = std::move(highResFrameQueue.front());
|
|
highResFrameQueue.pop();
|
|
if (frame->IsSaved()) continue;
|
|
lock.unlock();
|
|
recorder->Write(frame.get(), true);
|
|
lock.lock();
|
|
}
|
|
if (requestExit) break;
|
|
while(!frameQueue.empty())
|
|
{
|
|
if (frameQueue.size() > 3)
|
|
{
|
|
Logger::AR->warn("Falling behind saving frames, skipping ...");
|
|
//while(frameQueue.size() > 3) frameQueue.pop();
|
|
}
|
|
auto frame = std::move(frameQueue.front());
|
|
frameQueue.pop();
|
|
if (frame->IsSaved()) continue;
|
|
lock.unlock();
|
|
recorder->Write(frame.get(), false);
|
|
lock.lock();
|
|
}
|
|
}
|
|
while (!requestExit);
|
|
}
|
|
//endregion
|
|
}
|