import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import os
import json
import trimesh
import urllib.request
import zipfile

if __name__ == "__main__":

    root_dir = None
    # url = "https://shapenet.cs.stanford.edu/shapenet/obj-zip/ShapeNetCore.v2.zip"
    zip_path = os.path.join(root_dir, "ShapeNetCore.v2.PC15k.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.")