430 lines
15 KiB
C++
430 lines
15 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/Archive/ArchiveConfiguration.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 "Image/ExifBuilder.hpp"
|
|
#include "Image/XmpBuilder.hpp"
|
|
#include "Image/JpegWithTagsWriter.hpp"
|
|
#include <array>
|
|
#include <fstream>
|
|
#if __has_include("turbojpeg.h")
|
|
#include <turbojpeg.h>
|
|
#define TURBO_JPEG
|
|
#endif
|
|
|
|
namespace OpenVulkano::AR
|
|
{
|
|
namespace
|
|
{
|
|
constexpr int JPEG_QUALITY_DS = 90;
|
|
constexpr int JPEG_QUALITY_FS = 95;
|
|
|
|
std::filesystem::path GeneratePath(const std::filesystem::path& baseDir, const 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)
|
|
{
|
|
return fmt::format("{:07d}.{}", frameId, fileExtension);
|
|
}
|
|
|
|
std::string MakeXmpTag(ArFrame* arFrame)
|
|
{
|
|
Image::XmpBuilder xmpBuilder;
|
|
xmpBuilder.SetPose(arFrame->GetPose());
|
|
xmpBuilder.SetCreateDateNow();
|
|
xmpBuilder.SetExposureTime(arFrame->GetExposureTime());
|
|
xmpBuilder.SetFocalLength(arFrame->GetFocalLength());
|
|
xmpBuilder.SetCameraIntrinsic(arFrame->GetCameraIntrinsic());
|
|
return xmpBuilder.Finalize();
|
|
}
|
|
|
|
std::vector<uint8_t> MakeExifTag(ArFrame* arFrame)
|
|
{
|
|
Image::ExifBuilder exifBuilder;
|
|
exifBuilder.model = arFrame->GetLensModel();
|
|
exifBuilder.exposureTime = Image::RationalValue(1, 1.0f / arFrame->GetExposureTime());
|
|
exifBuilder.SetOrientation(atan2f(arFrame->GetCameraTransformation()[0][1], arFrame->GetCameraTransformation()[1][1]));
|
|
exifBuilder.dateTaken = exifBuilder.GetCurrentTimestamp();
|
|
exifBuilder.make = arFrame->GetArSession()->GetArType().GetVendorName();
|
|
exifBuilder.SetResolution();
|
|
exifBuilder.fNumber = arFrame->GetFNumber();
|
|
exifBuilder.focalLength = arFrame->GetFocalLength();
|
|
return exifBuilder.Build();
|
|
}
|
|
}
|
|
|
|
ArRecorder::ArRecorder(ArSession* session)
|
|
: m_session(session), m_asyncProcessor(this)
|
|
{
|
|
if (!session) { m_settings.asyncRecording = false; return; }
|
|
m_settings.path = GeneratePath(AppFolders::GetAppDataHomeDir(), "ar_recording");
|
|
session->OnNewFrameHighResolution += EventHandler(this, &ArRecorder::SaveHighResolution);
|
|
}
|
|
|
|
ArRecorder::~ArRecorder()
|
|
{
|
|
if (!m_settings.asyncRecording) WriteMetadataFile();
|
|
else m_asyncProcessor.Close();
|
|
}
|
|
|
|
void ArRecorder::WriteColorImage(ArFrame* arFrame, IArchiveWriter* colorWriter, JpegWithTagsWriter* jpgWriter, bool highRes) const
|
|
{
|
|
//BlockProfiler profile("Save AR Frame - Image");
|
|
#ifndef TURBO_JPEG
|
|
std::string fileName = GetFileName(arFrame->GetFrameId(), "jpg");
|
|
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;
|
|
int jpegQuality = JPEG_QUALITY_FS;
|
|
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, img.luminescenceOrColor.rowPadding, img.uv.rowPadding);
|
|
resX /= 2;
|
|
resY /= 2;
|
|
buffers[0] = dataBuffer.get();
|
|
buffers[1] = buffers[0] + sizeLum / 4;
|
|
buffers[2] = buffers[1] + sizeUV / 4;
|
|
jpegQuality = JPEG_QUALITY_DS; // Use lower quality for downsampled images
|
|
}
|
|
else
|
|
{
|
|
dataBuffer = std::unique_ptr<uint8_t[]>(new uint8_t[sizeUV + sizeUV]);
|
|
YuvUtils::ChromaPlanesFromNV12((uint8_t*)img.uv.data, dataBuffer.get(), img.uv.resolution.x, img.uv.resolution.y, img.uv.rowPadding);
|
|
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, jpegQuality, TJFLAG_FASTDCT)) [[unlikely]]
|
|
Logger::AR->error("Failed to create JPEG! {}", tjGetErrorStr());
|
|
else [[likely]]
|
|
{
|
|
if (colorWriter) [[likely]]
|
|
{
|
|
std::string fileName = GetFileName(arFrame->GetFrameId(), "jpg");
|
|
colorWriter->AddFile(fileName.c_str(), outBuffer, size);
|
|
}
|
|
if (jpgWriter) [[unlikely]]
|
|
jpgWriter->WriteImageData({ outBuffer, size }, false); // Keep open, lifetime is managed outside
|
|
}
|
|
tjFree(outBuffer);
|
|
tjDestroy(handle);
|
|
#endif
|
|
}
|
|
|
|
void ArRecorder::WriteDepthImage(ArFrame* arFrame, IArchiveWriter* depthWriter, IArchiveWriter* confWriter)
|
|
{
|
|
std::string depthName = GetFileName(arFrame->GetFrameId(), "pfm");
|
|
std::string confName = GetFileName(arFrame->GetFrameId(), "pgm");
|
|
WriteDepthImage(arFrame, depthWriter, confWriter, depthName.c_str(), confName.c_str());
|
|
}
|
|
|
|
void ArRecorder::WriteDepthImage(ArFrame* arFrame, IArchiveWriter* depthWriter, IArchiveWriter* confWriter, const char* depthName, const char* confidenceName)
|
|
{
|
|
//BlockProfiler profile("Save AR Frame - Depth");
|
|
if (!depthWriter || !confWriter) return;
|
|
auto depthImg = arFrame->GetDepthImage();
|
|
std::array<std::pair<const void*, size_t>, 2> buffers;
|
|
{ // 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);
|
|
|
|
depthWriter->AddFile(depthName, buffers);
|
|
}
|
|
|
|
if (depthImg.confidence.resolution.x > 1 && depthImg.confidence.resolution.y > 1)
|
|
{
|
|
//BlockProfiler profile("Save AR Frame - Confi");
|
|
PnmHeader confidenceHeader(depthImg.confidence.resolution.x, 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);
|
|
|
|
confWriter->AddFile(confidenceName, buffers);
|
|
}
|
|
}
|
|
|
|
void ArRecorder::WriteMetadata(ArFrame* frame, IArchiveWriter* metaWriter)
|
|
{
|
|
std::string fileName = GetFileName(frame->GetFrameId(), "meta");
|
|
WriteMetadata(frame, metaWriter, fileName.c_str());
|
|
}
|
|
|
|
void ArRecorder::WriteMetadata(ArFrame* frame, IArchiveWriter* metaWriter, const char* fileName)
|
|
{
|
|
//BlockProfiler profileMeta("Save AR Frame - Meta");
|
|
std::string metaContent = frame->GetFrameMetadata().ToYaml();
|
|
metaWriter->AddFile(fileName, 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(), nullptr, highRes);
|
|
WriteDepthImage(frame,
|
|
useHighResWriter ? m_highResWriter.get() : m_depthWriter.get(),
|
|
useHighResWriter ? m_highResWriter.get() : m_confidenceWriter.get());
|
|
m_frameCount++;
|
|
}
|
|
|
|
void ArRecorder::Start()
|
|
{
|
|
if (!m_session) return;
|
|
if (!m_colorWriter)
|
|
{
|
|
m_colorWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path.string(), "color_{:05d}.tar", ArchiveConfig::TAR, m_settings.archiveSize, true);
|
|
m_depthWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path.string(), "depth_{:05d}.tar.zst", ArchiveConfig::TAR_ZSTD, m_settings.archiveSize * 2, true);
|
|
m_confidenceWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path.string(), "confidence_{:05d}.tar.zst", ArchiveConfig::TAR_ZSTD1, m_settings.archiveSize * 10, true);
|
|
m_metadataWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path.string(), "meta_{:05d}.tar.gz", ArchiveConfig::TAR_GZ, m_settings.archiveSize * 10, true);
|
|
m_highResWriter = std::make_unique<MultiPartArchiveWriter>(m_settings.path.string(), "highres_{:05d}.tar", ArchiveConfig::TAR, m_settings.archiveSize, true);
|
|
|
|
WriteMetadataFile();
|
|
}
|
|
m_recording = true;
|
|
OnRecordingStateChanged(this, m_recording);
|
|
m_timer.Start();
|
|
}
|
|
|
|
void ArRecorder::SplitWriters()
|
|
{
|
|
for(MultiPartArchiveWriter* writer : { m_colorWriter.get(), m_depthWriter.get(), m_confidenceWriter.get(), m_metadataWriter.get(), m_highResWriter.get() })
|
|
{
|
|
if (writer) writer->Split();
|
|
}
|
|
}
|
|
|
|
void ArRecorder::Stop()
|
|
{
|
|
if (!m_recording) return;
|
|
m_recording = false;
|
|
OnRecordingStateChanged(this, m_recording);
|
|
if (!m_settings.asyncRecording) SplitWriters();
|
|
m_timer.Tick();
|
|
m_timer.Start();
|
|
}
|
|
|
|
void ArRecorder::SetRecordingPath(const std::filesystem::path& path)
|
|
{
|
|
if (!m_session) return;
|
|
if (!m_colorWriter)
|
|
{
|
|
for (MultiPartArchiveWriter* writer: { m_colorWriter.get(), m_depthWriter.get(), m_confidenceWriter.get(), m_metadataWriter.get(), m_highResWriter.get() })
|
|
{
|
|
if (writer) writer->Move(path);
|
|
}
|
|
if (std::filesystem::exists(m_settings.path / ArSessionMetadata::RECORDING_METADATA_FILENAME))
|
|
std::filesystem::rename(m_settings.path / ArSessionMetadata::RECORDING_METADATA_FILENAME,
|
|
path / ArSessionMetadata::RECORDING_METADATA_FILENAME);
|
|
}
|
|
|
|
m_persistent = true;
|
|
m_settings.path = path;
|
|
}
|
|
|
|
void ArRecorder::SetRecordingMode(RecordingMode mode)
|
|
{
|
|
if (!m_session) return;
|
|
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 || !m_settings.saveHighResFrames) return;
|
|
if (m_settings.asyncRecording)
|
|
{
|
|
m_asyncProcessor.Queue(frame, true);
|
|
}
|
|
else Write(frame.get(), true);
|
|
}
|
|
|
|
void ArRecorder::SaveToFile(const std::shared_ptr<ArFrame>& frame, const std::filesystem::path& path, bool downsample, bool includeAux)
|
|
{
|
|
if (m_settings.asyncRecording) m_asyncProcessor.Queue(frame, path, downsample, includeAux);
|
|
else WriteToFile(frame, path, downsample, includeAux);
|
|
}
|
|
|
|
void ArRecorder::WriteToFile(const std::shared_ptr<ArFrame>& frame, const std::filesystem::path& path, bool downsample, bool includeAux)
|
|
{
|
|
JpegWithTagsWriter jpgWriter(path);
|
|
jpgWriter.WriteExifTag(MakeExifTag(frame.get()));
|
|
jpgWriter.WriteXmpTag(MakeXmpTag(frame.get()));
|
|
WriteColorImage(frame.get(), nullptr, &jpgWriter, !downsample);
|
|
if (includeAux)
|
|
{
|
|
ArchiveWriter writer(jpgWriter.GetFilePtr(), ArchiveConfiguration(ArchiveType::ZIP));
|
|
WriteMetadata(frame.get(), &writer, "metadata.yml");
|
|
if (frame->GetDepthImage().depth.resolution.x > 0 && frame->GetDepthImage().depth.resolution.y > 0)
|
|
WriteDepthImage(frame.get(), &writer, &writer, "depth.pfm", "confidence.pgm");
|
|
}
|
|
}
|
|
|
|
void ArRecorder::WriteMetadataFile()
|
|
{
|
|
if (!m_session) return;
|
|
m_timer.Tick();
|
|
m_session->GetSessionMetadata().recFrameCount = m_frameCount;
|
|
m_session->GetSessionMetadata().recSkippedFrames = m_skippedFrames;
|
|
m_session->GetSessionMetadata().recDuration = static_cast<int32_t>(m_timer.GetTotalSeconds());
|
|
std::ofstream platformInfoStream(m_settings.path / ArSessionMetadata::RECORDING_METADATA_FILENAME);
|
|
platformInfoStream << m_session->GetSessionMetadata().ToYaml();
|
|
platformInfoStream.close();
|
|
}
|
|
|
|
//region AsyncProcessor
|
|
ArRecorder::AsyncProcessor::AsyncProcessor(ArRecorder* recorder)
|
|
: recorder(recorder), processingThread(&ArRecorder::AsyncProcessor::Handler, this)
|
|
{}
|
|
|
|
ArRecorder::AsyncProcessor::~AsyncProcessor()
|
|
{
|
|
Close();
|
|
}
|
|
|
|
void ArRecorder::AsyncProcessor::Close()
|
|
{
|
|
if (requestExit) return;
|
|
requestExit = true;
|
|
newDataAvailable.notify_one();
|
|
if (std::this_thread::get_id() != processingThread.get_id())
|
|
{
|
|
if (processingThread.joinable()) processingThread.join();
|
|
} else processingThread.detach();
|
|
}
|
|
|
|
void ArRecorder::AsyncProcessor::Queue(const std::shared_ptr<ArFrame>& frame, bool highRes)
|
|
{
|
|
if (requestExit || !recorder->m_recording) 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::Queue(const Ptr<ArFrame>& frame, const std::filesystem::path& path, bool downsample, bool aux)
|
|
{
|
|
if (requestExit) return;
|
|
{
|
|
std::unique_lock lock(queueMutex);
|
|
toFile.emplace(frame, path, downsample, aux);
|
|
}
|
|
newDataAvailable.notify_all();
|
|
}
|
|
|
|
void ArRecorder::AsyncProcessor::Handler()
|
|
{
|
|
Utils::SetThreadName("ArRecorder");
|
|
if (!recorder->m_session) return;
|
|
std::unique_lock lock(queueMutex);
|
|
do
|
|
{
|
|
if (Empty()) newDataAvailable.wait(lock, [this]{ return !Empty() || requestExit; });
|
|
while(!toFile.empty())
|
|
{
|
|
auto request = std::move(toFile.front());
|
|
toFile.pop();
|
|
if (!request.frame) continue;
|
|
lock.unlock();
|
|
recorder->WriteToFile(request.frame, request.path, request.downsample, request.addAux);
|
|
lock.lock();
|
|
}
|
|
while(!highResFrameQueue.empty())
|
|
{
|
|
auto frame = std::move(highResFrameQueue.front());
|
|
highResFrameQueue.pop();
|
|
if (!frame || frame->IsSaved()) continue;
|
|
lock.unlock();
|
|
recorder->Write(frame.get(), true);
|
|
lock.lock();
|
|
}
|
|
if (requestExit) break;
|
|
if(!frameQueue.empty())
|
|
{
|
|
if (frameQueue.size() > 3)
|
|
{
|
|
Logger::AR->warn("Falling behind saving frames, skipping ...");
|
|
recorder->m_skippedFrames++;
|
|
frameQueue.pop();
|
|
//while(frameQueue.size() > 3) frameQueue.pop();
|
|
}
|
|
auto frame = std::move(frameQueue.front());
|
|
frameQueue.pop();
|
|
if (!frame || frame->IsSaved()) continue;
|
|
lock.unlock();
|
|
recorder->Write(frame.get(), false);
|
|
lock.lock();
|
|
}
|
|
if (!recorder->m_recording) recorder->SplitWriters();
|
|
}
|
|
while (!requestExit);
|
|
recorder->WriteMetadataFile();
|
|
}
|
|
//endregion
|
|
}
|