diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py index b7d1c79..6c59371 100644 --- a/lerobot/common/robot_devices/control_utils.py +++ b/lerobot/common/robot_devices/control_utils.py @@ -260,10 +260,11 @@ def control_loop( dataset.add_frame(frame) if display_cameras and not is_headless(): - image_keys = [key for key in observation if "image" in key] - for key in image_keys: - cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR)) - cv2.waitKey(1) + # image_keys = [key for key in observation if "image" in key] + # for key in image_keys: + # cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR)) + # cv2.waitKey(1) + display_observations_combined(observation) if fps is not None: dt_s = time.perf_counter() - start_loop_t @@ -277,6 +278,107 @@ def control_loop( events["exit_early"] = False break +import numpy as np +def display_observations_combined(observation, display_cameras=True): + """将摄像头画面组合在一个窗口中显示,优化以防止抖动""" + if display_cameras and not is_headless(): + image_keys = [key for key in observation if "image" in key] + if not image_keys: + return + + # 获取所有图像并转换为BGR格式 + images = [] + for key in image_keys: + img = observation[key].numpy() + # 确保每个图像尺寸为640×480 + if img.shape[0] != 480 or img.shape[1] != 640: + img = cv2.resize(img, (640, 480)) + img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) + # 添加图像标题 + cv2.putText(img, key, (10, 30), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) + images.append(img) + + # 根据实际摄像头数量确定布局 + num_cameras = len(images) + + # 为固定数量的摄像头定义固定布局 + margin = 10 # 图像之间的间距 + + if num_cameras == 1: + # 单个摄像头 - 直接显示 + grid_image = images[0] + elif num_cameras == 2: + # 两个摄像头 - 水平排列 + grid_width = 2 * 640 + margin + grid_height = 480 + grid_image = np.zeros((grid_height, grid_width, 3), dtype=np.uint8) + + # 放置两个图像 + grid_image[0:480, 0:640] = images[0] + grid_image[0:480, 640+margin:] = images[1] + elif num_cameras == 3: + # 三个摄像头 - 上面一行2个,下面一个 + grid_width = 2 * 640 + margin + grid_height = 2 * 480 + margin + grid_image = np.zeros((grid_height, grid_width, 3), dtype=np.uint8) + + # 放置三个图像 + grid_image[0:480, 0:640] = images[0] + grid_image[0:480, 640+margin:] = images[1] + grid_image[480+margin:, (grid_width-640)//2:(grid_width+640)//2] = images[2] # 居中放置 + else: + # 四个或更多摄像头 - 2×2网格 + grid_cols = 2 + grid_rows = 2 + grid_width = grid_cols * 640 + (grid_cols - 1) * margin + grid_height = grid_rows * 480 + (grid_rows - 1) * margin + grid_image = np.zeros((grid_height, grid_width, 3), dtype=np.uint8) + + # 最多显示4个 + for i, img in enumerate(images[:4]): + row = i // grid_cols + col = i % grid_cols + + y_start = row * (480 + margin) + y_end = y_start + 480 + x_start = col * (640 + margin) + x_end = x_start + 640 + + grid_image[y_start:y_end, x_start:x_end] = img + + # 创建一个固定名称的窗口 + window_name = "Camera Views" + + # 使用 getWindowProperty 检查窗口是否已经存在 + try: + # 尝试获取窗口属性 - 如果窗口不存在会抛出异常 + window_exists = cv2.getWindowProperty(window_name, cv2.WND_PROP_VISIBLE) >= 0 + except: + window_exists = False + + # 如果窗口不存在,创建并定位它 + if not window_exists: + cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) + + # 设置窗口位置和大小(只需要在创建时设置一次) + screen_width = 1920 # 您的屏幕宽度 + screen_height = 1080 # 您的屏幕高度 + + # 根据实际图像计算合适的显示尺寸 + scale_factor = min(screen_width / grid_width, screen_height / grid_height) * 0.95 + display_width = int(grid_width * scale_factor) + display_height = int(grid_height * scale_factor) + + cv2.resizeWindow(window_name, display_width, display_height) + cv2.moveWindow(window_name, (screen_width - display_width) // 2, + (screen_height - display_height) // 2) # 居中显示 + + # 显示图像 + cv2.imshow(window_name, grid_image) + cv2.waitKey(1) + + def reset_environment(robot, events, reset_time_s, fps): # TODO(rcadene): refactor warmup_record and reset_environment