path: root/src/argaze/AreaOfInterest/AOI3DScene.py

#!/usr/bin/env python

""" """

__author__ = "Théo de la Hogue"
__credits__ = []
__copyright__ = "Copyright 2023, Ecole Nationale de l'Aviation Civile (ENAC)"
__license__ = "BSD"

from typing import TypeVar, Tuple
import math
import re

from argaze import DataStructures
from argaze.AreaOfInterest import AOIFeatures, AOI2DScene

import numpy
import cv2 as cv

T0 = numpy.array([0., 0., 0.])
"""Define no translation vector."""

R0 = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])
"""Define no rotation matrix."""

K0 = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 0.]])
"""Define default camera intrinsic parameters matrix."""

D0 = numpy.array([0.0, 0.0, 0.0, 0.0, 0.0])
"""Define default camera distorsion coefficients vector."""

AOI3DSceneType = TypeVar('AOI3DScene', bound="AOI3DScene")
# Type definition for type annotation convenience

AOI2DSceneType = TypeVar('AOI2DScene', bound="AOI2DScene")
# Type definition for type annotation convenience

class AOI3DScene(AOIFeatures.AOIScene):
	"""Define AOI 3D scene."""

	def __init__(self, aois_3d = None):

		super().__init__(3, aois_3d)

	@classmethod
	def from_obj(self, obj_filepath: str) -> AOI3DSceneType:
		"""Load AOI3D scene from .obj file."""

		aois_3d = {}

		# regex rules for .obj file parsing
		OBJ_RX_DICT = {
			'comment': re.compile(r'#(.*)\n'),
			'name': re.compile(r'o (\w+)(.*)\n'),
			'vertice': re.compile(r'v ([+-]?[0-9]*[.]?[0-9]+) ([+-]?[0-9]*[.]?[0-9]+) ([+-]?[0-9]*[.]?[0-9]+)\n'),
			'face': re.compile(r'f (.*)\n')
		}

        # regex .obj line parser
		def __parse_obj_line(line):

			for key, rx in OBJ_RX_DICT.items():
				match = rx.search(line)
				if match:
					return key, match

			# if there are no matches
			return None, None
        
        # start parsing
		try:

			name = None
			vertices = []
			faces = {}

			# open the file and read through it line by line
			with open(obj_filepath, 'r') as file:

				line = file.readline()

				while line:

					# at each line check for a match with a regex
					key, match = __parse_obj_line(line)

					# extract comment
					if key == 'comment':
						pass

					# extract aoi3D name
					elif key == 'name':

						name = str(match.group(1))

					# fill vertices array
					elif key == 'vertice':

						vertices.append(tuple([float(match.group(1)), float(match.group(2)), float(match.group(3))]))

					# extract aoi3D vertice id
					elif key == 'face':

						faces[name] = [int(i) for i in match.group(1).split()]

					# go to next line
					line = file.readline()

				file.close()

				# retreive all aoi3D vertices
				for name, face in faces.items():
					aoi3D = AOIFeatures.AreaOfInterest([ vertices[i-1] for i in face ])
					aois_3d[name] = aoi3D

		except IOError:
			raise IOError(f'File not found: {obj_filepath}')

		return AOI3DScene(aois_3d)

	def to_obj(self, obj_filepath: str):
		"""Save AOI3D scene into .obj file."""

		with open(obj_filepath, 'w', encoding='utf-8') as file:

			file.write('# ArGaze OBJ file\n')

			vertices_count = 0

			for name, aoi3D in self.items():

				file.write(f'o {name}\n') 

				vertices_ids = 'f'

				for vertices in aoi3D:

					vertices_coords = 'v'

					for coord in vertices:

						vertices_coords += f' {coord:.6f}'

					file.write(vertices_coords + '\n')

					vertices_count += 1
					vertices_ids += f' {vertices_count}'

				file.write('s off\n')
				file.write(vertices_ids + '\n')

	def vision_cone(self, cone_radius, cone_height, cone_tip=[0., 0., 0.], cone_direction=[0., 0., 1.]) -> Tuple[AOI3DSceneType, AOI3DSceneType]:
		"""Get AOI which are inside and out a given cone field.

		.. note:: By default
		   The cone have its tip at origin and its base oriented to positive Z axis.

		* **Returns:** 
			- scene inside of the cone
			- scene outside of the cone
		"""

		# define cone tip and direction as numpy array
		cone_tip = numpy.array(cone_tip).astype(numpy.float32)
		cone_direction = numpy.array(cone_direction).astype(numpy.float32)

		# sort aoi
		aoi3D_scene_inside = AOI3DScene()
		aoi3D_scene_outside = AOI3DScene()

		for name, aoi3D in self.items():

			one_vertice_out = False
			for vertices in aoi3D:

				distance = numpy.dot(vertices - cone_tip, cone_direction)
				radius = (distance / cone_height) * cone_radius
				ortho_distance = numpy.linalg.norm((vertices - cone_tip) - distance * cone_direction)

				if ortho_distance > radius:
					one_vertice_out = True
					break

			# if no vertice is outside the cone, aoi is inside
			if not one_vertice_out:
				aoi3D_scene_inside[name] = aoi3D
			else:
				aoi3D_scene_outside[name] = aoi3D

		return aoi3D_scene_inside, aoi3D_scene_outside

	def project(self, T: numpy.array = T0, R: numpy.array = R0, K: numpy.array = K0, D: numpy.array = D0) -> AOI2DSceneType:
		"""Project 3D scene onto 2D scene according translation, rotation and optical parameters.

		* **Arguments:** 
			- translation vector
			- rotation vector
			- camera intrinsic parameters matrix
			- camera distorsion coefficients vector

		.. danger::
		   Camera distorsion coefficients could projects points which are far from the frame into it.

		.. note::
		   As gaze is mainly focusing on the frame center, where the distorsion is low, it could be acceptable to not use camera distorsion.
		"""

		aoi2D_scene = AOI2DScene.AOI2DScene()

		for name, aoi3D in self.items():

			vertices_2D, J = cv.projectPoints(aoi3D.astype(numpy.float32), R, T, numpy.array(K), numpy.array(D))
			
			aoi2D = vertices_2D.reshape((len(vertices_2D), 2)).view(AOIFeatures.AreaOfInterest)

			aoi2D_scene[name] = aoi2D

		return aoi2D_scene

	def transform(self, T: numpy.array = T0, R: numpy.array = R0) -> AOI3DSceneType:
		"""Translate and/or rotate 3D scene.		

		* **Arguments:** 
			- translation vector
			- rotation matrix
		"""

		aoi3D_scene = AOI3DScene()

		for name, aoi3D in self.items():

			aoi3D_scene[name] = aoi3D.dot(R.T) + T

		return aoi3D_scene