Skip to content

Commit

Permalink
Refactored ModelNet dataset to use KaolinDataset (#189)
Browse files Browse the repository at this point in the history 
* Formatting and fixes

Signed-off-by: TommyX12 <tommyx058@gmail.com>

* Initial pass of refactoring ModelNet dataset

Signed-off-by: TommyX12 <tommyx058@gmail.com>

* Fixed ModelNet tests

Signed-off-by: TommyX12 <tommyx058@gmail.com>

* Fixed minor bugs

Signed-off-by: TommyX12 <tommyx058@gmail.com>

* Bug fix

Signed-off-by: TommyX12 <tommyx058@gmail.com>
  • Loading branch information
@TommyX12
TommyX12 committed Apr 8, 2020
1 parent c8537a7 commit 19796d3
Show file tree
Hide file tree
Showing 5 changed files with 48 additions and 148 deletions.
1 change: 0 additions & 1 deletion .flake8
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ exclude = .git, tests/, build/,
examples/renderers/NMR,
examples/SuperResolution,
kaolin/cuda,
kaolin/datasets/modelnet.py,
kaolin/datasets/scannet.py,
kaolin/datasets/shapenet.py,
kaolin/datasets/shrec.py,
Expand Down
141 changes: 32 additions & 109 deletions kaolin/datasets/modelnet.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,63 +12,66 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Callable, Iterable, Optional, Union, List
from typing import Iterable, Optional

import torch
import os
from glob import glob
from tqdm import tqdm

from kaolin.rep.TriangleMesh import TriangleMesh
from kaolin.transforms import transforms as tfs
from .base import KaolinDataset


class ModelNet(object):
r""" Dataset class for the ModelNet dataset.
class ModelNet(KaolinDataset):
r"""Dataset class for the ModelNet dataset.
Args:
basedir (str): Path to the base directory of the ModelNet dataset.
root (str): Path to the base directory of the ModelNet dataset.
split (str, optional): Split to load ('train' vs 'test',
default: 'train').
categories (iterable, optional): List of categories to load
(default: ['chair']).
transform (callable, optional): A function/transform to apply on each
loaded example.
device (str or torch.device, optional): Device to use (cpu,
cuda, cuda:1, etc.). Default: 'cpu'
Examples:
>>> dataset = ModelNet(basedir='data/ModelNet')
>>> dataset = ModelNet(root='data/ModelNet')
>>> train_loader = DataLoader(dataset, batch_size=10, shuffle=True, num_workers=8)
>>> obj, label = next(iter(train_loader))
"""

def __init__(self, basedir: str,
split: Optional[str] = 'train',
categories: Optional[Iterable] = ['bed'],
transform: Optional[Callable] = None,
device: Optional[Union[torch.device, str]] = 'cpu'):
def initialize(self, root: str,
split: Optional[str] = 'train',
categories: Optional[Iterable] = None):
"""Initialize the dataset.
Args:
root (str): Path to the base directory of the ModelNet dataset.
split (str, optional): Split to load ('train' vs 'test',
default: 'train').
categories (iterable, optional): List of categories to load
(default: ['chair']).
"""

assert split.lower() in ['train', 'test']

self.basedir = basedir
self.transform = transform
self.device = device
if categories is None:
categories = ['chair']

self.root = root
self.categories = categories
self.names = []
self.filepaths = []
self.cat_idxs = []

if not os.path.exists(basedir):
raise ValueError('ModelNet was not found at "{0}".'.format(basedir))
if not os.path.exists(root):
raise ValueError('ModelNet was not found at "{0}".'.format(root))

available_categories = [p for p in os.listdir(basedir) if os.path.isdir(os.path.join(basedir, p))]
available_categories = [p for p in os.listdir(root) if os.path.isdir(os.path.join(root, p))]

for cat_idx, category in enumerate(categories):
assert category in available_categories, 'object class {0} not in list of available classes: {1}'.format(
category, available_categories)

cat_paths = glob(os.path.join(basedir, category, split.lower(), '*.off'))
cat_paths = glob(os.path.join(root, category, split.lower(), '*.off'))

self.cat_idxs += [cat_idx] * len(cat_paths)
self.names += [os.path.splitext(os.path.basename(cp))[0] for cp in cat_paths]
Expand All @@ -77,92 +80,12 @@ def __init__(self, basedir: str,
def __len__(self):
return len(self.names)

def __getitem__(self, index):
"""Returns the item at index idx. """
category = torch.tensor(self.cat_idxs[index], dtype=torch.long, device=self.device)
def _get_data(self, index):
data = TriangleMesh.from_off(self.filepaths[index])
data.to(self.device)
if self.transform:
data = self.transform(data)

return data, category


class ModelNetVoxels(object):
r""" Dataloader for downloading and reading from ModelNet.
return data

Args:
basedir (str): location the dataset should be downloaded to /loaded from
cache_dir (str, optional)
split (str, optional): Split to load ('train' vs 'test',
default: 'train').
categories (str, optional): list of object classes to be loaded
resolutions (list of int, optional): list of voxel grid resolutions to create,
default: [32]
device (str or torch.device, optional): Device to use (cpu,
cuda, cuda:1, etc.). Default: 'cpu'
Returns:
.. code-block::
dict: {
'attributes': {'name': str, 'class': str},
'data': {'voxels': torch.Tensor}
def _get_attributes(self, index):
category = torch.tensor(self.cat_idxs[index], dtype=torch.long)
return {
'category': category,
}
Examples:
>>> dataset = ModelNet(basedir='data/ModelNet', resolutions=[32])
>>> train_loader = DataLoader(dataset, batch_size=10, shuffle=True, num_workers=8)
>>> obj = next(iter(train_loader))
>>> obj['data']['32'].size()
torch.Size(32, 32, 32)
"""

def __init__(self, basedir: str, cache_dir: Optional[str] = None,
split: Optional[str] = 'train', categories: list = ['bed'],
resolutions: List[int] = [32],
device: Optional[Union[torch.device, str]] = 'cpu'):

self.basedir = basedir
self.device = torch.device(device)
self.cache_dir = cache_dir if cache_dir is not None else os.path.join(basedir, 'cache')
self.params = {'resolutions': resolutions}
self.cache_transforms = {}

mesh_dataset = ModelNet(basedir=basedir, split=split, categories=categories, device=device)

self.names = mesh_dataset.names
self.categories = mesh_dataset.categories
self.cat_idxs = mesh_dataset.cat_idxs

for res in self.params['resolutions']:
self.cache_transforms[res] = tfs.CacheCompose([
tfs.TriangleMeshToVoxelGrid(res, normalize=True, vertex_offset=0.5),
tfs.FillVoxelGrid(thresh=0.5),
tfs.ExtractProjectOdmsFromVoxelGrid()
], self.cache_dir)

desc = 'converting to voxels to resolution {0}'.format(res)
for idx in tqdm(range(len(mesh_dataset)), desc=desc, disable=False):
name = mesh_dataset.names[idx]
if name not in self.cache_transforms[res].cached_ids:
mesh, _ = mesh_dataset[idx]
mesh.to(device=device)
self.cache_transforms[res](name, mesh)

def __len__(self):
return len(self.names)

def __getitem__(self, index):
"""Returns the item at index idx. """
data = dict()
attributes = dict()
name = self.names[index]

for res in self.params['resolutions']:
data[str(res)] = self.cache_transforms[res](name)
attributes['name'] = name
attributes['category'] = self.categories[self.cat_idxs[index]]
return {'data': data, 'attributes': attributes}
5 changes: 4 additions & 1 deletion kaolin/datasets/shapenet.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,7 @@ def __getitem__(self, index):

class ShapeNet(KaolinDataset):
r"""ShapeNetV1 Dataset class for meshes.
Args:
root (str): path to ShapeNet root directory
categories (list): List of categories to load from ShapeNet. This list may
Expand All @@ -176,6 +177,7 @@ class ShapeNet(KaolinDataset):
attributes: {name: str, path: str, synset: str, label: str},
data: {vertices: torch.Tensor, faces: torch.Tensor}
}
Example:
>>> meshes = ShapeNet(root='../data/ShapeNet/')
>>> obj = meshes[0]
Expand All @@ -186,7 +188,8 @@ class ShapeNet(KaolinDataset):
"""

def initialize(self, root: str, categories: list, train: bool = True, split: float = .7):
"""Initialize the dataset
"""Initialize the dataset.
Args:
root (str): path to ShapeNet root directory
categories (list): List of categories to load from ShapeNet. This list may
Expand Down
44 changes: 8 additions & 36 deletions tests/datasets/test_ModelNet.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,48 +21,20 @@
import kaolin.transforms.transforms as tfs


MODELNET_ROOT = 'data/ModelNet10/'
MODELNET_ROOT = '/data/ModelNet10/'
CACHE_DIR = 'tests/datasets/cache'


# Tests below can only be run is a ShapeNet dataset is available
REASON = 'ShapeNet not found at default location: {}'.format(MODELNET_ROOT)
# Tests below can only be run if a ModelNet dataset is available
REASON = 'ModelNet not found at default location: {}'.format(MODELNET_ROOT)


@pytest.mark.parametrize('device', ['cpu', 'cuda'])
@pytest.mark.skipif(not os.path.exists(MODELNET_ROOT), reason=REASON)
def test_ModelNet(device):
models = kal.datasets.ModelNet(basedir=MODELNET_ROOT, categories=['bathtub'], split='test')
def test_ModelNet(device):
models = kal.datasets.ModelNet(root=MODELNET_ROOT, categories=['bathtub'], split='test')

assert len(models) == 50
for obj, category in models:
assert category.item() == 0
assert isinstance(obj, kal.rep.Mesh)


@pytest.mark.parametrize('device', ['cpu', 'cuda'])
@pytest.mark.skipif(not os.path.exists(MODELNET_ROOT), reason=REASON)
def test_ModelNetPointCloud(device):
transform = tfs.Compose([
tfs.TriangleMeshToPointCloud(num_samples=32),
tfs.NormalizePointCloud()
])
models = kal.datasets.ModelNet(basedir=MODELNET_ROOT, categories=['bathtub'], split='test', transform=transform)

assert len(models) == 50
for obj, category in models:
assert category.item() == 0
assert isinstance(obj, torch.Tensor)
assert obj.size(0) == 32


@pytest.mark.parametrize('device', ['cpu', 'cuda'])
@pytest.mark.skipif(not os.path.exists(MODELNET_ROOT), reason=REASON)
def test_ModelNetVoxels(device):
models = kal.datasets.ModelNetVoxels(basedir=MODELNET_ROOT, cache_dir=CACHE_DIR, categories=['bathtub'], split='test', resolutions=[30])

assert len(models) == 50
for obj in models:
assert obj['attributes']['category'] == 'bathtub'
assert set(obj['data']['30'].shape) == set([30, 30, 30])
shutil.rmtree(CACHE_DIR)
for item in models:
assert item['attributes']['category'].item() == 0
assert isinstance(item['data'], kal.rep.Mesh)
5 changes: 4 additions & 1 deletion tests/datasets/test_ShapeNet.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,12 @@
CACHE_DIR = 'tests/datasets/cache'


# Tests below can only be run is a ShapeNet dataset is available
# Tests below can only be run if a ShapeNet dataset is available
SHAPENET_NOT_FOUND = 'ShapeNet not found at default location: {}'.format(SHAPENET_ROOT)
SHAPENET_RENDERING_NOT_FOUND = 'ShapeNetRendering not found at default location: {}'.format(
SHAPENET_RENDERING_ROOT)


@pytest.mark.skipif(not Path(SHAPENET_ROOT).exists(), reason=SHAPENET_NOT_FOUND)
def test_Meshes():
meshes1 = shapenet.ShapeNet_Meshes(root=SHAPENET_ROOT,
Expand Down Expand Up @@ -65,6 +66,7 @@ def test_Voxels():

shutil.rmtree('tests/datasets/cache/voxels')


@pytest.mark.parametrize('categories', [['chair'], ['plane', 'bench', 'cabinet', 'car', 'chair',
'monitor', 'lamp', 'speaker', 'rifle',
'sofa', 'table', 'phone', 'watercraft']])
Expand All @@ -83,6 +85,7 @@ def test_Images(categories):
assert list(obj['data']['params']['cam_mat'].shape) == [3, 3]
assert list(obj['data']['params']['cam_pos'].shape) == [3]


@pytest.mark.skipif(not Path(SHAPENET_ROOT).exists(), reason=SHAPENET_NOT_FOUND)
def test_Surface_Meshes():
surface_meshes = shapenet.ShapeNet_Surface_Meshes(root=SHAPENET_ROOT, cache_dir=CACHE_DIR,
Expand Down

0 comments on commit 19796d3

Please sign in to comment.