Source code for concepts.simulator.cad_asset_utils.shapenet
#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File : shapenet.py
# Author : Jiayuan Mao
# Email : maojiayuan@gmail.com
# Date : 08/1/2024
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.
import os.path as osp
import jacinle.io as io
from jacinle.logging import get_logger
from concepts.simulator.cad_asset_utils.assets_base import CADAssetCollection
logger = get_logger(__file__)
[docs]
class ShapeNetCoreCollection(CADAssetCollection):
[docs]
def __init__(self, shapenet_root: str):
self.shapenet_root = shapenet_root
self.available_models = self.load_available_models()
[docs]
def load_available_models(self):
taxonomy_file = osp.join(self.shapenet_root, 'taxonomy.json')
taxonomy_data = io.load(taxonomy_file)
all_model_names = dict()
for synset in taxonomy_data:
synset_id = synset['synsetId']
if osp.isdir(osp.join(self.shapenet_root, synset_id)):
identifier = synset['name'].split(',')[0]
all_model_names[identifier] = dict(synset_id=synset_id, full_name=synset['name'], available_models=list())
for model_dir in io.lsdir(osp.join(self.shapenet_root, synset_id), '*'):
if osp.isdir(model_dir):
all_model_names[identifier]['available_models'].append(osp.basename(model_dir))
return all_model_names
[docs]
def build_urdf(self, model_dir: str, urdf_dir: str, identifier: str, model_identifier: str) -> str:
"""Build URDF for the specific model."""
io.mkdir(urdf_dir)
for file in io.lsdir(model_dir, '*', return_type='full'):
if osp.isfile(file):
io.link(file, osp.join(urdf_dir, osp.basename(file)), use_relative_path=True)
if not osp.isfile(osp.join(urdf_dir, 'model_normalized.obj')):
raise FileExistsError(f'No normalized obj found for {identifier}/{model_identifier}')
object_builder = self.make_urdf_builder(urdf_dir)
object_builder.build_urdf(osp.join(urdf_dir, 'model_normalized.obj'), force_overwrite=True, decompose_concave=True, force_decompose=False, center='mass')
return osp.join(urdf_dir, 'model_normalized.obj.urdf')
[docs]
def get_urdf(self, identifier: str, model_identifier: str) -> str:
synset_id = self.available_models[identifier]['synset_id']
asset_dir = osp.join(self.shapenet_root, synset_id, model_identifier)
model_dir = osp.join(asset_dir, 'models')
urdf_dir = osp.join(asset_dir, 'urdf')
if osp.isdir(urdf_dir):
if osp.isfile(osp.join(urdf_dir, 'model_normalized.obj.urdf')):
return osp.join(urdf_dir, 'model_normalized.obj.urdf')
logger.warning(f'No normalized urdf found for {identifier}/{model_identifier}, re-building one...')
io.remove(osp.join(model_dir, 'urdf'))
return self.build_urdf(model_dir, urdf_dir, identifier, model_identifier)