path: root/src/argaze/utils/contexts/OpenCV.py

"""Define OpenCV window display context"""

"""
This program is free software: you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by the Free Software
Foundation, either version 3 of the License, or (at your option) any later
version.
This program is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
You should have received a copy of the GNU General Public License along with
this program. If not, see <https://www.gnu.org/licenses/>.
"""

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

import logging
import threading
import time

import cv2

from argaze import ArFeatures, DataFeatures


class Cursor(ArFeatures.ArContext):
	"""Capture cursor position over OpenCV window.

	!!! warning
	    It is assumed that an OpenCV window with the same name than the context is used to display context's pipeline image.
	"""

	@DataFeatures.PipelineStepInit
	def __init__(self, **kwargs):

		# Init DataCaptureContext class
		super().__init__()

	@DataFeatures.PipelineStepEnter
	def __enter__(self):

		logging.info('OpenCV.Cursor context starts...')

		# Create a window
		cv2.namedWindow(self.name, cv2.WINDOW_AUTOSIZE)

		# Init timestamp
		self._start_time = time.time()

		# Attach mouse event callback to window
		cv2.setMouseCallback(self.name, self.__on_mouse_event)
		
		return self

	@DataFeatures.PipelineStepExit
	def __exit__(self, exception_type, exception_value, exception_traceback):

		logging.info('OpenCV.Cursor context stops...')

		# Delete window
		cv2.destroyAllWindows()

	def __on_mouse_event(self, event, x, y, flags, param):
		"""Process pointer position."""

		logging.debug('OpenCV.Cursor.on_mouse_event %i %i', x, y)

		if not self.is_paused():

			# Process timestamped gaze position
			self._process_gaze_position(timestamp = int((time.time() - self._start_time) * 1e3), x = x, y = y)


class Movie(Cursor, ArFeatures.DataPlaybackContext):
	"""Playback movie images and capture cursor position over OpenCV window.

	!!! warning
	    It is assumed that an OpenCV window with the same name than the context is used to display context's pipeline image.
	"""
	@DataFeatures.PipelineStepInit
	def __init__(self, **kwargs):

		# Init Cursor class
		super().__init__()

		# Init private attributes
		self.__path = None
		self.__movie = None
		self.__movie_fps = None
		self.__movie_width = None
		self.__movie_height = None
		self.__movie_length = None

		self.__current_image_index = None
		self.__next_image_index = None
		self.__refresh = False

	@property
	def path(self) -> str:
		"""Movie file path."""
		return self.__path

	@path.setter
	def path(self, path: str):

		self.__path = path 

		# Load movie
		self.__movie = cv2.VideoCapture(self.__path)
		self.__movie_fps = self.__movie.get(cv2.CAP_PROP_FPS)
		self.__movie_width = int(self.__movie.get(cv2.CAP_PROP_FRAME_WIDTH))
		self.__movie_height = int(self.__movie.get(cv2.CAP_PROP_FRAME_HEIGHT))
		self.__movie_length = self.__movie.get(cv2.CAP_PROP_FRAME_COUNT)

	@DataFeatures.PipelineStepEnter
	def __enter__(self):

		logging.info('OpenCV.Movie context starts...')

		# Enter in Cursor context
		super().__enter__()

		# Open reading thread
		self.__reading_thread = threading.Thread(target=self.__read)

		logging.debug('> starting reading thread...')
		self.__reading_thread.start()
		
		return self

	def __read(self):
		"""Iterate on movie images."""

		# Init image selection
		_, current_image = self.__movie.read()
		current_image_time = self.__movie.get(cv2.CAP_PROP_POS_MSEC)
		self.__next_image_index = 0 #int(self.__start * self.__movie_fps)

		while self.is_running():

			# Check pause event (and stop event)
			while self.is_paused() and self.is_running():

				logging.debug('> reading is paused at %i', current_image_time)

				time.sleep(1)

			# Select a new image and detect markers once
			if self.__next_image_index != self.__current_image_index or self.__refresh:

				self.__movie.set(cv2.CAP_PROP_POS_FRAMES, self.__next_image_index)

				success, image = self.__movie.read()

				if success:

					# Refresh once
					self.__refresh = False

					self.__current_image_index = self.__movie.get(cv2.CAP_PROP_POS_FRAMES) - 1
					current_image_time = self.__movie.get(cv2.CAP_PROP_POS_MSEC)

					# Timestamp image
					image = DataFeatures.TimestampedImage(image, timestamp=current_image_time)

					# Process movie image
					self._process_camera_image(timestamp=current_image_time, image=image)

			# Wait
			time.sleep(1 / self.__movie_fps)

	@DataFeatures.PipelineStepExit
	def __exit__(self, exception_type, exception_value, exception_traceback):

		logging.info('OpenCV.Movie context stops...')

		# Exit from Cursor context
		super().__exit__(exception_type, exception_value, exception_traceback)

		# Close data capture
		self.stop()

		# Stop reading thread
		threading.Thread.join(self.__reading_thread)

	def refresh(self):
		"""Refresh current frame."""
		self.__refresh = True

	def previous(self):
		"""Go to previous frame."""
		self.__next_image_index -= 1

		# Clip image index
		if self.__next_image_index < 0:
			self.__next_image_index = 0

	def next(self):
		"""Go to next frame."""

		self.__next_image_index += 1

		# Clip image index
		if self.__next_image_index < 0:
			self.__next_image_index = 0

	@property
	def duration(self) -> int|float:
		"""Get movie duration."""

		return self.__movie_length / self.__movie_fps

	@property
	def progression(self) -> float:
		"""Get movie playback progression between 0 and 1."""

		if self.__current_image_index is not None:

			return self.__current_image_index / self.__movie_length

		else:

			return 0.

class Camera(Cursor, ArFeatures.DataCaptureContext):
	"""Capture camera images and capture cursor position over OpenCV window.

	!!! warning
	    It is assumed that an OpenCV window with the same name than the context is used to display context's pipeline image.
	"""
	@DataFeatures.PipelineStepInit
	def __init__(self, **kwargs):

		# Init Cursor class
		super().__init__()

		# Init private attributes
		self.__camera_id = None
		self.__camera = None
		self.__video_fps = None
		self.__video_width = None
		self.__video_height = None

	@property
	def identifier(self) -> int:
		"""Camera device id."""
		return self.__camera_id

	@identifier.setter
	def identifier(self, camera_id: int):

		self.__camera_id = camera_id

		# Load movie
		self.__camera = cv2.VideoCapture(self.__camera_id)
		self.__video_fps = self.__camera.get(cv2.CAP_PROP_FPS)
		self.__video_width = int(self.__camera.get(cv2.CAP_PROP_FRAME_WIDTH))
		self.__video_height = int(self.__camera.get(cv2.CAP_PROP_FRAME_HEIGHT))

	@DataFeatures.PipelineStepEnter
	def __enter__(self):

		logging.info('OpenCV.Movie context starts...')

		# Enter in Cursor context
		super().__enter__()

		# Open reading thread
		self.__reading_thread = threading.Thread(target=self.__read)

		logging.debug('> starting reading thread...')
		self.__reading_thread.start()
		
		return self

	def __read(self):
		"""Iterate on camera images."""

		while self.is_running():

			# Check pause event (and stop event)
			while self.is_paused() and self.is_running():

				logging.debug('> reading is paused at %i', current_image_time)

				time.sleep(1)

			# Select a new image
			success, image = self.__camera.read()
			image_time = self.__camera.get(cv2.CAP_PROP_POS_MSEC)

			if success:

				# Timestamp image
				image = DataFeatures.TimestampedImage(image, timestamp=image_time)

				# Process movie image
				self._process_camera_image(timestamp=image_time, image=image)

			# Wait
			time.sleep(1 / self.__video_fps)

	@DataFeatures.PipelineStepExit
	def __exit__(self, exception_type, exception_value, exception_traceback):

		logging.info('OpenCV.Movie context stops...')

		# Exit from Cursor context
		super().__exit__(exception_type, exception_value, exception_traceback)

		# Close data capture
		self.stop()

		# Stop reading thread
		threading.Thread.join(self.__reading_thread)