Skip to content

Commit

Permalink
Solve the legacy issue: support arbitrary file names
Browse files Browse the repository at this point in the history
  • Loading branch information
@yunshengtian
yunshengtian committed Mar 18, 2024
1 parent ca03563 commit 29db7df
Show file tree
Hide file tree
Showing 14 changed files with 87 additions and 157 deletions.
4 changes: 2 additions & 2 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ This repository contains the official code and dataset of [Assemble Them All: Ph

**News**:

- `translation.json` is no longer required for your assembly assets as long as the meshes are already in assembled states (03/2024).
- Easier input preparation for your custom assemblies: `translation.json` is no longer required for your assembly assets as long as the meshes are already in assembled states, and file names can be arbitrary (do not have to be integer anymore) (03/2024).

- We provide instructions [here](https://github.com/yunshengtian/Assemble-Them-All#custom-assembly-usage) for applying the algorithm on your custom meshes (03/2023).

Expand Down Expand Up @@ -216,7 +216,7 @@ We provide a pre-processing script `assets/process_mesh.py` that rescales the cu
python assets/process_mesh.py --source-dir source_dir/ --target-dir target_dir/ --subdivide
```

The output directory `target_dir/` will contain: 1) all the pre-processed meshes with names `0.obj`, `1.obj`, ..., `n.obj`; 2) `id_map.json` that stores the mappings to the names of original meshes under `source_dir/`.
Then, all the pre-processed meshes will be written to the output directory `target_dir/`.

To run our algorithm on top of them, just specify `--dir` and `--id` accordingly for the scripts in the above [Experiments section](https://github.com/yunshengtian/Assemble-Them-All#experiments) to make sure they can find `target_dir/`.

Expand Down
45 changes: 18 additions & 27 deletions assets/color.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,31 +5,22 @@
import numpy as np


def get_joint_color(index, normalize=True):
'''
Get color for 2-part assembly
'''
index = int(index)
assert index in [0, 1]
colors = np.array([
[107, 166, 161, 255],
[209, 184, 148, 255],
], dtype=int)
def get_color(part_ids, normalize=True):
color_map = {}
if len(part_ids) <= 2:
colors = np.array([
[107, 166, 161, 255],
[209, 184, 148, 255],
], dtype=int)
else:
colors = np.array([
[210, 87, 89, 255],
[237, 204, 73, 255],
[60, 167, 221, 255],
[190, 126, 208, 255],
[108, 192, 90, 255],
], dtype=int)
if normalize: colors = colors.astype(float) / 255.0
return colors[int(index)]


def get_multi_color(index, normalize=True):
'''
Get color for multi-part assembly
'''
index = int(index)
colors = np.array([
[210, 87, 89, 255],
[237, 204, 73, 255],
[60, 167, 221, 255],
[190, 126, 208, 255],
[108, 192, 90, 255],
], dtype=int)
if normalize: colors = colors.astype(float) / 255.0
return colors[index % 5]
for i, part_id in enumerate(part_ids):
color_map[part_id] = colors[i % len(colors)]
return color_map
16 changes: 8 additions & 8 deletions assets/load.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
import json
import trimesh

from assets.color import get_joint_color, get_multi_color
from assets.color import get_color
from assets.transform import get_transform_matrix, transform_pts_by_state


Expand All @@ -28,9 +28,9 @@ def load_translation(obj_dir, rotvec=None):
coms = json.load(fp)
new_coms = {}
for key, val in coms.items():
new_coms[int(key)] = np.array(val)
new_coms[key] = np.array(val)
if rotvec is not None:
new_coms[int(key)] = transform_pts_by_state(new_coms[int(key)], np.concatenate([np.zeros(3), rotvec]))
new_coms[key] = transform_pts_by_state(new_coms[key], np.concatenate([np.zeros(3), rotvec]))
coms = new_coms
return coms

Expand Down Expand Up @@ -68,7 +68,7 @@ def load_part_ids(obj_dir):
for file_name in os.listdir(obj_dir):
if file_name.endswith('.obj'):
try:
obj_id = int(file_name.replace('.obj', ''))
obj_id = file_name.replace('.obj', '')
except:
continue
obj_ids.append(obj_id)
Expand All @@ -88,9 +88,9 @@ def load_assembly(obj_dir, translate=True, rotvec=None, return_names=False):
coms = None

obj_ids = load_part_ids(obj_dir)
get_color = get_joint_color if len(obj_ids) <= 2 else get_multi_color
color_map = get_color(obj_ids, normalize=False)

for seq, obj_id in enumerate(obj_ids):
for obj_id in obj_ids:
obj_name = f'{obj_id}.obj'
obj_path = os.path.join(obj_dir, obj_name)
mesh = trimesh.load_mesh(obj_path, process=False, maintain_order=True)
Expand All @@ -105,7 +105,7 @@ def load_assembly(obj_dir, translate=True, rotvec=None, return_names=False):
coms[obj_id] = transform_pts_by_state(coms[obj_id], np.concatenate([np.zeros(3), rotvec]))
com_transform = com_to_transform(coms[obj_id])
mesh.apply_transform(com_transform)
mesh.visual.face_colors = get_color(seq, normalize=False)
mesh.visual.face_colors = color_map[obj_id]
meshes.append(mesh)
names.append(obj_name)

Expand All @@ -121,7 +121,7 @@ def load_paths(path_dir):
'''
paths = {}
for step in os.listdir(path_dir):
obj_id = int(step.split('_')[1])
obj_id = step.split('_')[1]
step_dir = os.path.join(path_dir, step)
if os.path.isdir(step_dir):
path = []
Expand Down
16 changes: 4 additions & 12 deletions assets/process_mesh.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ def get_oriented_bounding_box(mesh):
return extents


def process_mesh(source_dir, target_dir, subdivide, max_edge=0.5, keep_id=False, verbose=False):
def process_mesh(source_dir, target_dir, subdivide, max_edge=0.5, verbose=False):
'''
1. Check watertight
2. Scale to unit bounding box
Expand All @@ -67,11 +67,8 @@ def process_mesh(source_dir, target_dir, subdivide, max_edge=0.5, keep_id=False,
obj_ids = {}
watertight = True
for i, source_file in enumerate(source_files):
if keep_id:
source_id = int(source_file.replace('.obj', ''))
obj_ids[source_id] = source_file
else:
obj_ids[i] = source_file
source_id = source_file.replace('.obj', '')
obj_ids[source_id] = source_file
source_path = os.path.join(source_dir, source_file)
mesh = trimesh.load_mesh(source_path, process=False, maintain_order=True)
if not mesh.is_watertight:
Expand Down Expand Up @@ -101,10 +98,6 @@ def process_mesh(source_dir, target_dir, subdivide, max_edge=0.5, keep_id=False,
if verbose:
print(f'Processed obj written to {obj_target_path}')

# save obj id map
with open(os.path.join(target_dir, 'id_map.json'), 'w') as fp:
json.dump(obj_ids, fp)

return True


Expand All @@ -116,8 +109,7 @@ def process_mesh(source_dir, target_dir, subdivide, max_edge=0.5, keep_id=False,
parser.add_argument('--target-dir', type=str, required=True)
parser.add_argument('--subdivide', default=False, action='store_true')
parser.add_argument('--max-edge', type=float, default=0.5)
parser.add_argument('--keep-id', default=False, action='store_true')
args = parser.parse_args()

success = process_mesh(args.source_dir, args.target_dir, args.subdivide, max_edge=args.max_edge, keep_id=args.keep_id, verbose=True)
success = process_mesh(args.source_dir, args.target_dir, args.subdivide, max_edge=args.max_edge, verbose=True)
print(f'Success: {success}')
15 changes: 8 additions & 7 deletions baselines/run_joint_plan.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,9 @@
from pyplanners.targetless_rrt import targetless_rrt
from pyplanners.smoothing import smooth_path

from assets.load import load_assembly
from assets.load import load_assembly, load_part_ids
from assets.save import interpolate_path, save_path
from assets.color import get_joint_color
from assets.color import get_color
from assets.transform import transform_pts_by_state, get_transform_matrix
from assets.mesh_distance import compute_move_mesh_distance

Expand All @@ -38,19 +38,21 @@ def __init__(self, assembly_dir, move_id, still_ids, rotation=False, max_collisi
# load data
meshes, names = load_assembly(assembly_dir, return_names=True)
self.move_id, self.still_ids = move_id, still_ids
part_ids = load_part_ids(assembly_dir)
color_map = get_color(part_ids, normalize=False)

# build meshes
self.viz_mesh_move, self.viz_meshes_still = None, []
self.mesh_move, self.meshes_still = None, []

for mesh, name in zip(meshes, names):
mesh_id = int(name.replace('.obj', ''))
mesh_id = name.replace('.obj', '')
mesh.visual.face_colors = color_map[mesh_id]

sdf_load_path = os.path.join(assembly_dir, name.replace('.obj', '.sdf'))
sdf_save_path = sdf_load_path if save_sdf else ''

if mesh_id == move_id:
mesh.visual.face_colors = get_joint_color(0, normalize=False)
self.viz_mesh_move = mesh
if body_type == 'bvh':
self.mesh_move = redmax.BVHMesh(mesh.vertices.T, mesh.faces.T)
Expand All @@ -60,7 +62,6 @@ def __init__(self, assembly_dir, move_id, still_ids, rotation=False, max_collisi
raise NotImplementedError

elif mesh_id in still_ids:
mesh.visual.face_colors = get_joint_color(1, normalize=False)
self.viz_meshes_still.append(mesh)
if body_type == 'bvh':
self.meshes_still.append(redmax.BVHMesh(mesh.vertices.T, mesh.faces.T))
Expand Down Expand Up @@ -281,8 +282,8 @@ def plan(self, planner_name, step_size, max_time, seed=1, return_path=False, sim
parser = ArgumentParser()
parser.add_argument('--dir', type=str, default='joint_assembly', help='directory storing all assemblies')
parser.add_argument('--id', type=str, required=True, help='assembly id (e.g. 00000)')
parser.add_argument('--move-id', type=int, default=0)
parser.add_argument('--still-ids', type=int, nargs='+', default=[1])
parser.add_argument('--move-id', type=str, default='0')
parser.add_argument('--still-ids', type=str, nargs='+', default=['1'])
parser.add_argument('--rotation', default=False, action='store_true')
parser.add_argument('--planner', type=str, required=True, choices=['rrt', 'rrt-connect', 'birrt', 'trrt', 'matevec-trrt'])
parser.add_argument('--max-collision', type=float, default=1e-2)
Expand Down
4 changes: 2 additions & 2 deletions baselines/run_joint_plan_batch.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,8 @@ def log_results(log_dir, result_status, verbose=False):

parser = ArgumentParser()
parser.add_argument('--dir', type=str, default='joint_assembly')
parser.add_argument('--move-id', type=int, default=0)
parser.add_argument('--still-ids', type=int, nargs='+', default=[1])
parser.add_argument('--move-id', type=str, default='0')
parser.add_argument('--still-ids', type=str, nargs='+', default=['1'])
parser.add_argument('--rotation', default=False, action='store_true')
parser.add_argument('--planner', type=str, required=True, choices=['rrt', 'rrt-connect', 'birrt', 'trrt', 'matevec-trrt'])
parser.add_argument('--max-collision', type=float, default=1e-2)
Expand Down
6 changes: 3 additions & 3 deletions baselines/run_multi_plan.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ def __init__(self, assembly_dir):

meshes, names = load_assembly(assembly_dir, return_names=True)

part_ids = [int(name.replace('.obj', '')) for name in names]
part_ids = [name.replace('.obj', '') for name in names]
for i in range(len(part_ids)):
self.graph.add_node(part_ids[i])

Expand Down Expand Up @@ -111,7 +111,7 @@ def plan_sequence(self, path_planner_name,

if status in self.success_status:
self.graph.remove_node(move_id)
sequence.append(int(move_id))
sequence.append(move_id)

if len(self.graph.nodes) == 1:
seq_status = 'Success'
Expand Down Expand Up @@ -178,7 +178,7 @@ def plan_sequence(self, path_planner_name,

if status in self.success_status:
self.graph.remove_node(move_id)
sequence.append(int(move_id))
sequence.append(move_id)
else:
inactive_queue.append(move_id)

Expand Down
15 changes: 8 additions & 7 deletions examples/run_joint_plan.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

from assets.load import load_translation, load_assembly, load_part_ids
from assets.save import save_path, clear_saved_sdfs
from assets.color import get_joint_color, get_multi_color
from assets.color import get_color
from assets.transform import transform_pts_by_state
from assets.mesh_distance import compute_move_mesh_distance
from utils.renderer import SimRenderer
Expand Down Expand Up @@ -138,10 +138,11 @@ def arr_to_str(arr):


def get_xml_string(assembly_dir, move_id, still_ids, move_joint_type, body_type, sdf_dx, col_th, save_sdf):
part_ids = load_part_ids(assembly_dir)
translation = load_translation(assembly_dir)
if translation is None: translation = {part_id: [0, 0, 0] for part_id in load_part_ids(assembly_dir)}
if translation is None: translation = {part_id: [0, 0, 0] for part_id in part_ids}
body_type = body_type.upper()
get_color = get_joint_color if len(translation.keys()) <= 2 else get_multi_color
color_map = get_color(part_ids)
sdf_args = 'load_sdf="true" save_sdf="true"' if save_sdf else ''
string = f'''
<redmax model="assemble">
Expand All @@ -158,7 +159,7 @@ def get_xml_string(assembly_dir, move_id, still_ids, move_joint_type, body_type,
<robot>
<link name="part{part_id}">
<joint name="part{part_id}" type="{joint_type}" axis="0. 0. 0." pos="{arr_to_str(translation[part_id])}" quat="1 0 0 0" frame="WORLD" damping="0"/>
<body name="part{part_id}" type="{body_type}" filename="{assembly_dir}/{part_id}.obj" {sdf_args} pos="0 0 0" quat="1 0 0 0" scale="1 1 1" transform_type="OBJ_TO_JOINT" density="1" dx="{sdf_dx}" col_th="{col_th}" mu="0" rgba="{arr_to_str(get_color(part_id))}"/>
<body name="part{part_id}" type="{body_type}" filename="{assembly_dir}/{part_id}.obj" {sdf_args} pos="0 0 0" quat="1 0 0 0" scale="1 1 1" transform_type="OBJ_TO_JOINT" density="1" dx="{sdf_dx}" col_th="{col_th}" mu="0" rgba="{arr_to_str(color_map[part_id])}"/>
</link>
</robot>
'''
Expand Down Expand Up @@ -196,7 +197,7 @@ def __init__(self, asset_folder, assembly_dir, move_id, still_ids,
move_mesh = None
still_meshes = []
for mesh, name in zip(meshes, names):
obj_id = int(name.replace('.obj', ''))
obj_id = name.replace('.obj', '')
if body_type == 'bvh':
phys_mesh = redmax.BVHMesh(mesh.vertices.T, mesh.faces.T)
elif body_type == 'sdf':
Expand Down Expand Up @@ -699,8 +700,8 @@ def get_planner(name):
parser.add_argument('--planner', type=str, required=True, choices=['bfs', 'bk-rrt'])
parser.add_argument('--id', type=str, required=True, help='assembly id (e.g. 00000)')
parser.add_argument('--dir', type=str, default='joint_assembly', help='directory storing all assemblies')
parser.add_argument('--move-id', type=int, default=0)
parser.add_argument('--still-ids', type=int, nargs='+', default=[1])
parser.add_argument('--move-id', type=str, default='0')
parser.add_argument('--still-ids', type=str, nargs='+', default=['1'])
parser.add_argument('--rotation', default=False, action='store_true')
parser.add_argument('--body-type', type=str, default='sdf', choices=['bvh', 'sdf'], help='simulation type of body')
parser.add_argument('--sdf-dx', type=float, default=0.05, help='grid resolution of SDF')
Expand Down
4 changes: 2 additions & 2 deletions examples/run_joint_plan_batch.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,8 +56,8 @@ def log_results(log_dir, result_status, verbose=False):
parser = ArgumentParser()
parser.add_argument('--planner', type=str, required=True, choices=['bfs', 'bk-rrt'])
parser.add_argument('--dir', type=str, default='joint_assembly', help='directory storing all assemblies')
parser.add_argument('--move-id', type=int, default=0)
parser.add_argument('--still-ids', type=int, nargs='+', default=[1])
parser.add_argument('--move-id', type=str, default='0')
parser.add_argument('--still-ids', type=str, nargs='+', default=['1'])
parser.add_argument('--rotation', default=False, action='store_true')
parser.add_argument('--body-type', type=str, default='sdf', choices=['bvh', 'sdf'], help='simulation type of body')
parser.add_argument('--sdf-dx', type=float, default=0.05, help='grid resolution of SDF')
Expand Down
8 changes: 4 additions & 4 deletions examples/run_multi_plan.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ def __init__(self, asset_folder, assembly_dir):

meshes, names = load_assembly(assembly_dir, return_names=True)

part_ids = [int(name.replace('.obj', '')) for name in names]
part_ids = [name.replace('.obj', '') for name in names]
for i in range(len(part_ids)):
self.graph.add_node(part_ids[i])

Expand Down Expand Up @@ -119,7 +119,7 @@ def plan_sequence(self, path_planner_name, rotation, body_type, sdf_dx, collisio

if status in self.success_status:
self.graph.remove_node(move_id)
sequence.append(int(move_id))
sequence.append(move_id)

if len(self.graph.nodes) == 1:
seq_status = 'Success'
Expand Down Expand Up @@ -193,7 +193,7 @@ def plan_sequence(self, path_planner_name, rotation, body_type, sdf_dx, collisio

if status in self.success_status:
self.graph.remove_node(move_id)
sequence.append(int(move_id))
sequence.append(move_id)
else:
inactive_queue.append(move_id)

Expand Down Expand Up @@ -300,7 +300,7 @@ def plan_sequence(self, path_planner_name, rotation, body_type, sdf_dx, collisio

if status in self.success_status:
self.graph.remove_node(move_id)
sequence.append(int(move_id))
sequence.append(move_id)
else:
inactive_queue.append([move_id, max_depth + 1])

Expand Down
Loading

0 comments on commit 29db7df

Please sign in to comment.