240 lines
8.7 KiB
Python
Executable File
240 lines
8.7 KiB
Python
Executable File
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() |