Skip to content

Commit

Permalink
update to remove some dependencies
Browse files Browse the repository at this point in the history
  • Loading branch information
@beneisner
beneisner committed Aug 2, 2023
1 parent f1fd88f commit 0d81cc4
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 256 deletions.
40 changes: 10 additions & 30 deletions pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,23 +4,19 @@ version = "0.1.0"
description = "Partnet-Mobility Utils"
readme = "README.md"
requires-python = ">=3.6"
license = {file = "LICENSE.txt"}
authors = [
{email = "[email protected]", name = "Ben Eisner"}
]
license = { file = "LICENSE.txt" }
authors = [{ email = "[email protected]", name = "Ben Eisner" }]
dependencies = [
"numpy",
"rpad-core @ git+https://github.com/r-pad/core.git",
"rpad-pybullet-libs @ git+https://github.com/r-pad/pybullet_libs.git",
"scipy",
"trimesh",
"typer",
]

[build-system]
requires = [
"setuptools >= 62.3.2,<63",
"setuptools-scm",
"wheel",
]
requires = ["setuptools >= 62.3.2,<63", "setuptools-scm", "wheel"]
build-backend = "setuptools.build_meta"

[project.optional-dependencies]
Expand All @@ -34,20 +30,10 @@ develop = [
"pre-commit",
"jupyterlab",
]
build_docs = [
"mkdocs-material",
"mkdocstrings[python]",
]
pybullet = [
"wheel",
"pybullet",
]
sapien = [
"sapien",
]
notebooks = [
"rpad-visualize-3d @ git+https://github.com/r-pad/visualize_3d",
]
build_docs = ["mkdocs-material", "mkdocstrings[python]"]
pybullet = ["wheel", "pybullet"]
sapien = ["sapien"]
notebooks = ["rpad-visualize-3d @ git+https://github.com/r-pad/visualize_3d"]

# Make setuptools only look in the src folder.
[tool.setuptools.packages.find]
Expand All @@ -68,11 +54,5 @@ namespace_packages = true
explicit_package_bases = true

[[tool.mypy.overrides]]
module = [
"pybullet.*",
"pybullet_data.*",
"trimesh.*",
"scipy.*",
"sapien.*",
]
module = ["pybullet.*", "pybullet_data.*", "trimesh.*", "scipy.*", "sapien.*"]
ignore_missing_imports = true
238 changes: 12 additions & 226 deletions src/rpad/partnet_mobility_utils/render/pybullet.py
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
import functools
import os
import sys
from contextlib import contextmanager
from typing import Dict, List, Literal, Mapping, Optional, Tuple, Union

import numpy as np
from rpad.core.distributed import NPSeed
from rpad.pybullet_libs.camera import Camera, Render
from rpad.pybullet_libs.utils import get_obj_z_offset, isnotebook, suppress_stdout
from scipy.spatial.transform import Rotation as R

from rpad.partnet_mobility_utils.data import PMObject
Expand Down Expand Up @@ -74,217 +73,6 @@ def randomize_camera(env, seed=None):
return


def get_obj_z_offset(object_id, sim_id, starting_min=0.0):
bboxes = [p.getAABB(object_id, physicsClientId=sim_id)] # Add the base one.
for i in range(p.getNumJoints(object_id, physicsClientId=sim_id)):
aabb = p.getAABB(object_id, i, physicsClientId=sim_id) # Add the links.
bboxes.append(aabb)
minz = functools.reduce(lambda a, b: min(a, b[0][2]), bboxes, starting_min)
return minz


def get_obj_bbox_xy(object_id, sim_id):
bboxes = [p.getAABB(object_id, physicsClientId=sim_id)] # Add the base one.
for i in range(p.getNumJoints(object_id, physicsClientId=sim_id)):
aabb = p.getAABB(object_id, i, physicsClientId=sim_id) # Add the links.
bboxes.append(aabb)
# xmin = np.array(bboxes)[:, :, 0].min()
xmin = functools.reduce(lambda a, b: min(a, b[0][0]), bboxes, 0.0)
# xmax = np.array(bboxes)[:, :, 0].max()
xmax = functools.reduce(lambda a, b: max(a, b[1][0]), bboxes, 0.0)
# ymin = np.array(bboxes)[:, :, 1].min()
ymin = functools.reduce(lambda a, b: min(a, b[0][1]), bboxes, 0.0)
# ymax = np.array(bboxes)[:, :, 1].max()
ymax = functools.reduce(lambda a, b: max(a, b[1][1]), bboxes, 0.0)
return [[xmin, xmax], [ymin, ymax]]


def isnotebook():
try:
shell = get_ipython().__class__.__name__
if shell == "ZMQInteractiveShell":
return True # Jupyter notebook or qtconsole
elif shell == "TerminalInteractiveShell":
return False # Terminal running IPython
else:
return False # Other type (?)
except NameError:
return False # Probably standard Python interpreter


@contextmanager
def suppress_stdout():
fd = sys.stdout.fileno()

def _redirect_stdout(to):
sys.stdout.close() # + implicit flush()
os.dup2(to.fileno(), fd) # fd writes to 'to' file
sys.stdout = os.fdopen(fd, "w") # Python writes to fd

with os.fdopen(os.dup(fd), "w") as old_stdout:
with open(os.devnull, "w") as file:
_redirect_stdout(to=file)
try:
yield # allow code to be run with the redirected stdout
finally:
_redirect_stdout(to=old_stdout) # restore stdout.
# buffering and flags such as
# CLOEXEC may be different


RENDER_WIDTH = 640
RENDER_HEIGHT = 480
CAMERA_INTRINSICS = np.array(
[
[450, 0, RENDER_WIDTH / 2],
[0, 450, RENDER_HEIGHT / 2],
[0, 0, 1],
]
)

T_CAMGL_2_CAM = np.array(
[
[1, 0, 0, 0],
[0, -1, 0, 0],
[0, 0, -1, 0],
[0, 0, 0, 1],
]
)


def get_pointcloud(depth, intrinsics):
"""Get 3D pointcloud from perspective depth image.
Args:
depth: HxW float array of perspective depth in meters.
intrinsics: 3x3 float array of camera intrinsics matrix.
Returns:
points: HxWx3 float array of 3D points in camera coordinates.
"""
height, width = depth.shape
xlin = np.linspace(0, width - 1, width)
ylin = np.linspace(0, height - 1, height)
px, py = np.meshgrid(xlin, ylin)
px = (px - intrinsics[0, 2]) * (depth / intrinsics[0, 0])
py = (py - intrinsics[1, 2]) * (depth / intrinsics[1, 1])
points = np.float32([px, py, depth]).transpose(1, 2, 0)
return points


class Camera:
def __init__(
self,
pos,
render_height=RENDER_HEIGHT,
render_width=RENDER_WIDTH,
znear=0.01,
zfar=6,
intrinsics=CAMERA_INTRINSICS,
target=None,
):
#######################################
# First, compute the projection matrix.
#######################################
self.intrinsics = intrinsics
focal_length = intrinsics[0][0]
self.znear, self.zfar = znear, zfar
self.fovh = (np.arctan((render_height / 2) / focal_length) * 2 / np.pi) * 180
self.render_width = render_width
self.render_height = render_height

# Notes: 1) FOV is vertical FOV 2) aspect must be float
aspect_ratio = render_width / render_height
self.proj_list = p.computeProjectionMatrixFOV(
self.fovh, aspect_ratio, self.znear, self.zfar
)

#######################################
# Next, compute the view matrix.
#######################################
if target is None:
target = [0, 0, 0.5]
self.target = target
self.view_list = self.__view_list(pos, target)

@property
def view_list(self):
return self._view_list

@view_list.setter
def view_list(self, value):
self._view_list = value
self.T_camgl2world = np.asarray(value).reshape(4, 4).T
self.T_world2camgl = np.linalg.inv(self.T_camgl2world)
self.T_world2cam = self.T_world2camgl @ T_CAMGL_2_CAM

@staticmethod
def __view_list(eye, target):
up = [0.0, 0.0, 1.0]
target = target
view_list = p.computeViewMatrix(eye, target, up)
return view_list

def set_camera_position(self, pos):
self.view_list = self.__view_list(pos, self.target)

def render(
self, client_id, return_prgb=False, has_plane=True, link_seg=True
) -> Tuple[np.ndarray, ...]:
if link_seg:
_, _, rgb, zbuffer, seg = p.getCameraImage(
RENDER_WIDTH,
RENDER_HEIGHT,
self.view_list,
self.proj_list,
flags=p.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,
physicsClientId=client_id,
)
else:
_, _, rgb, zbuffer, seg = p.getCameraImage(
RENDER_WIDTH,
RENDER_HEIGHT,
self.view_list,
self.proj_list,
physicsClientId=client_id,
)

# Sometimes on mac things get weird.
if isinstance(rgb, tuple):
rgb = np.asarray(rgb).reshape(RENDER_HEIGHT, RENDER_WIDTH, 4)
zbuffer = np.asarray(zbuffer).reshape(RENDER_HEIGHT, RENDER_WIDTH)
seg = np.asarray(seg).reshape(RENDER_HEIGHT, RENDER_WIDTH)

zfar, znear = self.zfar, self.znear
depth = zfar + znear - (2.0 * zbuffer - 1.0) * (zfar - znear)
depth = (2.0 * znear * zfar) / depth

P_cam = get_pointcloud(depth, self.intrinsics)
foreground_ixs = seg > 0 if has_plane else seg > -1
pc_seg = seg[foreground_ixs].flatten()
P_cam = P_cam[foreground_ixs]
P_cam = P_cam.reshape(-1, 3)
P_rgb = rgb[foreground_ixs]
P_rgb = P_rgb[:, :3].reshape(-1, 3)

Ph_cam = np.concatenate([P_cam, np.ones((len(P_cam), 1))], axis=1)
Ph_world = (self.T_world2cam @ Ph_cam.T).T
P_world = Ph_world[:, :3]

# Undoing the bitmask so we can get the obj_id, link_index
segmap: Optional[Dict]
if link_seg:
segmap = {
label: ((label & ((1 << 24) - 1)), (label >> 24) - 1)
for label in np.unique(seg)
}
else:
segmap = None

if return_prgb:
return rgb, depth, seg, P_cam, P_world, P_rgb, pc_seg, segmap # type: ignore

return rgb, depth, seg, P_cam, P_world, pc_seg, segmap # type: ignore


class PMRenderEnv:
def __init__(
self,
Expand Down Expand Up @@ -358,17 +146,8 @@ def __init__(
if info[2] == p.JOINT_REVOLUTE or info[2] == p.JOINT_PRISMATIC:
self.jn_to_ix[joint_name] = _id

def render(self, return_prgb=False, link_seg=True):
if not return_prgb:
rgb, depth, seg, P_cam, P_world, pc_seg, segmap = self.camera.render(
self.client_id, return_prgb, self.with_plane, link_seg
)
return rgb, depth, seg, P_cam, P_world, pc_seg, segmap
else:
rgb, depth, seg, P_cam, P_world, P_rgb, pc_seg, segmap = self.camera.render(
self.client_id, return_prgb, self.with_plane, link_seg
)
return rgb, depth, seg, P_cam, P_world, P_rgb, pc_seg, segmap
def render(self, return_prgb=False, link_seg=True) -> Render:
return self.camera.render(self.client_id, self.with_plane, link_seg)

def set_camera(
self,
Expand Down Expand Up @@ -520,7 +299,14 @@ def render(
if camera_xyz is not None:
self._render_env.set_camera(camera_xyz, seed=seed2)

rgb, depth, seg, P_cam, P_world, pc_seg, segmap = self._render_env.render()
obs = self._render_env.render()
rgb = obs["rgb"]
depth = obs["depth"]
seg = obs["seg"]
P_cam = obs["P_cam"]
P_world = obs["P_world"]
pc_seg = obs["pc_seg"]
segmap = obs["segmap"]

# Reindex the segmentation.
pc_seg_obj = np.ones_like(pc_seg) * -1
Expand Down

0 comments on commit 0d81cc4

Please sign in to comment.