Expand image decoding to allow for YUV and NV12

openVulkanoCpp/AR/Provider/Playback/ArPlaybackReader.cpp

@@ -6,14 +6,18 @@
 
 #include "ArPlaybackReader.hpp"
 
+#include "Base/BlockProfiler.hpp"
+#include "Base/Wrapper.hpp"
+#include "Image/YuvUtils.hpp"
+
 #if __has_include("turbojpeg.h")
 #include <turbojpeg.h>
 
 namespace OpenVulkano::AR::Playback
 {
-	ColorImg ArPlaybackReader::ReadColorImage()
+	void ArPlaybackReader::ReadColorImage(ColorImg& img)
 	{
-		ColorImg img;
+		BlockProfiler profiler("Load jpeg");
 		auto file = m_archiveColor.GetNextFile();
 
 		long unsigned int jpegSize = file->second.Size();
@@ -23,21 +27,50 @@ namespace OpenVulkano::AR::Playback
 		tjhandle jpegDecompressor = tjInitDecompress();
 		tjDecompressHeader2(jpegDecompressor, compressedImage, jpegSize, &img.cols, &img.rows, &jpegSubsamp);
 
-		img.channels = 4;
+		if (img.UseRGB())
-		img.dataPtr = std::shared_ptr<uint8_t>(new uint8_t[img.cols * img.rows * img.channels]);
+		{
-		img.data = img.dataPtr.get();
+			img.channels = 4;
+			if (!img.data) img.Allocate();
 
-		//TODO is it better to not map to rgb? to keep the same pipeline as on device
+			tjDecompress2(jpegDecompressor, compressedImage, jpegSize, img.data, img.cols, 0/*pitch*/, img.rows, TJPF_RGBA, TJFLAG_FASTDCT);
-		tjDecompress2(jpegDecompressor, compressedImage, jpegSize, img.data, img.cols, 0/*pitch*/, img.rows, TJPF_RGBA, TJFLAG_FASTDCT);
+		}
+		else if (img.channels == 1)
+		{
+			//TODO handle grayscale
+		}
+		else if (img.channels != 0)
+		{
+			if (!img.data) img.Allocate();
+
+			std::array<uint8_t*, 3> planes;
+			size_t chromaSize = img.cols * img.rows / 4;
+			planes[0] = img.data;
+			if (!img.dataUV) img.dataUV = img.data + (img.cols * img.rows);
+			Unique<uint8_t[]> chromaTmp;
+			if (img.channels == ColorImg::NV12 && !img.dataPtr)
+			{
+				chromaTmp = std::make_unique<uint8_t[]>(img.cols * img.rows / 2);
+				planes[1] = chromaTmp.get();
+			}
+			else
+			{
+				planes[1] = img.dataUV;
+				if (img.channels == ColorImg::NV12) planes[1] += chromaSize;
+			}
+			planes[2] = planes[1] + chromaSize;
+			std::array<int, 3> rowSize = { img.cols, img.cols / 2, img.cols / 2 };
+
+			tjDecompressToYUVPlanes(jpegDecompressor, compressedImage, jpegSize, planes.data(),
+				img.cols, rowSize.data(), img.rows, TJFLAG_FASTDCT);
+
+			if (img.channels == ColorImg::NV12)
+			{
+				YuvUtils::NV12FromChromaPlanes(planes[1], img.data + (img.cols * img.rows), chromaSize);
+			}
+		}
 		tjDestroy(jpegDecompressor);
 
-		//auto buff = new uint8_t[img.cols * img.rows * 3];
-
-		//YuvUtils::NV12FromChromaPlanes(buff, img.data + (img.cols * img.rows), img.cols * img.rows / 4);
-
 		m_imgReadSize += 1000 + file.value().first.size;
-
-		return img;
 	}
 }
 

openVulkanoCpp/AR/Provider/Playback/ArPlaybackReader.hpp

@@ -22,9 +22,21 @@ namespace OpenVulkano::AR::Playback
 
 	struct ColorImg
 	{
-		std::shared_ptr<uint8_t> dataPtr;
+		enum ChannelCount { GRAY = 1, RGBA = 4, YUV = 2, NV12 = -2 };
-		uint8_t* data;
+
-		int cols, rows, channels;
+		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; }
 	};
 
 	class ArPlaybackReader final
@@ -55,7 +67,15 @@ namespace OpenVulkano::AR::Playback
 			return std::move(m_archiveMetadata.GetNextFile()->second);
 		}
 
-		ColorImg ReadColorImage();
+		ColorImg ReadColorImage()
+		{
+			ColorImg img;
+			img.channels = ColorImg::RGBA; // TODO test if defaulting to NV12
+			ReadColorImage(img);
+			return img;
+		}
+
+		void ReadColorImage(ColorImg& img);
 
 		DepthImage ReadDepthImage()
 		{