from typing import Optional, Callable, Dict import enum import time import cv2 import numpy as np import multiprocessing as mp from threadpoolctl import threadpool_limits from multiprocessing.managers import SharedMemoryManager from umi.common.timestamp_accumulator import get_accumulate_timestamp_idxs from umi.shared_memory.shared_memory_ring_buffer import SharedMemoryRingBuffer from umi.shared_memory.shared_memory_queue import SharedMemoryQueue, Full, Empty from umi.real_world.video_recorder import VideoRecorder from umi.common.usb_util import reset_usb_device class Command(enum.Enum): RESTART_PUT = 0 START_RECORDING = 1 STOP_RECORDING = 2 class UvcCamera(mp.Process): """ Call umi.common.usb_util.reset_all_elgato_devices if you are using Elgato capture cards. Required to workaround firmware bugs. """ MAX_PATH_LENGTH = 4096 # linux path has a limit of 4096 bytes def __init__( self, shm_manager: SharedMemoryManager, # v4l2 device file path # e.g. /dev/video0 # or /dev/v4l/by-id/usb-Elgato_Elgato_HD60_X_A00XB320216MTR-video-index0 dev_video_path, resolution=(1280, 720), capture_fps=60, put_fps=None, put_downsample=True, get_max_k=30, receive_latency=0.0, cap_buffer_size=1, num_threads=2, transform: Optional[Callable[[Dict], Dict]] = None, vis_transform: Optional[Callable[[Dict], Dict]] = None, recording_transform: Optional[Callable[[Dict], Dict]] = None, video_recorder: Optional[VideoRecorder] = None, verbose=False ): super().__init__() if put_fps is None: put_fps = capture_fps # create ring buffer resolution = tuple(resolution) shape = resolution[::-1] examples = { 'color': np.empty( shape=shape+(3,), dtype=np.uint8) } examples['camera_capture_timestamp'] = 0.0 examples['camera_receive_timestamp'] = 0.0 examples['timestamp'] = 0.0 examples['step_idx'] = 0 vis_ring_buffer = SharedMemoryRingBuffer.create_from_examples( shm_manager=shm_manager, examples=examples if vis_transform is None else vis_transform(dict(examples)), get_max_k=1, get_time_budget=0.2, put_desired_frequency=capture_fps ) ring_buffer = SharedMemoryRingBuffer.create_from_examples( shm_manager=shm_manager, examples=examples if transform is None else transform(dict(examples)), get_max_k=get_max_k, get_time_budget=0.2, put_desired_frequency=put_fps ) # create command queue examples = { 'cmd': Command.RESTART_PUT.value, 'put_start_time': 0.0, 'video_path': np.array('a'*self.MAX_PATH_LENGTH), 'recording_start_time': 0.0, } command_queue = SharedMemoryQueue.create_from_examples( shm_manager=shm_manager, examples=examples, buffer_size=128 ) # create video recorder if video_recorder is None: # default to nvenc GPU encoder video_recorder = VideoRecorder.create_hevc_nvenc( shm_manager=shm_manager, fps=capture_fps, input_pix_fmt='bgr24', bit_rate=6000*1000) assert video_recorder.fps == capture_fps # copied variables self.shm_manager = shm_manager self.dev_video_path = dev_video_path self.resolution = resolution self.capture_fps = capture_fps self.put_fps = put_fps self.put_downsample = put_downsample self.receive_latency = receive_latency self.cap_buffer_size = cap_buffer_size self.transform = transform self.vis_transform = vis_transform self.recording_transform = recording_transform self.video_recorder = video_recorder self.verbose = verbose self.put_start_time = None self.num_threads = num_threads # shared variables self.stop_event = mp.Event() self.ready_event = mp.Event() self.ring_buffer = ring_buffer self.vis_ring_buffer = vis_ring_buffer self.command_queue = command_queue # ========= context manager =========== def __enter__(self): self.start() return self def __exit__(self, exc_type, exc_val, exc_tb): self.stop() # ========= user API =========== def start(self, wait=True, put_start_time=None): self.put_start_time = put_start_time shape = self.resolution[::-1] data_example = np.empty(shape=shape+(3,), dtype=np.uint8) self.video_recorder.start( shm_manager=self.shm_manager, data_example=data_example) # must start video recorder first to create share memories super().start() if wait: self.start_wait() def stop(self, wait=True): self.video_recorder.stop() self.stop_event.set() if wait: self.end_wait() def start_wait(self): self.ready_event.wait() self.video_recorder.start_wait() def end_wait(self): self.join() self.video_recorder.end_wait() @property def is_ready(self): return self.ready_event.is_set() def get(self, k=None, out=None): if k is None: return self.ring_buffer.get(out=out) else: return self.ring_buffer.get_last_k(k, out=out) def get_vis(self, out=None): return self.vis_ring_buffer.get(out=out) def start_recording(self, video_path: str, start_time: float=-1): path_len = len(video_path.encode('utf-8')) if path_len > self.MAX_PATH_LENGTH: raise RuntimeError('video_path too long.') self.command_queue.put({ 'cmd': Command.START_RECORDING.value, 'video_path': video_path, 'recording_start_time': start_time }) def stop_recording(self): self.command_queue.put({ 'cmd': Command.STOP_RECORDING.value }) def restart_put(self, start_time): self.command_queue.put({ 'cmd': Command.RESTART_PUT.value, 'put_start_time': start_time }) # ========= interval API =========== def run(self): # limit threads threadpool_limits(self.num_threads) cv2.setNumThreads(self.num_threads) # open VideoCapture cap = cv2.VideoCapture(self.dev_video_path, cv2.CAP_V4L2) try: # set resolution and fps w, h = self.resolution fps = self.capture_fps cap.set(cv2.CAP_PROP_FRAME_WIDTH, w) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, h) # set fps cap.set(cv2.CAP_PROP_BUFFERSIZE, self.cap_buffer_size) cap.set(cv2.CAP_PROP_FPS, fps) # put frequency regulation put_idx = None put_start_time = self.put_start_time if put_start_time is None: put_start_time = time.time() # reuse frame buffer iter_idx = 0 t_start = time.time() while not self.stop_event.is_set(): ts = time.time() ret = cap.grab() assert ret # directly write into shared memory to avoid copy frame = self.video_recorder.get_img_buffer() ret, frame = cap.retrieve(frame) t_recv = time.time() assert ret mt_cap = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000 t_cap = mt_cap - time.monotonic() + time.time() t_cal = t_recv - self.receive_latency # calibrated latency # record frame if self.video_recorder.is_ready(): self.video_recorder.write_img_buffer(frame, frame_time=t_cal) data = dict() data['camera_receive_timestamp'] = t_recv data['camera_capture_timestamp'] = t_cap data['color'] = frame # apply transform put_data = data if self.transform is not None: put_data = self.transform(dict(data)) if self.put_downsample: # put frequency regulation local_idxs, global_idxs, put_idx \ = get_accumulate_timestamp_idxs( timestamps=[t_cal], start_time=put_start_time, dt=1/self.put_fps, # this is non in first iteration # and then replaced with a concrete number next_global_idx=put_idx, # continue to pump frames even if not started. # start_time is simply used to align timestamps. allow_negative=True ) for step_idx in global_idxs: put_data['step_idx'] = step_idx put_data['timestamp'] = t_cal self.ring_buffer.put(put_data, wait=False) else: step_idx = int((t_cal - put_start_time) * self.put_fps) put_data['step_idx'] = step_idx put_data['timestamp'] = t_cal self.ring_buffer.put(put_data, wait=False) # signal ready if iter_idx == 0: self.ready_event.set() # put to vis vis_data = data if self.vis_transform == self.transform: vis_data = put_data elif self.vis_transform is not None: vis_data = self.vis_transform(dict(data)) self.vis_ring_buffer.put(vis_data, wait=False) # perf t_end = time.time() duration = t_end - t_start frequency = np.round(1 / duration, 1) t_start = t_end if self.verbose: print(f'[UvcCamera {self.dev_video_path}] FPS {frequency}') # fetch command from queue try: commands = self.command_queue.get_all() n_cmd = len(commands['cmd']) except Empty: n_cmd = 0 # execute commands for i in range(n_cmd): command = dict() for key, value in commands.items(): command[key] = value[i] cmd = command['cmd'] if cmd == Command.RESTART_PUT.value: put_idx = None put_start_time = command['put_start_time'] elif cmd == Command.START_RECORDING.value: video_path = str(command['video_path']) start_time = command['recording_start_time'] if start_time < 0: start_time = None self.video_recorder.start_recording(video_path, start_time=start_time) elif cmd == Command.STOP_RECORDING.value: self.video_recorder.stop_recording() iter_idx += 1 finally: self.video_recorder.stop() # When everything done, release the capture cap.release()