离心机加试管初步配置
This commit is contained in:
@@ -89,59 +89,6 @@ import numpy as np
|
||||
|
||||
"""use gpu compute to record video"""
|
||||
|
||||
class MultiCameraRecorder:
|
||||
def __init__(self, env, camera_names: list[str], env_indices: list[int], output_dir: str, fps: int = 30):
|
||||
self.env = env
|
||||
self.camera_names = camera_names
|
||||
self.env_indices = env_indices
|
||||
self.output_dir = output_dir
|
||||
self.fps = fps
|
||||
self.frames = {cam_name: {env_idx: [] for env_idx in env_indices} for cam_name in camera_names}
|
||||
os.makedirs(self.output_dir, exist_ok=True)
|
||||
|
||||
self.cameras = {}
|
||||
for name in camera_names:
|
||||
if name in self.env.unwrapped.scene.keys():
|
||||
self.cameras[name] = self.env.unwrapped.scene[name]
|
||||
print(f"[INFO] Camera {name} linked.")
|
||||
|
||||
def record_step(self):
|
||||
"""保持在 GPU 上克隆数据"""
|
||||
for cam_name, camera_obj in self.cameras.items():
|
||||
# 获取数据前强制同步一次（防止后端丢失）
|
||||
rgb_data = camera_obj.data.output["rgb"]
|
||||
|
||||
for env_idx in self.env_indices:
|
||||
# 使用 .clone() 保持在 GPU，但要注意显存占用
|
||||
self.frames[cam_name][env_idx].append(rgb_data[env_idx].clone())
|
||||
|
||||
def save_videos(self, filename_suffix=""):
|
||||
print(f"[INFO] Saving videos from GPU to Disk...")
|
||||
for cam_name, env_dict in self.frames.items():
|
||||
for env_idx, frame_list in env_dict.items():
|
||||
if not frame_list: continue
|
||||
|
||||
# 转换为 torchvision 格式 (T, C, H, W)
|
||||
video_tensor = torch.stack(frame_list)
|
||||
if video_tensor.shape[-1] == 4: # RGBA -> RGB
|
||||
video_tensor = video_tensor[..., :3]
|
||||
|
||||
# 移动到 CPU 并保存
|
||||
video_cpu = video_tensor.cpu()
|
||||
output_path = os.path.join(self.output_dir, f"{cam_name}_env{env_idx}_{filename_suffix}.mp4")
|
||||
|
||||
# 使用 torchvision 写入 (T, H, W, C)
|
||||
torchvision.io.write_video(output_path, video_cpu, fps=self.fps)
|
||||
|
||||
# 【关键】保存后立即释放显存
|
||||
del video_tensor
|
||||
del video_cpu
|
||||
frame_list.clear()
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
"""use cpu compute to record video"""
|
||||
# # 2. 修改 MultiCameraRecorder 类
|
||||
# class MultiCameraRecorder:
|
||||
# def __init__(self, env, camera_names: list[str], env_indices: list[int], output_dir: str, fps: int = 30):
|
||||
# self.env = env
|
||||
@@ -151,57 +98,110 @@ class MultiCameraRecorder:
|
||||
# self.fps = fps
|
||||
# self.frames = {cam_name: {env_idx: [] for env_idx in env_indices} for cam_name in camera_names}
|
||||
# os.makedirs(self.output_dir, exist_ok=True)
|
||||
|
||||
# self.cameras = {}
|
||||
# for name in camera_names:
|
||||
# try:
|
||||
# if name in self.env.unwrapped.scene.keys():
|
||||
# self.cameras[name] = self.env.unwrapped.scene[name]
|
||||
# print(f"[INFO][MultiCameraRecorder] Found camera: {name}")
|
||||
# except KeyError:
|
||||
# print(f"[WARN][MultiCameraRecorder] Camera '{name}' not found!")
|
||||
# print(f"[INFO] Camera {name} linked.")
|
||||
|
||||
# def record_step(self):
|
||||
# """在每个仿真步调用"""
|
||||
# """保持在 GPU 上克隆数据"""
|
||||
# for cam_name, camera_obj in self.cameras.items():
|
||||
# # [关键修改] 获取数据前先确保数据已同步
|
||||
# # 这可以防止读取到正在渲染中的内存导致 access violation
|
||||
# rgb_data = camera_obj.data.output["rgb"]
|
||||
# # 获取数据前强制同步一次（防止后端丢失）
|
||||
# rgb_data = camera_obj.data.output["rgb"]
|
||||
|
||||
# for env_idx in self.env_indices:
|
||||
# if env_idx >= rgb_data.shape[0]: continue
|
||||
|
||||
# # 转换为 CPU 上的 numpy，这种方式通常比 torchvision 的 tensor 堆叠更稳
|
||||
# frame = rgb_data[env_idx].clone().detach().cpu().numpy()
|
||||
# self.frames[cam_name][env_idx].append(frame)
|
||||
# # 使用 .clone() 保持在 GPU，但要注意显存占用
|
||||
# self.frames[cam_name][env_idx].append(rgb_data[env_idx].clone())
|
||||
|
||||
# def save_videos(self, filename_suffix=""):
|
||||
# """循环结束后调用"""
|
||||
# print(f"[INFO][MultiCameraRecorder] Saving videos...")
|
||||
# print(f"[INFO] Saving videos from GPU to Disk...")
|
||||
# for cam_name, env_dict in self.frames.items():
|
||||
# for env_idx, frame_list in env_dict.items():
|
||||
# if not frame_list: continue
|
||||
|
||||
# print(f" -> Saving {cam_name} (Env {env_idx})...")
|
||||
# # 转换为 torchvision 格式 (T, C, H, W)
|
||||
# video_tensor = torch.stack(frame_list)
|
||||
# if video_tensor.shape[-1] == 4: # RGBA -> RGB
|
||||
# video_tensor = video_tensor[..., :3]
|
||||
|
||||
# # 处理格式并使用 imageio 保存
|
||||
# processed_frames = []
|
||||
# for img in frame_list:
|
||||
# # [0, 1] -> [0, 255]
|
||||
# if img.dtype != np.uint8:
|
||||
# if img.max() <= 1.01: img = (img * 255).astype(np.uint8)
|
||||
# else: img = img.astype(np.uint8)
|
||||
# # 去掉 Alpha 通道
|
||||
# if img.shape[-1] == 4: img = img[:, :, :3]
|
||||
# processed_frames.append(img)
|
||||
# # 移动到 CPU 并保存
|
||||
# video_cpu = video_tensor.cpu()
|
||||
# output_path = os.path.join(self.output_dir, f"{cam_name}_env{env_idx}_{filename_suffix}.mp4")
|
||||
|
||||
# # 使用 torchvision 写入 (T, H, W, C)
|
||||
# torchvision.io.write_video(output_path, video_cpu, fps=self.fps)
|
||||
|
||||
# # 【关键】保存后立即释放显存
|
||||
# del video_tensor
|
||||
# del video_cpu
|
||||
# frame_list.clear()
|
||||
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
# fname = f"{cam_name}_env{env_idx}_{filename_suffix}.mp4"
|
||||
# output_path = os.path.join(self.output_dir, fname)
|
||||
"""use cpu compute to record video"""
|
||||
# # 2. 修改 MultiCameraRecorder 类
|
||||
class MultiCameraRecorder:
|
||||
def __init__(self, env, camera_names: list[str], env_indices: list[int], output_dir: str, fps: int = 30):
|
||||
self.env = env
|
||||
self.camera_names = camera_names
|
||||
self.env_indices = env_indices
|
||||
self.output_dir = output_dir
|
||||
self.fps = fps
|
||||
self.frames = {cam_name: {env_idx: [] for env_idx in env_indices} for cam_name in camera_names}
|
||||
os.makedirs(self.output_dir, exist_ok=True)
|
||||
self.cameras = {}
|
||||
for name in camera_names:
|
||||
try:
|
||||
self.cameras[name] = self.env.unwrapped.scene[name]
|
||||
print(f"[INFO][MultiCameraRecorder] Found camera: {name}")
|
||||
except KeyError:
|
||||
print(f"[WARN][MultiCameraRecorder] Camera '{name}' not found!")
|
||||
|
||||
def record_step(self):
|
||||
"""在每个仿真步调用"""
|
||||
for cam_name, camera_obj in self.cameras.items():
|
||||
# [关键修改] 获取数据前先确保数据已同步
|
||||
# 这可以防止读取到正在渲染中的内存导致 access violation
|
||||
rgb_data = camera_obj.data.output["rgb"]
|
||||
|
||||
for env_idx in self.env_indices:
|
||||
if env_idx >= rgb_data.shape[0]: continue
|
||||
|
||||
# try:
|
||||
# # 使用 imageio 写入视频
|
||||
# imageio.mimsave(output_path, processed_frames, fps=self.fps)
|
||||
# print(f" Saved: {output_path}")
|
||||
# except Exception as e:
|
||||
# print(f" [ERROR] Failed to save {fname}: {e}")
|
||||
# 转换为 CPU 上的 numpy，这种方式通常比 torchvision 的 tensor 堆叠更稳
|
||||
frame = rgb_data[env_idx].clone().detach().cpu().numpy()
|
||||
self.frames[cam_name][env_idx].append(frame)
|
||||
|
||||
def save_videos(self, filename_suffix=""):
|
||||
"""循环结束后调用"""
|
||||
print(f"[INFO][MultiCameraRecorder] Saving videos...")
|
||||
for cam_name, env_dict in self.frames.items():
|
||||
for env_idx, frame_list in env_dict.items():
|
||||
if not frame_list: continue
|
||||
|
||||
print(f" -> Saving {cam_name} (Env {env_idx})...")
|
||||
|
||||
# 处理格式并使用 imageio 保存
|
||||
processed_frames = []
|
||||
for img in frame_list:
|
||||
# [0, 1] -> [0, 255]
|
||||
if img.dtype != np.uint8:
|
||||
if img.max() <= 1.01: img = (img * 255).astype(np.uint8)
|
||||
else: img = img.astype(np.uint8)
|
||||
# 去掉 Alpha 通道
|
||||
if img.shape[-1] == 4: img = img[:, :, :3]
|
||||
processed_frames.append(img)
|
||||
|
||||
fname = f"{cam_name}_env{env_idx}_{filename_suffix}.mp4"
|
||||
output_path = os.path.join(self.output_dir, fname)
|
||||
|
||||
try:
|
||||
# 使用 imageio 写入视频
|
||||
imageio.mimsave(output_path, processed_frames, fps=self.fps)
|
||||
print(f" Saved: {output_path}")
|
||||
except Exception as e:
|
||||
print(f" [ERROR] Failed to save {fname}: {e}")
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user