shrec11_dataset.py
import shutil
import os
import sys
import random
import numpy as np
import torch
from torch.utils.data import Dataset
import potpourri3d as pp3d
sys.path.append(os.path.join(os.path.dirname(__file__), "../../src/"))
import fannet
from fannet.utils import toNP
class Shrec11MeshDataset_Original(Dataset):
# NOTE: Original data from the challenge, not simplified models
# The original SHREC11 models were previously distributed via NIST [here](https://www.nist.gov/itl/iad/shrec-2011-datasets), but that page seems to have been lost to the sands of time. We provide a zip of the old dataset page here: https://drive.google.com/uc?export=download&id=1O_P03aAxhjCOKQH2n71j013-EfSmEp5e. The relevant files are in `SHREC11_test_database_new.zip`, which is password protected with the password `SHREC11@NIST`.
# Unzip it like
# unzip -P SHREC11@NIST SHREC11_test_database_new.zip -d [DATA_ROOT]/raw
def __init__(self, root_dir, split_size, k_eig, exclude_dict=None, op_cache_dir=None, spoke_length=[0.00,0.01,0.02,0.03]):
self.root_dir = root_dir
self.n_class = 30
self.split_size = split_size # pass None to take all entries (except those in exclude_dict)
self.k_eig = k_eig
self.op_cache_dir = op_cache_dir
self.spoke_length = spoke_length
self.class_names = []
self.entries = {}
# store in memory
self.verts_list = []
self.faces_list = []
self.labels_list = []
raw_path = os.path.join(self.root_dir, 'raw')
## Parse the categories file
cat_path = os.path.join(self.root_dir, 'categories.txt')
with open(cat_path) as cat_file:
cat_file.readline() # skip the first two lines
cat_file.readline()
for i_class in range(30):
cat_file.readline()
class_name, _, count = cat_file.readline().strip().split()
count = int(count)
self.class_names.append(class_name)
mesh_list = []
for j in range(20):
mesh_list.append(cat_file.readline().strip())
# Randomly grab samples for this split. If given, disallow any samples in commmon with exclude_dict (ie making sure train set is distinct from test).
order = np.random.permutation(len(mesh_list))
added = 0
self.entries[class_name] = set()
for ind in order:
if(split_size is not None and added == split_size): continue
name = mesh_list[ind]
if exclude_dict is not None and name in exclude_dict[class_name]:
continue
path = os.path.join(root_dir, "raw", "T{}.off".format(name))
verts, faces = pp3d.read_mesh(path)
verts = torch.tensor(verts).float()
faces = torch.tensor(faces)
#print(verts)
#print(faces)
# center and unit scale
verts = fannet.geometry.normalize_positions(verts)
self.verts_list.append(verts)
self.faces_list.append(faces)
self.labels_list.append(i_class)
self.entries[class_name].add(name)
added += 1
print(class_name + " -- " + " ".join([p for p in self.entries[class_name]]))
if(split_size is not None and added < split_size):
raise ValueError("could not find enough entries to generate requested split")
for ind, label in enumerate(self.labels_list):
self.labels_list[ind] = torch.tensor(label)
def __len__(self):
return len(self.verts_list)
def __getitem__(self, idx):
verts = self.verts_list[idx]
faces = self.faces_list[idx]
label = self.labels_list[idx]
frames, mass, L, evals, evecs, gradX, gradY, fanIndices = fannet.geometry.get_operators(verts, faces, k_eig=self.k_eig, op_cache_dir=self.op_cache_dir, spoke_length=self.spoke_length)
return verts, faces, frames, mass, L, evals, evecs, gradX, gradY, label , fanIndices
class Shrec11MeshDataset_Simplified(Dataset):
# NOTE: Remeshed data from MeshCNN authors.
# Can be downloaded here (link from the MeshCNN authors): https://www.dropbox.com/s/w16st84r6wc57u7/shrec_16.tar.gz. Note that despite the filename, this really is the shapes from the SHREC 2011 dataset. Extract it to the `[ROOT_DIR]/raw/` directory.
def __init__(self, root_dir, split_size, k_eig, exclude_dict=None, op_cache_dir=None, spoke_length=[0.00,0.01,0.02,0.03]):
self.root_dir = root_dir
self.n_class = 30
self.split_size = split_size # pass None to take all entries (except those in exclude_dict)
self.k_eig = k_eig
self.op_cache_dir = op_cache_dir
self.spoke_length = spoke_length
self.class_names = [ 'alien', 'ants', 'armadillo', 'bird1', 'bird2', 'camel', 'cat', 'centaur', 'dinosaur', 'dino_ske', 'dog1', 'dog2', 'flamingo', 'glasses', 'gorilla', 'hand', 'horse', 'lamp', 'laptop', 'man', 'myScissor', 'octopus', 'pliers', 'rabbit', 'santa', 'shark', 'snake', 'spiders', 'two_balls', 'woman']
self.entries = {}
# store in memory
self.verts_list = []
self.faces_list = []
self.labels_list = []
raw_path = os.path.join(self.root_dir, 'raw', "shrec_16")
for class_idx, class_name in enumerate(self.class_names):
# load both train and test subdirectories; we are manually regenerating random splits to do multiple trials
mesh_files = []
for t in ['test', 'train']:
files = os.listdir(os.path.join(raw_path, class_name, t))
for f in files:
full_f = os.path.join(raw_path, class_name, t, f)
mesh_files.append(full_f)
# Randomly grab samples for this split. If given, disallow any samples in commmon with exclude_dict (ie making sure train set is distinct from test).
order = np.random.permutation(len(mesh_files))
added = 0
self.entries[class_name] = set()
for ind in order:
if(split_size is not None and added == split_size): continue
path = mesh_files[ind]
if exclude_dict is not None and path in exclude_dict[class_name]:
continue
verts, faces = pp3d.read_mesh(path)
verts = torch.tensor(verts).float()
faces = torch.tensor(faces)
#print(verts.shape)
#print(faces.shape)
# center and unit scale
verts = fannet.geometry.normalize_positions(verts)
self.verts_list.append(verts)
self.faces_list.append(faces)
self.labels_list.append(class_idx)
self.entries[class_name].add(path)
added += 1
print(class_name + " -- " + " ".join([os.path.basename(p) for p in self.entries[class_name]]))
if(split_size is not None and added < split_size):
raise ValueError("could not find enough entries to generate requested split")
for ind, label in enumerate(self.labels_list):
self.labels_list[ind] = torch.tensor(label)
# Precompute operators
self.frames_list, self.massvec_list, self.L_list, self.evals_list, self.evecs_list, self.gradX_list, self.gradY_list, self.fanIndices = fannet.geometry.get_all_operators(self.verts_list,
self.faces_list, k_eig=self.k_eig, op_cache_dir=self.op_cache_dir, spoke_length=self.spoke_length)
def __len__(self):
return len(self.verts_list)
def __getitem__(self, idx):
return self.verts_list[idx], self.faces_list[idx], self.frames_list[idx], self.massvec_list[idx], self.L_list[idx], self.evals_list[idx], self.evecs_list[idx], self.gradX_list[idx], self.gradY_list[idx], self.labels_list[idx], self.fanIndices[idx]
