Helpers.hpp

#pragma once

#include <limits>
#include <cstdint>
#include <vector>
#include <algorithm>
#include <stdexcept>

struct Annotation
{
public:
	int m_X;
	int m_Y;


	float x;
	float y;
	float z;
	Annotation(void)
		: m_X(-1)
		, m_Y(-1)
	{

	};

	Annotation(int X, int Y)
		: m_X(X)
		, m_Y(Y)
	{

	};
};

std::vector<std::string> &split(const std::string &s, char delim, std::vector<std::string> &elems) {
	std::stringstream ss(s);
	std::string item;
	while (std::getline(ss, item, delim)) {
		elems.push_back(item);
	}
	return elems;
}

std::vector<std::string> split(const std::string &s, char delim) {
	std::vector<std::string> elems;
	split(s, delim, elems);
	return elems;
}

struct AnnotationList
{
public:
	std::vector<Annotation> m_AnnoList;
	int m_Size;
	bool m_readDepthImage;

	AnnotationList(void)
		: m_Size(0)
	{
		m_readDepthImage = true;
	};

	AnnotationList(int Size)
	{
		m_Size = Size;

		if (m_Size < 1)
			return;

		for (int i = 0; i < m_Size; ++i)
			m_AnnoList.push_back(Annotation());

		m_readDepthImage = true;
	}

	// For printing/writing to file
	friend std::ostream& operator<<(std::ostream &s, const AnnotationList &p)
	{
		for (int i = 0; i < p.m_AnnoList.size(); ++i)
			s << p.m_AnnoList[i].m_X << ", " << p.m_AnnoList[i].m_Y << "; ";

		s << "\n";
		return s;
	};

	// For reading from file
	void load(std::string& Line)
	{
		// Tokenize by ;
		std::vector<std::string> FrameAnnos = split(Line, ';');
		if (FrameAnnos.size() == 0)
			return;
		for (int i = 0; i < FrameAnnos.size() - 1; ++i) // There is a stray ; in the end so -1
		{
			// std::cout << "FrameAnnos[i]: " << FrameAnnos[i] << std::endl;
			// Tokenize by ,
			std::vector<std::string> Elems = split(FrameAnnos[i], ',');
			if (Elems.size() != 2)
				throw std::runtime_error("Something went wrong reading annotations.");
			m_AnnoList.push_back(Annotation(std::atoi(Elems[0].c_str()), std::atoi(Elems[1].c_str())));
			// std::cout << m_AnnoList[i].m_X << ", " << m_AnnoList[i].m_Y << "\n";
		}
	}
};

std::string g_BaseDataDir;
QStringList g_ImageList;
QStringList g_ImageListColor;
std::vector<AnnotationList> g_ImageAnnoLists; // One list per image
AnnotationList * g_CurrImageAnnoList = NULL;
Annotation g_LastAnno;
std::fstream g_OutFile;
std::fstream g_OutFile3D;

std::vector<float> intrinsics;

int g_FileCount;
int g_CurrFileCtr = 0;
cv::Mat g_CurrImage, g_CurrDispImage, g_CurrColorImage;
std::string g_MainWindTitle = "AnnoTool";
std::string g_ColorWindTitle = "ColorImage";
int g_ABWidth = 16;
int g_MaxAnnoSize = 8; // 5 fingers + 3 cuboid

inline cv::Vec3b HSV2BGR(const cv::Vec3b& hsv)
{
	// This is from the OpenCV source code directly
	// opencv/3rdparty/openexr/Imath/ImathColorAlgo.cpp
	float hue = float(hsv[0]) * 0.005555; // Div by 180.0 NOT 255.0 because of OpenCV convention
	float sat = float(hsv[1]) * 0.003921; // Div by 255.0
	float val = float(hsv[2]) * 0.003921;

	float x = 0.0, y = 0.0, z = 0.0;

	if (hue == 1) hue = 0;
	else hue *= 6;

	int i = int(floor(hue));
	float f = hue - i;
	float p = val*(1 - sat);
	float q = val*(1 - (sat*f));
	float t = val*(1 - (sat*(1 - f)));

	switch (i)
	{
	case 0: x = val; y = t; z = p; break;
	case 1: x = q; y = val; z = p; break;
	case 2: x = p; y = val; z = t; break;
	case 3: x = p; y = q; z = val; break;
	case 4: x = t; y = p; z = val; break;
	case 5: x = val; y = p; z = q; break;
	}

	return cv::Vec3b(uint8_t(z * 255),
		uint8_t(y * 255),
		uint8_t(x * 255));
};

inline cv::Vec3b GetDepth2ColorMap(int DepthVal, int zNear = 50, int zFar = 3000) // Default near/far are in mm
{
	// OPTIMIZE: This can be made a lookup if runtime is an issue
	// Assuming all 3 arguments are in mm
	// We divide the total depth variation in bins of size 256
	// Then for each bin we assign a "major" color
	// The output is a color gradient for the input DepthVal (within the color of it's bin)
	// We do the computation in HSV space but convert the image later

	cv::Vec3b HSV(64, 64, 64); // Pre-computed in BGR for grayscale to save time
	if (DepthVal > zFar || DepthVal < zNear)
		return HSV;

	float HueGap = 180.0f;
	float HueBinSz = float(zFar - zNear) / HueGap;
	int HueBin = floor(float(DepthVal) / HueBinSz); // Zero indexed bin ID
	float BinWeight = (float(DepthVal) / HueBinSz) - HueBin; // How bright is the depth in this bin

	// We change only hue (depth bin) and saturation (brightness within bin)
	// OpenCV: H: 0 - 180, S: 0 - 255, V: 0 - 255
	uint8_t Offset = 240;
	HSV = cv::Vec3b(uint8_t(HueBin * (180.0 / HueGap)), Offset + uint8_t(BinWeight * (255.0f - Offset)), 255); // We provide a constant offset for saturation
	return HSV2BGR(HSV);
};

void ReadAnnotations(void)
{
	std::string line;
	int Ctr = 0;
	for (int i = 0; i < g_FileCount; ++i)
	{
		std::getline(g_OutFile, line);
		// std::cout << line << std::endl;
		g_ImageAnnoLists[i].load(line);
		Ctr++;
	}

	std::cout << "[ INFO ]: Loaded annotations from: " << Ctr << " lines." << std::endl;
}

void WriteAnnotations(void)
{
	/*for (int i = 0; i < g_ImageAnnoLists.size(); ++i)
	{
	g_OutFile << g_ImageAnnoLists[i];
	g_ImageAnnoLists[i].m_changedSinceLastWrite = false;
	}*/
	AnnotationList emptyLine(g_MaxAnnoSize);
	for (int i = 0; i < g_MaxAnnoSize; ++i)
	{
		emptyLine.m_AnnoList[i].m_X = -1;
		emptyLine.m_AnnoList[i].m_Y = -1;
		
	}

	// Back project points to 3D and write to file
	for (int i = 0; i < g_ImageAnnoLists.size(); ++i)
	{
	
		if (g_ImageAnnoLists[i].m_AnnoList.size() == 0)
		{
			//g_OutFile3D << "\n";
			g_OutFile << emptyLine;

			for (int i = 0; i < g_MaxAnnoSize; ++i)
			{
				g_OutFile3D << 0 << ", " << 0 << ", " << 0 << "; ";

			}
			
			g_OutFile3D << "\n";
			continue;
		
		}

		//write 2d annotations
		g_OutFile << g_ImageAnnoLists[i];

		if (g_ImageAnnoLists[i].m_readDepthImage)
		{
			// Load appropriate image
			cv::Mat LocImage = cv::imread(g_BaseDataDir + "//depth//" + QDir::separator().toAscii() + g_ImageList[i].toUtf8().constData(), -1);
			if (g_ImageAnnoLists[i].m_AnnoList.size() == 0)
				for (int i = 0; i < g_MaxAnnoSize;++i)
				{
					g_OutFile3D << 0 << ", " << 0 << ", " << 0 << "; ";
					
				}

			for (int kk = 0; kk < g_ImageAnnoLists[i].m_AnnoList.size(); ++kk)
			{
				

				int c = g_ImageAnnoLists[i].m_AnnoList[kk].m_X;
				int r = g_ImageAnnoLists[i].m_AnnoList[kk].m_Y;
				int Depth = LocImage.at<int16_t>(r, c);

				cv::Mat Pt2D_Hom = cv::Mat::ones(3, 1, CV_32FC1);
				Pt2D_Hom.at<float>(0, 0) = float(c);
				Pt2D_Hom.at<float>(1, 0) = float(r);
				cv::Mat Pt3D(3, 1, CV_32FC1);
				cv::Mat Intrinsics = cv::Mat::eye(3, 3, CV_32FC1);
				Intrinsics.at<float>(0, 0) = intrinsics[0];
				Intrinsics.at<float>(1, 1) = intrinsics[1];
				Intrinsics.at<float>(0, 2) = intrinsics[2];
				Intrinsics.at<float>(1, 2) = intrinsics[3];


				//std::cout << Intrinsics;

				Pt3D = Intrinsics.inv() * Pt2D_Hom;

				// Normalize
				Pt3D.at<float>(0, 0) *= (float(Depth) / Pt3D.at<float>(2, 0)); // *-1.0; // Not sure why the -1 is needed
				Pt3D.at<float>(1, 0) *= (float(Depth) / Pt3D.at<float>(2, 0)); // *-1.0;
				Pt3D.at<float>(2, 0) *= (float(Depth) / Pt3D.at<float>(2, 0));


				g_ImageAnnoLists[i].m_AnnoList[kk].x = Pt3D.at<float>(0, 0);
				g_ImageAnnoLists[i].m_AnnoList[kk].y = Pt3D.at<float>(1, 0);
				g_ImageAnnoLists[i].m_AnnoList[kk].z = Pt3D.at<float>(2, 0);

				g_OutFile3D << Pt3D.at<float>(0, 0) << ", " << Pt3D.at<float>(1, 0) << ", " << Pt3D.at<float>(2, 0) << "; ";

				g_ImageAnnoLists[i].m_readDepthImage = false;
				//std::cout << "File: " << g_BaseDataDir+ "//depth//" + QDir::separator().toAscii() + g_ImageList[i].toUtf8().constData() << std::endl;
				//  std::cout << "[ x, y ]: " << Pt2D_Hom << std::endl; // c << ", " << r << std::endl;
				//   std::cout << "Depth: " << Depth << std::endl;
				// std::cout << Intrinsics.inv() << std::endl;
				// std::cout << Pt3D.at<float>(0, 0) << ", " << Pt3D.at<float>(1, 0) << ", " << Pt3D.at<float>(2, 0) << std::endl;
				// std::cin.get();

				// // Check by projecting back
				// cv::Mat Pt2D_Hom_Check = cv::Mat::ones(3, 1, CV_32FC1);
				// Pt2D_Hom_Check = Intrinsics.inv() * Pt3D;
				// cv::Point ProjPt2D = 

				// std::cout << "Clicked point: " << Pt2D_Hom << std::endl;
				// std::cout << "Projected point: " << Pt2D_Hom << std::endl;
			}

		}
		else
			for (int kk = 0; kk < g_ImageAnnoLists[i].m_AnnoList.size(); ++kk)
			{

				g_OutFile3D << g_ImageAnnoLists[i].m_AnnoList[kk].x << ", " << g_ImageAnnoLists[i].m_AnnoList[kk].y << ", " << g_ImageAnnoLists[i].m_AnnoList[kk].z << "; ";
			}


		g_OutFile3D << "\n";
	}
	g_OutFile << "\n";
	g_OutFile3D << "\n";

}

void AnnotateImage(int event, int x, int y, int, void*)
{
	if (event == CV_EVENT_LBUTTONDOWN)
		g_LastAnno = Annotation(x, y);
	// else if(event == CV_EVENT_RBUTTONDOWN)
	// 	g_LastAnno = Annotation();
}

void DrawBox(Annotation& Anno, cv::Scalar Color = cv::Scalar(0.0, 0.0, 255.0), int OptionalDisplayNumber = -1)
{
	cv::Rect Box = cv::Rect(Anno.m_X - (g_ABWidth / 2), Anno.m_Y - (g_ABWidth / 2), g_ABWidth, g_ABWidth);
	cv::rectangle(g_CurrDispImage, Box, Color, 1);
	cv::circle(g_CurrDispImage, cv::Point(Anno.m_X, Anno.m_Y), 1, Color, -1);

	if (OptionalDisplayNumber > 0)
	{
		// Display frame count
		std::stringstream ss;
		ss << OptionalDisplayNumber;
		cv::putText(g_CurrDispImage, ss.str(), cv::Point(Anno.m_X + 10, Anno.m_Y + 10), cv::FONT_HERSHEY_COMPLEX_SMALL, 0.5, Color, 1, CV_AA);
	}
}

void MakeDisplayImage(const cv::Mat& Depth16U, cv::Mat& DisplayImage)
{
	cv::Mat Depth8U = cv::Mat(Depth16U.size(), CV_8UC1);
	DisplayImage = cv::Mat(Depth16U.size(), CV_8UC3);
	// NOTE: Don't parallelize this! It slows down for some reason
	int16_t * DepthFramePtr = (int16_t *)Depth16U.data;
	uint8_t * DepthFrame8UPtr = (uint8_t *)Depth8U.data;
	uint8_t * DepthFrame8UC3Ptr = (uint8_t *)DisplayImage.data;
	for (int i = 0; i < Depth16U.rows; i++)
	{
		for (int j = 0; j < Depth16U.cols; j++)
		{
			int LocDatCtr = i*Depth16U.cols + j;
			int16_t ImVal = DepthFramePtr[LocDatCtr];

			bool isBG = false;
			if (ImVal > 31999)
				isBG = true;

			float zNear = 100; // In mm
			// 8UC1 images lack bit depth so 600 is tops for good close range visualization
			// Use the 8UC3 color images for visualization
			float zFar = 800;
			float FImVal = float(ImVal);
			if (ImVal > zFar) // NOTE: It's ImVal not FImVal
				FImVal = 255.0;
			else if (ImVal < zNear)
				FImVal = 255.0;
			else
				FImVal = ((FImVal - zNear) / (zFar - zNear)) * 255.0;
			DepthFrame8UPtr[LocDatCtr] = uint8_t(FImVal);  // Invalid points are white. Near is black, far is white-ish

			// // This is the raw point cloud in mm
			// Vector3VP Pt3D = BackProject(Vector2VP(j, i), float(ImVal));

			cv::Vec3b PxCol = GetDepth2ColorMap(int(ImVal), 100, 800); // Near/Far parameters are arbitrary but don't matter too much
			if (isBG)
			{
				DepthFrame8UC3Ptr[LocDatCtr * 3 + 0] = 0;
				DepthFrame8UC3Ptr[LocDatCtr * 3 + 1] = 0;
				DepthFrame8UC3Ptr[LocDatCtr * 3 + 2] = 0;
			}
			else
			{
				DepthFrame8UC3Ptr[LocDatCtr * 3 + 0] = PxCol[0];
				DepthFrame8UC3Ptr[LocDatCtr * 3 + 1] = PxCol[1];
				DepthFrame8UC3Ptr[LocDatCtr * 3 + 2] = PxCol[2];
			}
		}
	}

	// Testing edge detection
	cv::Mat Edges, EdgesColor, Depth8UCol;
	// cv::cvtColor(Depth8U, Depth8UCol, CV_GRAY2BGR);
	cv::Canny(DisplayImage, Edges, 280, 390, 3);
	cv::cvtColor(Edges, EdgesColor, CV_GRAY2BGR);
	// cv::addWeighted(Depth8UCol/*DisplayImage*/, 0.8, EdgesColor, 0.2, 0, DisplayImage);
	cv::addWeighted(DisplayImage, 0.6, EdgesColor, 0.4, 0, DisplayImage);

	// Display frame count
	std::stringstream ss;
	ss << g_CurrFileCtr + 1 << " / " << g_FileCount;
	cv::putText(g_CurrColorImage, ss.str(), cvPoint(440, 460), cv::FONT_HERSHEY_COMPLEX_SMALL, 1.5, cvScalar(20, 20, 255), 1, CV_AA);
}