hofee/nbv_grasping
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

success

Browse Source
This commit is contained in:
hofee
2024-10-09 16:13:22 +00:00
commit 0ea3f048dc
437 changed files with 44406 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

240
data_generation/tools/get_view.py Executable file
View File

@@ -0,0 +1,240 @@
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()