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)