File size: 3,068 Bytes
48f95a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import os
import json
import trimesh
import urllib.request
import zipfile
from .shapenet_data_pc import ShapeNet15kPointClouds
class ShapeNetV2Dataset(Dataset):
def __init__(self, root_dir, categories=None, num_points=2048, split='train'):
self.root_dir = root_dir
self.num_points = num_points
self.split = split
# Download and extract the dataset if it doesn't exist
if not os.path.exists(root_dir):
self.download_shapenet(root_dir)
# Load the taxonomy file
with open(os.path.join(root_dir, 'taxonomy.json'), 'r') as f:
self.taxonomy = json.load(f)
# Filter categories if specified
if categories is not None:
self.taxonomy = [t for t in self.taxonomy if t['synsetId'] in categories]
# Load the split file
split_file = os.path.join(root_dir, f'{split}.txt')
with open(split_file, 'r') as f:
self.models = [line.strip() for line in f]
self.category_to_synsetId = {t['name']: t['synsetId'] for t in self.taxonomy}
def __len__(self):
return len(self.models)
def __getitem__(self, idx):
model = self.models[idx]
synsetId, model_id = model.split('-')
# Load the point cloud
mesh_path = os.path.join(self.root_dir, synsetId, model_id, 'models', 'model_normalized.obj')
mesh = trimesh.load(mesh_path)
points = mesh.sample(self.num_points)
return torch.FloatTensor(points)
@staticmethod
def download_shapenet(root_dir):
url = "https://shapenet.cs.stanford.edu/shapenet/obj-zip/ShapeNetCore.v2.zip"
zip_path = os.path.join(root_dir, "ShapeNetCore.v2.zip")
# Create the directory if it doesn't exist
os.makedirs(root_dir, exist_ok=True)
# Download the dataset
print("Downloading ShapeNetV2 dataset...")
urllib.request.urlretrieve(url, zip_path)
# Extract the dataset
print("Extracting ShapeNetV2 dataset...")
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
zip_ref.extractall(root_dir)
# Remove the zip file
os.remove(zip_path)
print("ShapeNetV2 dataset downloaded and extracted successfully.")
def get_dataset(dataroot, npoints, category):
tr_dataset = ShapeNet15kPointClouds(root_dir=dataroot,
categories=category, split='train',
tr_sample_size=npoints,
te_sample_size=npoints,
scale=1.,
normalize_per_shape=False,
normalize_std_per_axis=False,
random_subsample=False)
te_dataset = ShapeNet15kPointClouds(root_dir=dataroot,
categories=category, split='val',
tr_sample_size=npoints,
te_sample_size=npoints,
scale=1.,
normalize_per_shape=False,
normalize_std_per_axis=False,
all_points_mean=tr_dataset.all_points_mean,
all_points_std=tr_dataset.all_points_std,
)
return tr_dataset, te_dataset
|