diff --git a/ar_tag_pose/detector_board_main.py b/ar_tag_pose/detector_board_main.py index 3c11e18..41121d5 100644 --- a/ar_tag_pose/detector_board_main.py +++ b/ar_tag_pose/detector_board_main.py @@ -10,6 +10,7 @@ from ar_tag_pose.utils import to_tf, to_euler, draw_text from detector_board import BoardDetector from detector_board_aruco import BoardDetectorAruco from detector_board_charuco import BoardDetectorCharuco +from cv2_enumerate_cameras import enumerate_cameras RESULTS_FOLDER = "results" @@ -19,6 +20,7 @@ class Detector(abc.ABC): self.detector = detector self.out_file = out_file self.rotation_order = rotation_order + self.camera_name = "UnknownCamera" self.device_id = None self.src_video = None self.src_images = None @@ -37,13 +39,30 @@ class Detector(abc.ABC): self.src_images = f"{os.path.dirname(source)}/{name}_%0{len(numeric_part)}d.{suffix}" Path.mkdir(Path(RESULTS_FOLDER), exist_ok=True) - self.camera_name = "YourCameraNameHere" + + self.cap = None + if self.device_id is not None: + # Open camera + self.cap = cv.VideoCapture(self.device_id) + self.cap_init() + elif self.src_video is not None: + # Open video file + self.cap = cv.VideoCapture(self.src_video) + elif self.src_images is not None: + # Open image file sequence + self.cap = cv.VideoCapture(self.src_images, cv.CAP_IMAGES) + def params_save(self): with open(f"{RESULTS_FOLDER}/{self.params["name"]}.json", 'w') as fp: json.dump(self.params, fp, indent=4) - def cap_init(self, cap: cv.VideoCapture): + def cap_init(self): + backends = {'V4L2': cv.CAP_V4L} + cam_dict = enumerate_cameras(backends[self.cap.getBackendName()]) + self.camera_name = cam_dict[self.device_id].name + print(self.camera_name) + camera_params_default = { cv.CAP_PROP_FRAME_WIDTH: 1920 * 2, cv.CAP_PROP_FRAME_HEIGHT: 1080 * 2, @@ -60,7 +79,7 @@ class Detector(abc.ABC): key = item value = camera[key] print(f"{item}: {value}") - cap.set(key, value) + self.cap.set(key, value) self.params['camera'] = {self.camera_name: camera} self.params_save() @@ -98,6 +117,7 @@ class Detector(abc.ABC): calib = {self.camera_name: {str(det): {'mtx': mtx.tolist(), 'dist': dist.tolist()}}} det.calibration_set(mtx, dist, None, None) self.params['calib'] = calib + self.params_save() print(f"Loaded {filename} successfully!") except FileNotFoundError: pass @@ -105,18 +125,7 @@ class Detector(abc.ABC): pass def process(self): - cap = None - if self.device_id is not None: - # Open camera - cap = cv.VideoCapture(self.device_id) - self.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) - + cap = self.cap if cap is None or not cap.isOpened(): print("Cannot open media") exit() diff --git a/pyproject.toml b/pyproject.toml index 491e81a..2af8a44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,8 @@ dependencies = [ "opencv-python", "opencv-contrib-python", "scipy", - "reedsolo" + "reedsolo", + "cv2_enumerate_cameras" ] readme = "README.md"