diff --git a/ar_tag_pose/detector_board_main.py b/ar_tag_pose/detector_board_main.py index 6dd95b5..2a074dc 100644 --- a/ar_tag_pose/detector_board_main.py +++ b/ar_tag_pose/detector_board_main.py @@ -100,6 +100,7 @@ class Detector(abc.ABC): frame_go = True frame_num = 0 + out = None while cap.isOpened(): if self.src_images: cap.set(cv.CAP_PROP_POS_FRAMES, frame_num) @@ -151,12 +152,31 @@ class Detector(abc.ABC): img_scaled = cv.resize(img, dsize=None, fx=0.5, fy=0.5) cv.imshow('img', img_scaled) + out = out_writer(out, img_scaled) if frame_go: frame_num += 1 + cap.release() + out_writer(out, None) + cv.destroyAllWindows() +def out_writer(writer: cv.VideoWriter|None, frame: np.ndarray|None) -> cv.VideoWriter: + if writer is None: + fourcc = cv.VideoWriter.fourcc(*"mp4v") + frame_size = (frame.shape[1], frame.shape[0]) + writer = cv.VideoWriter(filename=f"{RESULTS_FOLDER}/out.mp4", fourcc=fourcc, fps=30.0, frameSize=frame_size) + + if writer is not None: + if frame is not None: + writer.write(frame) + else: + writer.release() + writer = None + + return writer + def main(): # construct the argument parser and parse the arguments