import os import pickle import pybullet as p import time import pybullet_data import numpy as np import matplotlib.pyplot as plt from PIL import Image import open3d as o3d import cv2 import math import json class GenerateScene: def __init__(self, dataset_path, scene_path, output_path, camera_params) -> None: self._init_variables() self._load_object_dataset(dataset_path) self._load_scene(scene_path) self._set_output_path(output_path) self._set_camera_params(camera_params) def _init_variables(self): self.object_paths = {} self.scene = {} self.object_model_scale = [0.001, 0.001, 0.001] self.output_path = None self.camera_params = None self.segmentation_labels = {} def _extract_model_name(self, path): # Specify it according to specific file directory NUM_SLASH_BEFORE_NAME = 1 num_slash = 0 object_name_str = [] for i in range(len(path)): index = len(path) -1 - i char = path[index] if(num_slash == NUM_SLASH_BEFORE_NAME): object_name_str.append(char) if(char == "/"): num_slash += 1 object_name_str.reverse() object_name_str = ''.join(object_name_str[1:]) return object_name_str def _load_object_dataset(self, dataset_path): if(dataset_path[-1] == "/"): dataset_path = dataset_path[:-1] for root, dirs, files in os.walk(dataset_path, topdown=False): for name in dirs: path = os.path.join(root, name) if(os.path.join(root, name)[-4:] == "Scan"): name = self._extract_model_name(path) self.object_paths[name] = path+"/Simp.obj" def _load_scene(self, scene_path): if(scene_path[-1] == "/"): scene_path = scene_path[:-1] scene_path = scene_path + "/scene.pickle" self.scene = pickle.load(open(scene_path, 'rb')) def _set_output_path(self, output_path): self.output_path = output_path if(self.output_path[-1] == "/"): self.output_path = self.output_path[:-1] def _set_camera_params(self, camera_params): self.camera_params = camera_params def load_camera_pose_from_frame(self, camera_params_path): with open(camera_params_path, "r") as f: camera_params = json.load(f) view_transform = camera_params["cameraViewTransform"] print(view_transform) view_transform = np.resize(view_transform, (4,4)) view_transform = np.linalg.inv(view_transform).T offset = np.mat([[1,0,0,0],[0,-1,0,0],[0,0,-1,0],[0,0,0,1]]) view_transform = view_transform.dot(offset) print(view_transform) return view_transform def _load_obj_to_pybullet(self, obj_file_path, position, orientation, scale): visual_ind = p.createVisualShape( shapeType=p.GEOM_MESH, fileName=obj_file_path, rgbaColor=[1, 1, 1, 1], specularColor=[0.4, 0.4, 0], visualFramePosition=[0, 0, 0], meshScale=scale) p.createMultiBody( baseMass=1, baseVisualShapeIndex=visual_ind, basePosition=position, baseOrientation=orientation, useMaximalCoordinates=True) def _render_image(self, camera_pose): width = self.camera_params["width"] height = self.camera_params["height"] fov = self.camera_params["fov"] aspect = width / height near = self.camera_params["near"] far = self.camera_params["far"] T = np.mat([[1,0,0,0], [0,1,0,0], [0,0,1,1], [0,0,0,1]]) look_at_T = camera_pose.dot(T) view_matrix = p.computeViewMatrix([camera_pose[0,3], camera_pose[1,3], camera_pose[2,3]], [look_at_T[0,3], look_at_T[1,3], look_at_T[2,3]], [-camera_pose[0,1], -camera_pose[1,1], -camera_pose[2,1]]) projection_matrix = p.computeProjectionMatrixFOV(fov, aspect, near, far) # Get depth values using the OpenGL renderer images = p.getCameraImage(width, height, view_matrix, projection_matrix, renderer=p.ER_BULLET_HARDWARE_OPENGL) rgb = images[2] depth = images[3] seg = images[4] rgb_image = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB) cv2.imwrite(self.output_path+'/rgb.jpg',rgb_image) depth_image = far * near / (far - (far - near) * depth) depth_image = np.asanyarray(depth_image).astype(np.float32) * 1000.0 depth_image_array = depth_image depth_image = (depth_image.astype(np.uint16)) depth_image = Image.fromarray(depth_image) depth_image.save(self.output_path+'/depth.png') cv2.imwrite(self.output_path+'/seg.jpg', seg) id = 0 for object_name in self.scene.keys(): self.segmentation_labels[str(id+1)] = object_name id += 1 with open(self.output_path+"/seg_labels.json", 'w') as seg_labels: json.dump(self.segmentation_labels, seg_labels) with open(self.output_path+"/cam_intrinsics.json", 'w') as cam_intrinsics: json.dump(self.camera_params, cam_intrinsics) def generate_images(self, camera_pose): physicsClient = p.connect(p.GUI) p.setAdditionalSearchPath(pybullet_data.getDataPath()) p.setGravity(0,0,0) p.loadURDF("plane100.urdf") for obj_name in self.scene.keys(): orientation = self.scene[obj_name]["rotation"] position = self.scene[obj_name]["position"] self._load_obj_to_pybullet(obj_file_path=self.object_paths[obj_name], position=position, orientation=orientation, scale=self.object_model_scale) self._render_image(camera_pose) p.stepSimulation() p.disconnect() def visualize_pcd(self): color_image = o3d.io.read_image(self.output_path+'/rgb.jpg') depth_image = o3d.io.read_image(self.output_path+'/depth.png') rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(color_image, depth_image,convert_rgb_to_intensity=False) intrinsic = o3d.camera.PinholeCameraIntrinsic( o3d.camera.PinholeCameraIntrinsicParameters.Kinect2DepthCameraDefault ) intrinsic.set_intrinsics(width=self.camera_params["width"], height=self.camera_params["height"], fx=self.camera_params["fx"], fy=self.camera_params["fy"], cx=self.camera_params["cx"], cy=self.camera_params["cy"]) point_cloud = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, intrinsic) o3d.visualization.draw_geometries([point_cloud]) def print_scene(self): print("================= Scene Objects: =================") print(self.scene.keys()) print("==================================================") DATASET_PATH = "/home/hzx/Downloads/OmniObject3d-simplified/output" SCENE_PATH = "/home/hzx/Projects/ActivePerception/data_generation/output/scene_0/" OUTPUT_PATH = "/home/hzx/Projects/ActivePerception/data_generation/output/" FRAME_PATH = "/home/hzx/Projects/ActivePerception/data_generation/output/scene_0/camera_params_0119.json" ISAAC_SIM_CAM_H_APERTURE = 20.955 # Isaac Sim里读取的 ISAAC_SIM_CAM_V_APERTURE = 15.2908 # Isaac Sim里读取的 ISAAC_SIM_FOCAL_LENGTH = 39 # 试出来的，其实Isaac Sim里原本是24 ISAAC_SIM_CAM_D_APERTURE = math.sqrt(ISAAC_SIM_CAM_H_APERTURE**2 + ISAAC_SIM_CAM_V_APERTURE**2) CAM_WIDTH = 640 CAM_HEIGHT = 480 CAM_FOV = 2*math.atan(ISAAC_SIM_CAM_D_APERTURE/(2*ISAAC_SIM_FOCAL_LENGTH))/math.pi*180 CAM_NEAR = 0.001 # 成像最近距离 CAM_FAR = 10 # 成像最远距离 CAM_CX = CAM_WIDTH/2 CAM_CY = CAM_HEIGHT/2 CAM_FX = 1/math.tan(CAM_FOV*math.pi/180.0/2.0)*CAM_WIDTH/2 CAM_FY = 1/(CAM_HEIGHT/CAM_WIDTH*math.tan(CAM_FOV*math.pi/180.0/2.0))*CAM_HEIGHT/2 CAMERA_PARAMS = { "width": CAM_WIDTH, "height": CAM_HEIGHT, "fov": CAM_FOV, "near": CAM_NEAR, "far": CAM_FAR, "cx": CAM_CX, "cy": CAM_CY, "fx": CAM_FX, "fy": CAM_FY, } if __name__ == "__main__": gs = GenerateScene(DATASET_PATH, SCENE_PATH, OUTPUT_PATH, CAMERA_PARAMS) gs.print_scene() cam_pose = gs.load_camera_pose_from_frame(FRAME_PATH) gs.generate_images(cam_pose) # 在OUTPUT_PATH路径下生成rgb、depth、segmentation图 # gs.visualize_pcd()