Z3D/source/assets/data_loaders/kitti_loader.cpp

#include "assets/data_loaders/kitti_loader.hpp"

#include "glm/glm.hpp"

#include <charconv>
#include <fstream>
#include <glm/ext/matrix_transform.hpp>

#include "assets/components/point_cloud_vertex_components.hpp"
#include "util/binary_ifstream.hpp"
#include "util/logger.hpp"


ztu::result<std::string_view> kitti_loader::parent_directory(const std::string_view path)
{
	const auto sep_index = path.rfind(std::filesystem::path::preferred_separator);

	if (sep_index == std::string_view::npos)
	{
		return std::unexpected(std::make_error_code(std::errc::no_such_file_or_directory));
	}

	return path.substr(0, sep_index);
};

std::error_code kitti_loader::prefetch(
	const file_dir_list& paths,
	prefetch_queue& queue
) {

	// Directories can simply be passed on
	queue.kitti_pose_queue.directories.push_back(paths.directories);

	// For files, we just forward the files directory
	for (const auto file : queue.kitti_pose_queue.files)
	{
		if (const auto base_directory = parent_directory(file).and_then(parent_directory))
		{
			queue.kitti_pose_queue.directories.push_back(*base_directory);
		}
		else
		{
			// TODO remove from list
			ztu::logger::error("Malformed kitti file path: %.", file);
		}
	}

	return {};
}

std::error_code kitti_loader::load(
   dynamic_point_cloud_buffer& buffer,
   const file_dir_list& paths,
   prefetch_lookup& id_lookup,
   dynamic_data_store& store,
   bool
) {
	namespace fs = std::filesystem;
	std::error_code error;

	std::vector<dynamic_pose_store::iterator_type> pose_its;
	pose_its.reserve(paths.files.size());

	auto processed_filenames = ztu::string_list{};

	auto path_buffer = fs::path{};

	const auto preprocess_filename = [&](
		std::string_view path,
		const auto& directory,
		std::string_view filename,
		const pose_prefetch_lookup::directory_iterator& dir_it
	) {
		const auto pose_index = frame_id_from_filename(filename);
		if (not pose_index) [[unlikely]]
		{
			ztu::logger::error("Could not parse frame id from kitti file path: %.", filename);
			return;
		}

		const auto [ index_it, pose_id_match ] =  id_lookup.poses.find_index(dir_it, *pose_index);

		if (not pose_id_match) [[unlikely]]
		{
			ztu::logger::error("No matching pose index (%) found in directory (%).", directory, *pose_index);
			return;
		}

		const auto [ pose_it, pose_match ] = store.poses.find(pose_id_match);
		if (not pose_id_match) [[unlikely]]
		{
			ztu::logger::error("No matching pose found for id: %.", pose_id_match);
			return;
		}

		processed_filenames.push_back(path);
		pose_its.push_back(pose_it);
	};

	for (const auto file : paths.files)
	{
		path_buffer.assign(file.begin(), file.end());

		if (not fs::is_regular_file(path_buffer))
		{
			ztu::logger::error("Given kitti file does not exist: %.", path_buffer);
			continue;
		}

		const auto sep_index = file.rfind(fs::path::preferred_separator);

		if (sep_index == std::string_view::npos) [[unlikely]]
		{
			ztu::logger::error("Could not parse frame directory from kitti file path: %.", file);
			continue;
		}

		const auto base_directory = parent_directory(file.substr(0, sep_index));
		if (not base_directory) [[unlikely]]
		{
			ztu::logger::error("Could not parse base directory from kitti file path: %.", file);
			continue;
		}

		const auto [ dir_it, dir_match ] = id_lookup.poses.find_directory(*base_directory);
		if (not dir_match) [[unlikely]]
		{
			ztu::logger::error("No matching pose directory found for %.", file);
			continue;
		}

		const auto filename = file.substr(sep_index + 1);

		preprocess_filename(
		file,
			*base_directory,
			filename,
			dir_it
		);
	}

	for (const auto directory : paths.directories)
	{
		path_buffer.assign(directory.begin(), directory.end());

		const auto [ dir_it, dir_match ] = id_lookup.poses.find_directory(path_buffer);
		if (not dir_match) [[unlikely]]
		{
			ztu::logger::error("No matching pose directory found for %.", path_buffer);
			continue;
		}

		path_buffer /= frame_folder;

		if (not fs::is_directory(path_buffer))
		{
			ztu::logger::error("Given kitti directory does not exist: %.", directory);
			continue;
		}

		for (const auto& file : fs::directory_iterator{ path_buffer })
		{
			const auto file_path = std::string_view{ file.path().c_str() };

			const auto extension_begin = file_path.rfind('.');

			if (extension_begin == std::string_view::npos or file_path.substr(extension_begin) != ".bin")
			{
				continue;
			}

			auto filename_begin = file_path.rfind(fs::path::preferred_separator, extension_begin);
			filename_begin = filename_begin == std::string_view::npos ? 0 : filename_begin + 1;

			const auto filename_only = file_path.substr(filename_begin);

			const auto pose_index = frame_id_from_filename(filename_only);
			if (not pose_index) [[unlikely]]
			{
				ztu::logger::error("Could not parse frame id from kitti filename: %.", filename_only);
				continue;
			}

			const auto [ index_it, pose_id_match ] =  id_lookup.poses.find_index(dir_it, *pose_index);

			if (not pose_id_match) [[unlikely]]
			{
				ztu::logger::error("No matching pose index (%) found in directory (%).", directory, *pose_index);
				continue;
			}

			const auto [ pose_it, pose_match ] = store.poses.find(pose_id_match);
			if (not pose_id_match) [[unlikely]]
			{
				ztu::logger::error("No matching pose found for id: %.", pose_id_match);
				continue;
			}

			processed_filenames.push_back(file_path);
			pose_its.push_back(pose_it);

			preprocess_filename(
				file_path,
				directory,
				filename_only,
				dir_it
			);
		}
	}


	for (const auto [ filename, pose_it ] : std::ranges::views::zip_view(processed_filenames, pose_its))
	{
		buffer.clear();

		if ((error = load_point_file(filename, buffer)))
		{
			return error;
		}

		transform_point_cloud(buffer.positions(), *pose_it);

		store.point_clouds.add(buffer);
	}

	return {};
}

void kitti_loader::transform_point_cloud(
	std::span<components::point_cloud_vertex::position> points,
	const glm::mat4& pose
) {
	for (auto& [ x, y, z ] : points) {
		auto vec = glm::vec4{ x, y, z, 1.0f };
		vec = pose * vec;
		x = vec.x;
		y = vec.y;
		z = vec.z;
	}
}

std::error_code kitti_loader::load_point_file(
	const std::filesystem::path& filename,
	dynamic_point_cloud_buffer& point_cloud
) {

	auto in = binary_ifstream{};

	auto error = std::error_code{};

	if ((error == in.open(filename, true)))
	{
		return error;
	}

	const auto read_vector = [&in](auto& vector) -> std::error_code
	{
		for (auto& component : vector)
		{
			float component32;
			if (const auto e = in.read_ieee754<std::endian::little>(component32))
			{
				return e;
			}
			component = component32;
		}
		return {};
	};

	components::point_cloud_vertex::position position;

	auto& positions = point_cloud.positions();

	while (not ((error = read_vector(position)))) {
		positions.push_back(position);
		if ((error = in.skip<float>())) // TODO what am I skipping here?!?
		{
			break;
		}
	}

	if (static_cast<std::errc>(error.value()) != std::errc::result_out_of_range)
	{
		return error;
	}

	return {};
}

ztu::result<std::size_t> kitti_loader::frame_id_from_filename(
	std::string_view filename
) {
	std::size_t id;
	const auto result = std::from_chars(filename.cbegin(), filename.cend(), id);

	if (result.ec != std::errc{})
	{
		return std::unexpected(std::make_error_code(result.ec));
	}

	return id;
}