111 lines
3.2 KiB
C++
111 lines
3.2 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/.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include "IO/Archive/ArchiveReader.hpp"
|
|
#include "IO/Files/Pfm.hpp"
|
|
#include "IO/Files/Pnm.hpp"
|
|
#include <string>
|
|
#include <fstream>
|
|
|
|
namespace OpenVulkano::AR::Playback
|
|
{
|
|
struct DepthImage
|
|
{
|
|
PfmImage depth;
|
|
PnmImage confidence;
|
|
};
|
|
|
|
struct ColorImg
|
|
{
|
|
enum ChannelCount { GRAY = 1, RGBA = 4, YUV = 2, NV12 = -2 };
|
|
|
|
std::unique_ptr<uint8_t[]> dataPtr;
|
|
uint8_t* data = nullptr;
|
|
uint8_t* dataUV = nullptr;
|
|
int cols = 0, rows = 0, channels = 0;
|
|
|
|
void Allocate()
|
|
{
|
|
if (data) throw std::runtime_error("Data already set!");
|
|
dataPtr.reset(new uint8_t[cols * rows * abs(channels)]);
|
|
data = dataPtr.get();
|
|
}
|
|
|
|
bool UseRGB() const { return channels == RGBA; }
|
|
|
|
void Decode(const Array<char>& data);
|
|
};
|
|
|
|
class ArPlaybackReader final
|
|
{
|
|
static constexpr std::string_view TAR_EXTENSIONS_REGEX = R"(\.(tar(\.gz|\.bz2|\.zst)?|tgz|tbz|tb2|tbz2))";
|
|
|
|
ArchiveReader m_archiveMetadata, m_archiveColor, m_archiveDepth, m_archiveConfidence;
|
|
size_t m_imgTotalSize = 0, m_imgReadSize = 0;
|
|
bool m_hasConfidence = false;
|
|
|
|
public:
|
|
ArPlaybackReader(const std::filesystem::path& recDir)
|
|
{
|
|
const std::string extensions = R"((_\d+|\.part\d+)?)" + std::string(TAR_EXTENSIONS_REGEX);
|
|
m_archiveMetadata.Open(recDir, ".*meta(data)?" + extensions);
|
|
m_archiveColor.Open(recDir, ".*(color|image)" + extensions, &m_imgTotalSize);
|
|
m_archiveDepth.Open(recDir, ".*depth" + extensions);
|
|
m_hasConfidence = m_archiveConfidence.Open(recDir, ".*conf(idence)?" + extensions);
|
|
}
|
|
|
|
int GetNextFrameId()
|
|
{
|
|
std::string name = m_archiveMetadata.GetDescription().path;
|
|
return std::stoi(name.substr(0, name.length() - 5));
|
|
}
|
|
|
|
Array<char> ReadMetadata()
|
|
{
|
|
return std::move(m_archiveMetadata.GetNextFile()->second);
|
|
}
|
|
|
|
ColorImg ReadColorImage()
|
|
{
|
|
ColorImg img;
|
|
img.channels = ColorImg::RGBA; // TODO test if defaulting to NV12
|
|
std::optional<std::pair<FileDescription, Array<char>>> file = ReadColorImageRaw();
|
|
img.Decode(file->second);
|
|
m_imgReadSize += 1000 + file.value().first.size;
|
|
return img;
|
|
}
|
|
|
|
std::optional<std::pair<FileDescription, Array<char>>> ReadColorImageRaw() { return m_archiveColor.GetNextFile(); }
|
|
|
|
DepthImage ReadDepthImage()
|
|
{
|
|
DepthImage img;
|
|
m_archiveDepth.GetNextFileAsStream([&img](const FileDescription& desc, std::istream& stream) { img.depth.TryRead(stream); });
|
|
if (!m_archiveConfidence.GetNextFileAsStream([&img](const FileDescription&, std::istream& stream) { img.confidence.TryRead(stream); }))
|
|
{ // No confidence image available
|
|
img.confidence.header.width = img.confidence.header.height = 1;
|
|
img.confidence.image = std::make_unique<char[]>(m_imgTotalSize);
|
|
img.confidence.image[0] = 2; // TODO
|
|
}
|
|
|
|
return img;
|
|
}
|
|
|
|
[[nodiscard]] double GetProgress() const
|
|
{
|
|
return static_cast<double>(m_imgReadSize) / static_cast<double>(m_imgTotalSize);
|
|
}
|
|
|
|
[[nodiscard]] bool HasNext() const
|
|
{
|
|
return m_archiveMetadata.HasNext() && m_archiveDepth.HasNext() && m_archiveColor.HasNext()
|
|
&& (m_archiveConfidence.HasNext() || !m_hasConfidence);
|
|
}
|
|
};
|
|
}
|