Source code for sapien.asset

import requests
import zipfile
import io
import os
from pathlib import Path
import numpy as np
import sapien
from typing import Tuple

url = "https://sapien.ucsd.edu/api/download/compressed/{}.zip?token={}"


[docs] def download_partnet_mobility(model_id, token=None, directory=None): if not directory: directory = os.environ.get("PARTNET_MOBILITY_DATASET") if not directory: directory = "partnet-mobility-dataset" urdf_file = os.path.join(directory, str(model_id), "mobility.urdf") # return if file exists if os.path.exists(urdf_file): return urdf_file if token is None: if "SAPIEN_ACCESS_TOKEN" not in os.environ: raise Exception( "To download the model, you need to provide the token or set the SAPIEN_ACCESS_TOKEN environment variable." ) token = os.environ["SAPIEN_ACCESS_TOKEN"] # download file print("downloading from", url.format(model_id, token)) r = requests.get(url.format(model_id, token), stream=True) if not r.ok: raise Exception( "Download PartNet-Mobility failed. " "Please check your token and IP address." "Also make sure sure the model id is valid" ) z = zipfile.ZipFile(io.BytesIO(r.content)) os.makedirs(directory, exist_ok=True) z.extractall(directory) return urdf_file
[docs] def create_dome_envmap( filename="/tmp/sapien_dome.ktx", sky_color=[0.6, 0.6, 0.6], ground_color=[0.2, 0.2, 0.2], blend=0.3, resolution=256, ): assert str(filename).endswith(".ktx"), "generated envmap must end with .ktx" path = Path(filename) path.parent.mkdir(exist_ok=True) import numpy as np from PIL import Image import tempfile sky_color = np.array(sky_color) ground_color = np.array(ground_color) py = (np.ones((resolution, resolution, 3)) * sky_color * 255).astype(np.uint8) ny = (np.ones((resolution, resolution, 3)) * ground_color * 255).astype(np.uint8) other = np.zeros((resolution, resolution, 3), dtype=np.uint8) c = np.linspace(0.5, resolution - 0.5, resolution) * 2 / resolution - 1 coords = np.stack(np.meshgrid(c, c), 0) angles = np.arctan2(coords[1], np.sqrt(coords[0] ** 2 + 1)) r = (angles / blend).clip(-1, 1) * 0.5 + 0.5 other = (1 - r[..., None]) * sky_color + r[..., None] * ground_color other = (other.clip(0, 1) * 255).astype(np.uint8) with tempfile.TemporaryDirectory() as d: fpy = os.path.join(d, "py.png") fny = os.path.join(d, "ny.png") fother = os.path.join(d, "other.png") Image.fromarray(py).save(fpy) Image.fromarray(ny).save(fny) Image.fromarray(other).save(fother) cubemap = sapien.render.RenderCubemap(fother, fother, fpy, fny, fother, fother) cubemap.export(filename) return cubemap
# def create_checkerboard( # renderer: sapien.core.Renderer, # shape: tuple, # length=0.1, # color1=[1, 1, 1, 1], # color2=[0, 0, 0, 1], # ) -> Tuple[sapien.core.RenderMesh, sapien.core.RenderMaterial]: # array = np.zeros([*shape, 4]) # checker = np.indices(shape).sum(axis=0) % 2 # array[checker == 0] = color1 # array[checker == 1] = color2 # array = (array * 255).clip(0, 255).astype(np.uint8) # texture = renderer.create_texture_from_array(array, filter_mode="nearest") # mat = renderer.create_material() # mat.set_diffuse_texture(texture) # mat.set_base_color([1, 0, 1, 1]) # width = length * shape[0] # height = length * shape[1] # mesh = renderer.create_mesh( # np.array( # [ # [0, 0, 0], # [width, 0, 0], # [width, height, 0], # [0, height, 0], # ], # dtype=np.float32, # ), # np.array([[0, 1, 2], [0, 2, 3]], dtype=np.uint32), # ) # mesh.set_uvs(np.array([[0, 0], [0, 1], [1, 1], [1, 0]], dtype=np.float32)) # mesh.set_normals( # np.array([[0, 0, 1], [0, 0, 1], [0, 0, 1], [0, 0, 1]], dtype=np.float32) # ) # return mesh, mat # def bake_envmap( # renderer: sapien.core.SapienRenderer, # scene: sapien.core.Scene, # position: np.ndarray, # resolution: int, # output_dir="sapien_generated_envmap", # near=0.1, # far=100.0, # ): # cam = scene.add_camera("cam", resolution, resolution, np.pi / 2, near, far) # poses = [ # sapien.core.Pose(position, [0.7071068, 0, 0, -0.7071068]), # sapien.core.Pose(position, [0.7071068, 0, 0, 0.7071068]), # sapien.core.Pose(position, [0, 0.7071068, 0, 0.7071068]), # sapien.core.Pose(position, [0, 0.7071068, 0, -0.7071068]), # sapien.core.Pose(position, [0, 0, 0, 1]), # sapien.core.Pose(position, [1, 0, 0, 0]), # +x direction is -z in opengl # ] # filenames = [ # os.path.join(output_dir, name + ".png") # for name in ["px", "nx", "py", "ny", "pz", "nz"] # ] # os.makedirs(output_dir, exist_ok=True) # from PIL import Image # for pose, filename in zip(poses, filenames): # cam.set_pose(pose) # scene.update_render() # cam.take_picture() # img = cam.get_texture("Color")[:, ::-1] # Image.fromarray((np.clip(img, 0, 1) * 255).astype(np.uint8)).save(filename) # renderer.create_ktx_environment_map( # *filenames, os.path.join(output_dir, "envmap.ktx") # ) # scene.remove_camera(cam)