handle images and movies

ar_tag_pose/board_pose.py

@@ -1,9 +1,11 @@
+import os
 import abc
 import math
 import cv2 as cv
 import numpy as np
 from pathlib import Path
 from ar_tag_pose.utils import to_board_pose_euler
+import mimetypes
 
 RESULTS_FOLDER = "results"
 
@@ -18,8 +20,24 @@ class BoardPose(abc.ABC):
 	obj_points_list = []
 	img_points_list = []
 
-	def __init__(self, device_id: int):
-		self.device_id = device_id
+	def __init__(self, source: str):
+		self.device_id = None
+		self.src_video = None
+		self.src_images = None
+		if source.isnumeric():
+			self.device_id = int(source)
+		elif isinstance(source, str):
+			mime_type = mimetypes.guess_type(source)[0]
+			if "video" in mime_type:
+				self.src_video = source
+			if "image" in mime_type:
+				self.src_images = source
+				basename = os.path.basename(source)
+				name = basename.split('_')[0]
+				numeric_part = basename.split('_')[-1].split('.')[0]
+				suffix = basename.split('_')[-1].split('.')[1]
+				self.src_images = f"{os.path.dirname(source)}/{name}_%0{len(numeric_part)}d.{suffix}"
+
 		Path.mkdir(Path(RESULTS_FOLDER), exist_ok=True)
 
 	@staticmethod
@@ -62,7 +80,7 @@ class BoardPose(abc.ABC):
 															  rvecs=r_vecs, tvecs=t_vecs, flags=flags)
 
 			print(f"Camera calibration finished: ", end='')
-			if ret < 2:
+			if ret < 3:
 				print(f"Success ({ret})!")
 				success = True
 				self.mtx = mtx
@@ -75,29 +93,52 @@ class BoardPose(abc.ABC):
 		return success
 
 	def process(self):
+		cap = None
+		if self.device_id is not None:
+			# Open camera
+			cap = cv.VideoCapture(self.device_id)
+			BoardPose.cap_init(cap)
+		elif  self.src_video is not None:
+			# Open video file
+			cap = cv.VideoCapture(self.src_video)
+		elif self.src_images is not None:
+			# Open image file sequence
+			cap = cv.VideoCapture(self.src_images, cv.CAP_IMAGES)
 
-		# Open camera
-		cap = cv.VideoCapture(self.device_id)
-		BoardPose.cap_init(cap)
-
-		if not cap.isOpened():
-			print("Cannot open camera")
+		if cap is None or not cap.isOpened():
+			print("Cannot open media")
 			exit()
 
 		cal_img_size = None
-		last_rotation = [0, 0, 0]
-		while True:
+		frame_go = True
+		frame_num = 0
+		while cap.isOpened():
+			if self.src_images:
+				cap.set(cv.CAP_PROP_POS_FRAMES, frame_num)
 			# Read a new frame
 			ok, img = cap.read()
 			if not ok:
-				return
+				break
+
 			k = cv.waitKey(1) & 0xFF
 			if k == ord('q'):
 				break
+			elif k == ord('g'):
+				frame_go = True
+				frame_num = 0
 			elif k == ord('s'):
+				frame_go = False
+				frame_num = 0
+			elif k == ord('-'):
+				frame_go = False
+				frame_num = frame_num - 1
+			elif k == ord('+'):
+				frame_go = False
+				frame_num = frame_num + 1
+			elif k == ord('p'):
 				success, cal_img_size, obj_points, img_points = self.calibrate_points(img)
 				if success:
-					print("Frame captured")
+					print("Points captured")
 					self.obj_points_list.append(obj_points)
 					self.img_points_list.append(img_points)
 				else:
@@ -111,20 +152,21 @@ class BoardPose(abc.ABC):
 				self.dist = None
 				self.obj_points_list = []
 				self.img_points_list = []
-			else:
-				success, r_vec, t_vec = self._process(img)
-				if success:
-					r_vec_obj, _ = to_board_pose_euler(r_vec, t_vec)
-					rot = [round(180.0 / math.pi * v, 0) for v in r_vec_obj]
-					if rot != last_rotation:
-						print("\r                                                                                 ",
-							  end='')
-						print(f"\rTilt: {rot[0]:.1f}, Roll: {rot[1]:.1f}, Azimuth: {rot[2]:.1f}", end='')
-					last_rotation = rot
+
+			success, r_vec, t_vec = self._process(img)
+			if success:
+				r_vec_obj, _ = to_board_pose_euler(r_vec, t_vec)
+				rot = [round(180.0 / math.pi * v, 1) for v in r_vec_obj]
+				print("\r                                                                                 ", end='')
+				print(f"\rTilt: {rot[0]:.1f}, Roll: {rot[1]:.1f}, Azimuth: {rot[2]:.1f}", end='')
+
 			img_scaled = cv.resize(img, dsize=None, fx=0.5, fy=0.5)
 			cv.imshow('img', img_scaled)
-		cv.destroyAllWindows()
 
+			if frame_go:
+				frame_num += 1
+
+		cv.destroyAllWindows()
 
 	@abc.abstractmethod
 	def get_name(self) -> str:

ar_tag_pose/board_pose_aruco.py

@@ -15,7 +15,7 @@ def factory(dict_type: int, board_size: cv.typing.Size, marker_length: float=100
 	return detector, board, name
 
 class BoardPoseAruco(BoardPose):
-	def __init__(self, device_id: int, dict_type: int, board_size: tuple[int, int]):
+	def __init__(self, device_id: str, dict_type: int, board_size: tuple[int, int]):
 		BoardPose.__init__(self, device_id)
 		self.board_size = board_size
 		self.detector, self.board, self.name = factory(dict_type, board_size)
@@ -53,8 +53,6 @@ class BoardPoseAruco(BoardPose):
 		pose = False
 		r_vec = None
 		t_vec = None
-		if not self.is_calibrated():
-			return pose, None, None
 
 		gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
 		(corners, ids, rejected) = self.detector.detectMarkers(gray)
@@ -67,6 +65,9 @@ class BoardPoseAruco(BoardPose):
 		if draw_marker_id:
 			cv.aruco.drawDetectedMarkers(image, corners, ids)
 
+		if not self.is_calibrated():
+			return pose, None, None
+
 		# Draw marker axis
 		if draw_marker_axis:
 			_r_vecs, _t_vecs, _ = cv.aruco.estimatePoseSingleMarkers(corners, markerLength=0.1, cameraMatrix=self.mtx, distCoeffs=self.dist)
@@ -87,12 +88,12 @@ class BoardPoseAruco(BoardPose):
 def main():
 	# construct the argument parser and parse the arguments
 	ap = argparse.ArgumentParser()
-	ap.add_argument("-i", "--device_id", type=int,
-		default=0,
-		help="Camera device index ")
+	ap.add_argument("-s", "--source", type=str,
+		default='0',
+		help="Camera device index, movie file or first file of image sequence")
 	args = vars(ap.parse_args())
 
-	pose = BoardPoseAruco(args["device_id"], 10, (5, 5))
+	pose = BoardPoseAruco(args["source"], 10, (5, 5))
 	pose.board_image_generate(margin_length=10)
 	pose.calibrate_load()
 	pose.process()

ar_tag_pose/board_pose_charuco.py

@@ -15,7 +15,7 @@ def factory(dict_type: int, board_size: cv.typing.Size, square_length: float = 5
 	return detector, board, name
 
 class BoardPoseCharuco(BoardPose):
-	def __init__(self, device_id: int, dict_type: int, board_size: tuple[int, int]):
+	def __init__(self, device_id: str, dict_type: int, board_size: tuple[int, int]):
 		BoardPose.__init__(self, device_id)
 		self.board_size = board_size
 		self.detector, self.board, self.name = factory(dict_type, board_size)
@@ -85,13 +85,13 @@ class BoardPoseCharuco(BoardPose):
 def main():
 	# construct the argument parser and parse the arguments
 	ap = argparse.ArgumentParser()
-	ap.add_argument("-i", "--device_id", type=int,
-		default=0,
-		help="Camera device index ")
+	ap.add_argument("-s", "--source", type=str,
+		default='0',
+		help="Camera device index, movie file or first file of image sequence")
 	args = vars(ap.parse_args())
 
-	pose = BoardPoseCharuco(args["device_id"], 10, (6, 7))
-	pose.board_image_generate()
+	pose = BoardPoseCharuco(args["source"], 10, (6, 7))
+	pose.board_image_generate(margin_length=10)
 	pose.calibrate_load()
 	pose.process()