🪐 Objaverse-XL is a Universe of 10M+ 3D Objects. Contains API Scripts for Downloading and Processing!
"""A package for downloading and processing Objaverse."""

import glob
import gzip
import json
import multiprocessing
import os
import urllib.request
import warnings
from typing import Any, Dict, List, Optional, Tuple

from tqdm import tqdm

BASE_PATH = os.path.join(os.path.expanduser("~"), ".objaverse")

__version__ = "<REPLACE_WITH_VERSION>"
_VERSIONED_PATH = os.path.join(BASE_PATH, "hf-objaverse-v1")


def load_annotations(uids: Optional[List[str]] = None) -> Dict[str, Any]:
    """Load the full metadata of all objects in the dataset.

    Args:
        uids: A list of uids with which to load metadata. If None, it loads
        the metadata for all uids.

    Returns:
        A dictionary mapping the uid to the metadata.
    """
    metadata_path = os.path.join(_VERSIONED_PATH, "metadata")
    object_paths = _load_object_paths()
    dir_ids = (
        set(object_paths[uid].split("/")[1] for uid in uids)
        if uids is not None
        else [f"{i // 1000:03d}-{i % 1000:03d}" for i in range(160)]
    )
    if len(dir_ids) > 10:
        dir_ids = tqdm(dir_ids)
    out = {}
    for i_id in dir_ids:
        json_file = f"{i_id}.json.gz"
        local_path = os.path.join(metadata_path, json_file)
        if not os.path.exists(local_path):
            hf_url = f"https://huggingface.co/datasets/allenai/objaverse/resolve/main/metadata/{i_id}.json.gz"
            # wget the file and put it in local_path
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            urllib.request.urlretrieve(hf_url, local_path)
        with gzip.open(local_path, "rb") as f:
            data = json.load(f)
        if uids is not None:
            data = {uid: data[uid] for uid in uids if uid in data}
        out.update(data)
        if uids is not None and len(out) == len(uids):
            break
    return out


def _load_object_paths() -> Dict[str, str]:
    """Load the object paths from the dataset.

    The object paths specify the location of where the object is located
    in the Hugging Face repo.

    Returns:
        A dictionary mapping the uid to the object path.
    """
    object_paths_file = "object-paths.json.gz"
    local_path = os.path.join(_VERSIONED_PATH, object_paths_file)
    if not os.path.exists(local_path):
        hf_url = f"https://huggingface.co/datasets/allenai/objaverse/resolve/main/{object_paths_file}"
        # wget the file and put it in local_path
        os.makedirs(os.path.dirname(local_path), exist_ok=True)
        urllib.request.urlretrieve(hf_url, local_path)
    with gzip.open(local_path, "rb") as f:
        object_paths = json.load(f)
    return object_paths


def load_uids() -> List[str]:
    """Load the uids from the dataset.

    Returns:
        A list of uids.
    """
    return list(_load_object_paths().keys())


def _download_object(
    uid: str,
    object_path: str,
    total_downloads: float,
    start_file_count: int,
) -> Tuple[str, str]:
    """Download the object for the given uid.

    Args:
        uid: The uid of the object to load.
        object_path: The path to the object in the Hugging Face repo.

    Returns:
        The local path of where the object was downloaded.
    """
    # print(f"downloading {uid}")
    local_path = os.path.join(_VERSIONED_PATH, object_path)
    tmp_local_path = os.path.join(_VERSIONED_PATH, object_path + ".tmp")
    hf_url = (
        f"https://huggingface.co/datasets/allenai/objaverse/resolve/main/{object_path}"
    )
    # wget the file and put it in local_path
    os.makedirs(os.path.dirname(tmp_local_path), exist_ok=True)
    urllib.request.urlretrieve(hf_url, tmp_local_path)

    os.rename(tmp_local_path, local_path)

    files = glob.glob(os.path.join(_VERSIONED_PATH, "glbs", "*", "*.glb"))
    print(
        "Downloaded",
        len(files) - start_file_count,
        "/",
        total_downloads,
        "objects",
    )

    return uid, local_path


def load_objects(uids: List[str], download_processes: int = 1) -> Dict[str, str]:
    """Return the path to the object files for the given uids.

    If the object is not already downloaded, it will be downloaded.

    Args:
        uids: A list of uids.
        download_processes: The number of processes to use to download the objects.

    Returns:
        A dictionary mapping the object uid to the local path of where the object
        downloaded.
    """
    object_paths = _load_object_paths()
    out = {}
    if download_processes == 1:
        uids_to_download = []
        for uid in uids:
            if uid.endswith(".glb"):
                uid = uid[:-4]
            if uid not in object_paths:
                warnings.warn(f"Could not find object with uid {uid}. Skipping it.")
                continue
            object_path = object_paths[uid]
            local_path = os.path.join(_VERSIONED_PATH, object_path)
            if os.path.exists(local_path):
                out[uid] = local_path
                continue
            uids_to_download.append((uid, object_path))
        if len(uids_to_download) == 0:
            return out
        start_file_count = len(
            glob.glob(os.path.join(_VERSIONED_PATH, "glbs", "*", "*.glb"))
        )
        for uid, object_path in uids_to_download:
            uid, local_path = _download_object(
                uid, object_path, len(uids_to_download), start_file_count
            )
            out[uid] = local_path
    else:
        args = []
        for uid in uids:
            if uid.endswith(".glb"):
                uid = uid[:-4]
            if uid not in object_paths:
                warnings.warn(f"Could not find object with uid {uid}. Skipping it.")
                continue
            object_path = object_paths[uid]
            local_path = os.path.join(_VERSIONED_PATH, object_path)
            if not os.path.exists(local_path):
                args.append((uid, object_paths[uid]))
            else:
                out[uid] = local_path
        if len(args) == 0:
            return out
        print(
            f"starting download of {len(args)} objects with {download_processes} processes"
        )
        start_file_count = len(
            glob.glob(os.path.join(_VERSIONED_PATH, "glbs", "*", "*.glb"))
        )
        args_list = [(*arg, len(args), start_file_count) for arg in args]
        with multiprocessing.Pool(download_processes) as pool:
            r = pool.starmap(_download_object, args_list)
            for uid, local_path in r:
                out[uid] = local_path
    return out


def load_lvis_annotations() -> Dict[str, List[str]]:
    """Load the LVIS annotations.

    If the annotations are not already downloaded, they will be downloaded.

    Returns:
        A dictionary mapping the LVIS category to the list of uids in that category.
    """
    hf_url = "https://huggingface.co/datasets/allenai/objaverse/resolve/main/lvis-annotations.json.gz"
    local_path = os.path.join(_VERSIONED_PATH, "lvis-annotations.json.gz")
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    if not os.path.exists(local_path):
        urllib.request.urlretrieve(hf_url, local_path)
    with gzip.open(local_path, "rb") as f:
        lvis_annotations = json.load(f)
    return lvis_annotations

objaverse/__init__.py
allenai/objaverse-xl
"""Utility functions for the objaverse_xl package."""

import hashlib
import os
import uuid


def get_uid_from_str(string: str) -> str:
    """Generates a UUID from a string.

    Args:
        string (str): String to generate a UUID from.

    Returns:
        str: UUID generated from the string.
    """
    namespace = uuid.NAMESPACE_DNS
    return str(uuid.uuid5(namespace, string))


def get_file_hash(file_path: str) -> str:
    """Get the sha256 hash of a file.

    Args:
        file_path (str): Path to the file.

    Returns:
        str: sha256 hash of the file.
    """
    # Check if the path is a symbolic link
    if os.path.islink(file_path):
        # Resolve the symbolic link
        resolved_path = os.readlink(file_path)
        # Check if the resolved path exists
        if not os.path.exists(resolved_path):
            raise FileNotFoundError(
                f"The symbolic link points to a file that doesn't exist: {resolved_path}"
            )
    sha256 = hashlib.sha256()
    # Read the file from the path
    with open(file_path, "rb") as f:
        # Loop till the end of the file
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256.update(byte_block)
    return sha256.hexdigest()

objaverse/utils.py
allenai/objaverse-xl
"""A package for downloading and processing Objaverse-XL."""

from typing import Callable, Dict, Optional

import pandas as pd

from objaverse.xl.github import GitHubDownloader
from objaverse.xl.sketchfab import SketchfabDownloader
from objaverse.xl.smithsonian import SmithsonianDownloader
from objaverse.xl.thingiverse import ThingiverseDownloader

downloaders = {
    "github": GitHubDownloader(),
    "thingiverse": ThingiverseDownloader(),
    "smithsonian": SmithsonianDownloader(),
    "sketchfab": SketchfabDownloader(),
}


def get_annotations(
    download_dir: str = "~/.objaverse", refresh: bool = False
) -> pd.DataFrame:
    """Loads the 3D object metadata as a Pandas DataFrame.

    Args:
        download_dir (str, optional): Directory to download the parquet metadata
            file. Supports all file systems supported by fsspec. Defaults to
            "~/.objaverse".
        refresh (bool, optional): Whether to refresh the annotations by downloading
            them from the remote source. Defaults to False.

    Returns:
        pd.DataFrame: Metadata of the 3D objects as a Pandas DataFrame with columns
            for the object "fileIdentifier", "license", "source", "fileType",
            "sha256", and "metadata".
    """
    annotations = [
        downloader.get_annotations(download_dir=download_dir, refresh=refresh)
        for downloader in downloaders.values()
    ]
    return pd.concat(annotations, ignore_index=True)


def download_objects(
    objects: pd.DataFrame,
    download_dir: Optional[str] = "~/.objaverse",
    processes: Optional[int] = None,
    handle_found_object: Optional[Callable] = None,
    handle_modified_object: Optional[Callable] = None,
    handle_missing_object: Optional[Callable] = None,
    **kwargs,
) -> Dict[str, str]:
    """Downloads all objects from the source.

    Args:
        objects (pd.DataFrame): Objects to download. Must have columns for
            the object "fileIdentifier", "source", and "sha256". Use the
            `get_annotations` function to get all objects as a DataFrame.
        download_dir (Optional[str], optional): Directory to download the objects
            to. Supports all file systems supported by fsspec. If None, the objects
            will be deleted after they are downloaded. Defaults to "~/.objaverse".
        processes (Optional[int], optional): Number of processes to use for
            downloading.  If None, will use the number of CPUs on the machine.
            Defaults to None.
        save_repo_format (Optional[Literal["zip", "tar", "tar.gz", "files"]],
            optional): Format to save the repository. If None, the repository will
            not be saved. If "files" is specified, each file will be saved
            individually. Otherwise, the repository can be saved as a "zip", "tar",
            or "tar.gz" file. Defaults to None.
        handle_found_object (Optional[Callable], optional): Called when an object is
            successfully found and downloaded. Here, the object has the same sha256
            as the one that was downloaded with Objaverse-XL. If None, the object
            will be downloaded, but nothing will be done with it. Args for the
            function include:
            - local_path (str): Local path to the downloaded 3D object.
            - file_identifier (str): File identifier of the 3D object.
            - sha256 (str): SHA256 of the contents of the 3D object.
            - metadata (Dict[Hashable, Any]): Metadata about the 3D object,
                including the GitHub organization and repo names.
            Return is not used. Defaults to None.
        handle_modified_object (Optional[Callable], optional): Called when a
            modified object is found and downloaded. Here, the object is
            successfully downloaded, but it has a different sha256 than the one that
            was downloaded with Objaverse-XL. This is not expected to happen very
            often, because the same commit hash is used for each repo. If None, the
            object will be downloaded, but nothing will be done with it. Args for
            the function include:
            - local_path (str): Local path to the downloaded 3D object.
            - file_identifier (str): File identifier of the 3D object.
            - new_sha256 (str): SHA256 of the contents of the newly downloaded 3D
                object.
            - old_sha256 (str): Expected SHA256 of the contents of the 3D object as
                it was when it was downloaded with Objaverse-XL.
            - metadata (Dict[Hashable, Any]): Metadata about the 3D object, which is
                particular to the souce.
            Return is not used. Defaults to None.
        handle_missing_object (Optional[Callable], optional): Called when an object
            that is in Objaverse-XL is not found. Here, it is likely that the
            repository was deleted or renamed. If None, nothing will be done with
            the missing object.
            Args for the function include:
            - file_identifier (str): File identifier of the 3D object.
            - sha256 (str): SHA256 of the contents of the original 3D object.
            - metadata (Dict[Hashable, Any]): Metadata about the 3D object, which is
                particular to the source.
            Return is not used. Defaults to None.

    Returns:
        Dict[str, str]: Mapping of file identifiers to local paths of the downloaded
            3D objects.
    """
    sources = set(objects["source"].unique().tolist())
    all_sources = {"github", "thingiverse", "smithsonian", "sketchfab"}

    if not sources.issubset(all_sources):
        raise ValueError(
            f"Invalid sources: {sources}. Must be a subset of {all_sources}."
        )

    downloaded_objects = {}
    for source in sources:
        source_downloads = downloaders[source].download_objects(
            objects[objects["source"] == source],
            download_dir,
            processes,
            handle_found_object,
            handle_modified_object,
            handle_missing_object,
            **kwargs,
        )
        downloaded_objects.update(source_downloads)

    return downloaded_objects

objaverse/xl/__init__.py
allenai/objaverse-xl
"""Abstract class for Objaverse-XL sources."""

from abc import ABC, abstractmethod
from typing import Callable, Dict, Optional

import pandas as pd


class ObjaverseSource(ABC):
    """Abstract class for Objaverse-XL sources."""

    @classmethod
    @abstractmethod
    def get_annotations(
        cls, download_dir: str = "~/.objaverse", refresh: bool = False
    ) -> pd.DataFrame:
        """Loads the 3D object metadata as a Pandas DataFrame.

        Args:
            download_dir (str, optional): Directory to download the parquet metadata
                file. Supports all file systems supported by fsspec. Defaults to
                "~/.objaverse".
            refresh (bool, optional): Whether to refresh the annotations by downloading
                them from the remote source. Defaults to False.

        Returns:
            pd.DataFrame: Metadata of the 3D objects as a Pandas DataFrame with columns
                for the object "fileIdentifier", "license", "source", "fileType",
                "sha256", and "metadata".
        """

    @classmethod
    @abstractmethod
    def download_objects(
        cls,
        objects: pd.DataFrame,
        download_dir: Optional[str] = "~/.objaverse",
        processes: Optional[int] = None,
        handle_found_object: Optional[Callable] = None,
        handle_modified_object: Optional[Callable] = None,
        handle_missing_object: Optional[Callable] = None,
        **kwargs
    ) -> Dict[str, str]:
        """Downloads all objects from the source.

        Args:
            objects (pd.DataFrame): Objects to download. Must have columns for
                the object "fileIdentifier" and "sha256". Use the `get_annotations`
                function to get the metadata.
            processes (Optional[int], optional): Number of processes to use for
                downloading.  If None, will use the number of CPUs on the machine.
                Defaults to None.
            download_dir (Optional[str], optional): Directory to download the objects
                to. Supports all file systems supported by fsspec. If None, the objects
                will be disregarded after they are downloaded and processed. Defaults to
                "~/.objaverse".
            save_repo_format (Optional[Literal["zip", "tar", "tar.gz", "files"]],
                optional): Format to save the repository. If None, the repository will
                not be saved. If "files" is specified, each file will be saved
                individually. Otherwise, the repository can be saved as a "zip", "tar",
                or "tar.gz" file. Defaults to None.
            handle_found_object (Optional[Callable], optional): Called when an object is
                successfully found and downloaded. Here, the object has the same sha256
                as the one that was downloaded with Objaverse-XL. If None, the object
                will be downloaded, but nothing will be done with it. Args for the
                function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the 3D object.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object,
                    including the GitHub organization and repo names.
                Return is not used. Defaults to None.
            handle_modified_object (Optional[Callable], optional): Called when a
                modified object is found and downloaded. Here, the object is
                successfully downloaded, but it has a different sha256 than the one that
                was downloaded with Objaverse-XL. This is not expected to happen very
                often, because the same commit hash is used for each repo. If None, the
                object will be downloaded, but nothing will be done with it. Args for
                the function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - new_sha256 (str): SHA256 of the contents of the newly downloaded 3D
                    object.
                - old_sha256 (str): Expected SHA256 of the contents of the 3D object as
                    it was when it was downloaded with Objaverse-XL.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object, which is
                    particular to the souce.
                Return is not used. Defaults to None.
            handle_missing_object (Optional[Callable], optional): Called when an object
                that is in Objaverse-XL is not found. Here, it is likely that the
                repository was deleted or renamed. If None, nothing will be done with
                the missing object.
                Args for the function include:
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the original 3D object.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object, which is
                    particular to the source.
                Return is not used. Defaults to None.

        Returns:
            Dict[str, str]: Mapping of file identifiers to local paths of the downloaded
                3D objects.
        """
        pass

objaverse/xl/abstract.py
allenai/objaverse-xl
"""Script to download objects from GitHub."""

import json
import multiprocessing
import os
import shutil
import subprocess
import tarfile
import tempfile
from multiprocessing import Pool
from typing import Callable, Dict, List, Literal, Optional

import fsspec
import pandas as pd
import requests
from loguru import logger
from tqdm import tqdm

from objaverse.utils import get_file_hash
from objaverse.xl.abstract import ObjaverseSource

FILE_EXTENSIONS = [
    ".obj",
    ".glb",
    ".gltf",
    ".usdz",
    ".usd",
    ".fbx",
    ".stl",
    ".usda",
    ".dae",
    ".ply",
    ".abc",
    ".blend",
]


class GitHubDownloader(ObjaverseSource):
    """Script to download objects from GitHub."""

    @classmethod
    def get_annotations(
        cls, download_dir: str = "~/.objaverse", refresh: bool = False
    ) -> pd.DataFrame:
        """Loads the GitHub 3D object metadata as a Pandas DataFrame.

        Args:
            download_dir (str, optional): Directory to download the parquet metadata
                file. Supports all file systems supported by fsspec. Defaults to
                "~/.objaverse".
            refresh (bool, optional): Whether to refresh the annotations by downloading
                them from the remote source. Defaults to False.

        Returns:
            pd.DataFrame: GitHub 3D object metadata as a Pandas DataFrame with columns
                for the object "fileIdentifier", "license", "source", "fileType",
                "sha256", and "metadata".
        """
        filename = os.path.join(download_dir, "github", "github.parquet")
        fs, path = fsspec.core.url_to_fs(filename)
        fs.makedirs(os.path.dirname(path), exist_ok=True)

        # download the parquet file if it doesn't exist
        if refresh or not fs.exists(path):
            url = "https://huggingface.co/datasets/allenai/objaverse-xl/resolve/main/github/github.parquet"
            logger.info(f"Downloading {url} to {filename}")
            response = requests.get(url)
            response.raise_for_status()
            with fs.open(path, "wb") as file:
                file.write(response.content)

        # load the parquet file with fsspec
        with fs.open(path) as f:
            df = pd.read_parquet(f)

        return df

    @classmethod
    def _get_repo_id_with_hash(cls, item: pd.Series) -> str:
        org, repo = item["fileIdentifier"].split("/")[3:5]
        commit_hash = item["fileIdentifier"].split("/")[6]
        return f"{org}/{repo}/{commit_hash}"

    @classmethod
    def _git_shallow_clone(cls, repo_url: str, target_directory: str) -> bool:
        """Helper function to shallow clone a repo with git.

        Args:
            repo_url (str): URL of the repo to clone.
            target_directory (str): Directory to clone the repo to.

        Returns:
            bool: True if the clone was successful, False otherwise.
        """
        return cls._run_command_with_check(
            ["git", "clone", "--depth", "1", repo_url, target_directory],
        )

    @classmethod
    def _run_command_with_check(
        cls, command: List[str], cwd: Optional[str] = None
    ) -> bool:
        """Helper function to run a command and check if it was successful.

        Args:
            command (List[str]): Command to run.
            cwd (Optional[str], optional): Current working directory to run the command
                in. Defaults to None.

        Returns:
            bool: True if the command was successful, False otherwise.
        """
        try:
            subprocess.run(
                command,
                cwd=cwd,
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            return True
        except subprocess.CalledProcessError as e:
            logger.error("Error:", e)
            logger.error(e.stdout)
            logger.error(e.stderr)
            return False

    @classmethod
    def _process_repo(
        cls,
        repo_id: str,
        fs: fsspec.AbstractFileSystem,
        base_dir: str,
        save_repo_format: Optional[Literal["zip", "tar", "tar.gz", "files"]],
        expected_objects: Dict[str, str],
        handle_found_object: Optional[Callable],
        handle_modified_object: Optional[Callable],
        handle_missing_object: Optional[Callable],
        handle_new_object: Optional[Callable],
        commit_hash: Optional[str],
    ) -> Dict[str, str]:
        """Process a single repo.

        Args:
            repo_id (str): GitHub repo ID in the format of organization/repo.
            fs (fsspec.AbstractFileSystem): File system to use for saving the repo.
            base_dir (str): Base directory to save the repo to.
            expected_objects (Dict[str, str]): Dictionary of objects that one expects to
                find in the repo. Keys are the "fileIdentifier" (i.e., the GitHub URL in
                this case) and values are the "sha256" of the objects.
            {and the rest of the args are the same as download_objects}

        Returns:
            Dict[str, str]: A dictionary that maps from the "fileIdentifier" to the path
                of the downloaded object.
        """
        # NOTE: assuming that the user has already checked that the repo doesn't exist,
        org, repo = repo_id.split("/")

        out = {}
        with tempfile.TemporaryDirectory() as temp_dir:
            # clone the repo to a temp directory
            target_directory = os.path.join(temp_dir, repo)
            successful_clone = cls._git_shallow_clone(
                f"https://github.com/{org}/{repo}.git", target_directory
            )
            if not successful_clone:
                logger.error(f"Could not clone {repo_id}")
                if handle_missing_object is not None:
                    for github_url, sha256 in expected_objects.items():
                        handle_missing_object(
                            file_identifier=github_url,
                            sha256=sha256,
                            metadata=dict(github_organization=org, github_repo=repo),
                        )
                return {}

            # use the commit hash if specified
            repo_commit_hash = cls._get_commit_hash_from_local_git_dir(target_directory)
            if commit_hash is not None:
                keep_going = True
                if repo_commit_hash != commit_hash:
                    # run git reset --hard && git checkout 37f4d8d287e201ce52c048bf74d46d6a09d26b2c
                    if not cls._run_command_with_check(
                        ["git", "fetch", "origin", commit_hash], target_directory
                    ):
                        logger.error(
                            f"Error in git fetch! Sticking with {repo_commit_hash=} instead of {commit_hash=}"
                        )
                        keep_going = False

                    if keep_going and not cls._run_command_with_check(
                        ["git", "reset", "--hard"], target_directory
                    ):
                        logger.error(
                            f"Error in git reset! Sticking with {repo_commit_hash=} instead of {commit_hash=}"
                        )
                        keep_going = False

                    if keep_going:
                        if cls._run_command_with_check(
                            ["git", "checkout", commit_hash], target_directory
                        ):
                            repo_commit_hash = commit_hash
                        else:
                            logger.error(
                                f"Error in git checkout! Sticking with {repo_commit_hash=} instead of {commit_hash=}"
                            )

            # pull the lfs files
            cls._pull_lfs_files(target_directory)

            # get all the files in the repo
            files = cls._list_files(target_directory)
            files_with_3d_extension = [
                file
                for file in files
                if any(file.lower().endswith(ext) for ext in FILE_EXTENSIONS)
            ]

            # get the sha256 for each file
            file_hashes = []
            for file in tqdm(files_with_3d_extension, desc="Handling 3D object files"):
                file_hash = get_file_hash(file)
                # remove the temp_dir from the file path
                github_url = file.replace(
                    target_directory,
                    f"https://github.com/{org}/{repo}/blob/{repo_commit_hash}",
                )
                file_hashes.append(dict(sha256=file_hash, fileIdentifier=github_url))

                # handle the object under different conditions
                if github_url in expected_objects:
                    out[github_url] = file[len(target_directory) + 1 :]
                    if expected_objects[github_url] == file_hash:
                        if handle_found_object is not None:
                            handle_found_object(
                                local_path=file,
                                file_identifier=github_url,
                                sha256=file_hash,
                                metadata=dict(
                                    github_organization=org, github_repo=repo
                                ),
                            )
                    else:
                        if handle_modified_object is not None:
                            handle_modified_object(
                                local_path=file,
                                file_identifier=github_url,
                                new_sha256=file_hash,
                                old_sha256=expected_objects[github_url],
                                metadata=dict(
                                    github_organization=org, github_repo=repo
                                ),
                            )
                elif handle_new_object is not None:
                    handle_new_object(
                        local_path=file,
                        file_identifier=github_url,
                        sha256=file_hash,
                        metadata=dict(github_organization=org, github_repo=repo),
                    )

            # save the file hashes to a json file
            with open(
                os.path.join(target_directory, ".objaverse-file-hashes.json"),
                "w",
                encoding="utf-8",
            ) as f:
                json.dump(file_hashes, f, indent=2)

            # remove the .git directory
            shutil.rmtree(os.path.join(target_directory, ".git"))

            if save_repo_format is None:
                # remove the paths, since it's not downloaded
                out = {}
            else:
                logger.debug(f"Saving {org}/{repo} as {save_repo_format}")
                # save the repo to a zip file
                if save_repo_format == "zip":
                    shutil.make_archive(target_directory, "zip", target_directory)
                elif save_repo_format == "tar":
                    with tarfile.open(
                        os.path.join(temp_dir, f"{repo}.tar"), "w"
                    ) as tar:
                        tar.add(target_directory, arcname=repo)
                elif save_repo_format == "tar.gz":
                    with tarfile.open(
                        os.path.join(temp_dir, f"{repo}.tar.gz"), "w:gz"
                    ) as tar:
                        tar.add(target_directory, arcname=repo)
                elif save_repo_format == "files":
                    pass
                else:
                    raise ValueError(
                        f"save_repo_format must be one of zip, tar, tar.gz, files. Got {save_repo_format}"
                    )

                dirname = os.path.join(base_dir, "repos", org)
                fs.makedirs(dirname, exist_ok=True)
                if save_repo_format != "files":
                    # move the repo to the correct location (with put)
                    fs.put(
                        os.path.join(temp_dir, f"{repo}.{save_repo_format}"),
                        os.path.join(dirname, f"{repo}.{save_repo_format}"),
                    )

                    for file_identifier in out.copy():
                        out[file_identifier] = os.path.join(
                            dirname, f"{repo}.{save_repo_format}", out[file_identifier]
                        )
                else:
                    # move the repo to the correct location (with put)
                    fs.put(target_directory, dirname, recursive=True)

                    for file_identifier in out.copy():
                        out[file_identifier] = os.path.join(
                            dirname, repo, out[file_identifier]
                        )

        # get each object that was missing from the expected objects
        if handle_missing_object is not None:
            obtained_urls = {x["fileIdentifier"] for x in file_hashes}
            for github_url, sha256 in expected_objects.items():
                if github_url not in obtained_urls:
                    handle_missing_object(
                        file_identifier=github_url,
                        sha256=sha256,
                        metadata=dict(github_organization=org, github_repo=repo),
                    )

        return out

    @classmethod
    def _list_files(cls, root_dir: str) -> List[str]:
        return [
            os.path.join(root, f)
            for root, dirs, files in os.walk(root_dir)
            for f in files
        ]

    @classmethod
    def _pull_lfs_files(cls, repo_dir: str) -> None:
        if cls._has_lfs_files(repo_dir):
            subprocess.run(["git", "lfs", "pull"], cwd=repo_dir, check=True)

    @classmethod
    def _has_lfs_files(cls, repo_dir: str) -> bool:
        gitattributes_path = os.path.join(repo_dir, ".gitattributes")
        if not os.path.exists(gitattributes_path):
            return False
        with open(gitattributes_path, "r", encoding="utf-8") as f:
            for line in f:
                if "filter=lfs" in line:
                    return True
        return False

    @classmethod
    def _get_commit_hash_from_local_git_dir(cls, local_git_dir: str) -> str:
        """Get the commit hash of the local git directory."""
        # get the git hash of the repo
        result = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=local_git_dir,
            capture_output=True,
            check=True,
        )
        commit_hash = result.stdout.strip().decode("utf-8")
        return commit_hash

    @classmethod
    def _parallel_process_repo(cls, args) -> Dict[str, str]:
        """Helper function to process a repo in parallel.

        Note: This function is used to parallelize the processing of repos. It is not
        intended to be called directly.

        Args:
            args (Tuple): Tuple of arguments to pass to _process_repo.

        Returns:
            Dict[str, str]: A dictionary that maps from the "fileIdentifier" to the path
                of the downloaded object.
        """

        (
            repo_id_hash,
            fs,
            base_dir,
            save_repo_format,
            expected_objects,
            handle_found_object,
            handle_modified_object,
            handle_missing_object,
            handle_new_object,
        ) = args
        repo_id = "/".join(repo_id_hash.split("/")[:2])
        commit_hash = repo_id_hash.split("/")[2]
        return cls._process_repo(
            repo_id=repo_id,
            fs=fs,
            base_dir=base_dir,
            save_repo_format=save_repo_format,
            expected_objects=expected_objects,
            handle_found_object=handle_found_object,
            handle_modified_object=handle_modified_object,
            handle_missing_object=handle_missing_object,
            handle_new_object=handle_new_object,
            commit_hash=commit_hash,
        )

    @classmethod
    def _process_group(cls, group):
        key, group_df = group
        return key, group_df.set_index("fileIdentifier")["sha256"].to_dict()

    @classmethod
    def download_objects(
        cls,
        objects: pd.DataFrame,
        download_dir: Optional[str] = "~/.objaverse",
        processes: Optional[int] = None,
        handle_found_object: Optional[Callable] = None,
        handle_modified_object: Optional[Callable] = None,
        handle_missing_object: Optional[Callable] = None,
        **kwargs,
    ) -> Dict[str, str]:
        """Download the specified GitHub objects.

        Args:
            objects (pd.DataFrame): GitHub objects to download. Must have columns for
                the object "fileIdentifier" and "sha256". Use the `get_annotations`
                function to get the metadata.
            download_dir (Optional[str], optional): Directory to download the GitHub
                objects to. Supports all file systems supported by fsspec. If None, the
                repository will not be saved (note that save_repo_format must also be
                None in this case, otherwise a ValueError is raised). Defaults to
                "~/.objaverse".
            processes (Optional[int], optional): Number of processes to use for
                downloading.  If None, will use the number of CPUs on the machine.
                Defaults to None.
            handle_found_object (Optional[Callable], optional): Called when an object is
                successfully found and downloaded. Here, the object has the same sha256
                as the one that was downloaded with Objaverse-XL. If None, the object
                will be downloaded, but nothing will be done with it. Args for the
                function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): GitHub URL of the 3D object.
                - sha256 (str): SHA256 of the contents of the 3D object.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used. Defaults to None.
            handle_modified_object (Optional[Callable], optional): Called when a
                modified object is found and downloaded. Here, the object is
                successfully downloaded, but it has a different sha256 than the one that
                was downloaded with Objaverse-XL. This is not expected to happen very
                often, because the same commit hash is used for each repo. If None, the
                object will be downloaded, but nothing will be done with it. Args for
                the function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): GitHub URL of the 3D object.
                - new_sha256 (str): SHA256 of the contents of the newly downloaded 3D
                    object.
                - old_sha256 (str): Expected SHA256 of the contents of the 3D object as
                    it was when it was downloaded with Objaverse-XL.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used. Defaults to None.
            handle_missing_object (Optional[Callable], optional): Called when an object
                that is in Objaverse-XL is not found. Here, it is likely that the
                repository was deleted or renamed. If None, nothing will be done with
                the missing object. Args for the function include:
                - file_identifier (str): GitHub URL of the 3D object.
                - sha256 (str): SHA256 of the contents of the original 3D object.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used. Defaults to None.
            save_repo_format (Optional[Literal["zip", "tar", "tar.gz", "files"]],
                optional): Format to save the repository. If None, the repository will
                not be saved. If "files" is specified, each file will be saved
                individually. Otherwise, the repository can be saved as a "zip", "tar",
                or "tar.gz" file. Defaults to None.
            handle_new_object (Optional[Callable], optional): Called when a new object
                is found. Here, the object is not used in Objaverse-XL, but is still
                downloaded with the repository. The object may have not been used
                because it does not successfully import into Blender. If None, the
                object will be downloaded, but nothing will be done with it. Args for
                the function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): GitHub URL of the 3D object.
                - sha256 (str): SHA256 of the contents of the 3D object.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used. Defaults to None.

        Raises:
            ValueError: If download_dir is None and save_repo_format is not None.
                Otherwise, we don't know where to save the repo!

        Returns:
            Dict[str, str]: A dictionary that maps from the "fileIdentifier" to the path
                of the downloaded object.
        """
        save_repo_format = kwargs.get("save_repo_format", None)
        handle_new_object = kwargs.get("handle_new_object", None)

        if processes is None:
            processes = multiprocessing.cpu_count()
        if download_dir is None:
            if save_repo_format is not None:
                raise ValueError(
                    f"If {save_repo_format=} is not None, {download_dir=} must be specified."
                )
            # path doesn't matter if we're not saving the repo
            download_dir = "~/.objaverse"

        base_download_dir = os.path.join(download_dir, "github")
        fs, path = fsspec.core.url_to_fs(base_download_dir)
        fs.makedirs(path, exist_ok=True)

        # Getting immediate subdirectories of root_path
        if save_repo_format == "files":
            downloaded_repo_dirs = fs.glob(base_download_dir + "/repos/*/*/")
            downloaded_repo_ids = {
                "/".join(x.split("/")[-2:]) for x in downloaded_repo_dirs
            }
        else:
            downloaded_repo_dirs = fs.glob(
                base_download_dir + f"/repos/*/*.{save_repo_format}"
            )
            downloaded_repo_ids = set()
            for x in downloaded_repo_dirs:
                org, repo = x.split("/")[-2:]
                repo = repo[: -len(f".{save_repo_format}")]
                repo_id = f"{org}/{repo}"
                downloaded_repo_ids.add(repo_id)

        # make copy of objects
        objects = objects.copy()

        # get the unique repoIds
        objects["repoIdHash"] = objects.apply(cls._get_repo_id_with_hash, axis=1)
        repo_id_hashes = set(objects["repoIdHash"].unique().tolist())
        repo_ids = {
            "/".join(repo_id_hash.split("/")[:2]) for repo_id_hash in repo_id_hashes
        }
        assert len(repo_id_hashes) == len(repo_ids), (
            f"More than 1 commit hash per repoId!"
            f" {len(repo_id_hashes)=}, {len(repo_ids)=}"
        )

        logger.info(
            f"Provided {len(repo_ids)} repoIds with {len(objects)} objects to process."
        )

        # remove repoIds that have already been downloaded
        repo_ids_to_download = repo_ids - downloaded_repo_ids
        repo_id_hashes_to_download = [
            repo_id_hash
            for repo_id_hash in repo_id_hashes
            if "/".join(repo_id_hash.split("/")[:2]) in repo_ids_to_download
        ]

        logger.info(
            f"Found {len(repo_ids_to_download)} repoIds not yet downloaded. Downloading now..."
        )

        # get the objects to download
        groups = list(objects.groupby("repoIdHash"))
        with Pool(processes=processes) as pool:
            out_list = list(
                tqdm(
                    pool.imap_unordered(cls._process_group, groups),
                    total=len(groups),
                    desc="Grouping objects by repository",
                )
            )
        objects_per_repo_id_hash = dict(out_list)

        all_args = [
            (
                repo_id_hash,
                fs,
                path,
                save_repo_format,
                objects_per_repo_id_hash[repo_id_hash],
                handle_found_object,
                handle_modified_object,
                handle_missing_object,
                handle_new_object,
            )
            for repo_id_hash in repo_id_hashes_to_download
        ]

        with Pool(processes=processes) as pool:
            # use tqdm to show progress
            out = list(
                tqdm(
                    pool.imap_unordered(cls._parallel_process_repo, all_args),
                    total=len(all_args),
                    desc="Downloading repositories",
                )
            )

        out_dict = {}
        for x in out:
            out_dict.update(x)

        return out_dict

objaverse/xl/github.py
allenai/objaverse-xl
"""Script to download objects from Objaverse 1.0."""

import gzip
import json
import multiprocessing
import os
import tempfile
import urllib.request
from multiprocessing import Pool
from typing import Any, Callable, Dict, List, Optional, Tuple

import fsspec
import pandas as pd
import requests
from loguru import logger
from tqdm import tqdm

from objaverse.utils import get_file_hash
from objaverse.xl.abstract import ObjaverseSource


class SketchfabDownloader(ObjaverseSource):
    """A class for downloading and processing Objaverse 1.0."""

    @classmethod
    def get_annotations(
        cls, download_dir: str = "~/.objaverse", refresh: bool = False
    ) -> pd.DataFrame:
        """Load the annotations from the given directory.

        Args:
            download_dir (str, optional): The directory to load the annotations from.
                Supports all file systems supported by fsspec. Defaults to
                "~/.objaverse".
            refresh (bool, optional): Whether to refresh the annotations by downloading
                them from the remote source. Defaults to False.

        Returns:
            pd.DataFrame: The annotations, which includes the columns "thingId", "fileId",
                "filename", and "license".
        """
        remote_url = "https://huggingface.co/datasets/allenai/objaverse-xl/resolve/main/sketchfab/sketchfab.parquet"
        download_path = os.path.join(download_dir, "sketchfab", "sketchfab.parquet")
        fs, path = fsspec.core.url_to_fs(download_path)

        if refresh or not fs.exists(path):
            fs.makedirs(os.path.dirname(path), exist_ok=True)
            logger.info(f"Downloading {remote_url} to {download_path}")
            response = requests.get(remote_url)
            response.raise_for_status()
            with fs.open(path, "wb") as file:
                file.write(response.content)

        # read the file with pandas and fsspec
        with fs.open(download_path, "rb") as f:
            annotations_df = pd.read_parquet(f)

        annotations_df["metadata"] = "{}"

        return annotations_df

    @classmethod
    def get_full_annotations(
        cls,
        uids: Optional[List[str]] = None,
        download_dir: str = "~/.objaverse",
    ) -> Dict[str, Any]:
        """Get the full metadata of all objects in the dataset.

        Args:
            uids: A list of uids with which to load metadata. If None, it loads
                the metadata for all uids.
            download_dir: The base directory to download the annotations to. Supports all
                file systems supported by fsspec. Defaults to "~/.objaverse".

        Returns:
            A dictionary of the metadata for each object. The keys are the uids and the
            values are the metadata for that object.
        """
        # make the metadata dir if it doesn't exist
        metadata_path = os.path.join(download_dir, "hf-objaverse-v1", "metadata")
        fs, _ = fsspec.core.url_to_fs(metadata_path)
        fs.makedirs(metadata_path, exist_ok=True)

        # get the dir ids that need to be loaded if only downloading a subset of uids
        object_paths = cls._get_object_paths(download_dir=download_dir)
        dir_ids = (
            {object_paths[uid].split("/")[1] for uid in uids}
            if uids is not None
            else {f"{i // 1000:03d}-{i % 1000:03d}" for i in range(160)}
        )

        # get the existing metadata files
        existing_metadata_files = fs.glob(
            os.path.join(metadata_path, "*.json.gz"), refresh=True
        )
        existing_dir_ids = {
            file.split("/")[-1].split(".")[0]
            for file in existing_metadata_files
            if file.endswith(".json.gz")  # note partial files end with .json.gz.tmp
        }
        downloaded_dir_ids = existing_dir_ids.intersection(dir_ids)
        logger.info(
            f"Found {len(downloaded_dir_ids)} metadata files already downloaded"
        )

        # download the metadata from the missing dir_ids
        dir_ids_to_download = dir_ids - existing_dir_ids
        logger.info(f"Downloading {len(dir_ids_to_download)} metadata files")

        # download the metadata file if it doesn't exist
        if len(dir_ids_to_download) > 0:
            for i_id in tqdm(dir_ids_to_download, desc="Downloading metadata files"):
                # get the path to the json file
                path = os.path.join(metadata_path, f"{i_id}.json.gz")

                # get the url to the remote json file
                hf_url = f"https://huggingface.co/datasets/allenai/objaverse/resolve/main/metadata/{i_id}.json.gz"

                # download the file to a tmp path to avoid partial downloads on interruption
                tmp_path = f"{path}.tmp"
                with fs.open(tmp_path, "wb") as f:
                    with urllib.request.urlopen(hf_url) as response:
                        f.write(response.read())
                fs.rename(tmp_path, path)

        out = {}
        for i_id in tqdm(dir_ids, desc="Reading metadata files"):
            # get the path to the json file
            path = os.path.join(metadata_path, f"{i_id}.json.gz")

            # read the json file of the metadata chunk
            with fs.open(path, "rb") as f:
                with gzip.GzipFile(fileobj=f) as gfile:
                    content = gfile.read()
                    data = json.loads(content)

            # filter the data to only include the uids we want
            if uids is not None:
                data = {uid: data[uid] for uid in uids if uid in data}

            # add the data to the out dict
            out.update(data)

        return out

    @classmethod
    def _get_object_paths(cls, download_dir: str) -> Dict[str, str]:
        """Load the object paths from the dataset.

        The object paths specify the location of where the object is located in the
        Hugging Face repo.

        Returns:
            A dictionary mapping the uid to the object path.
        """
        object_paths_file = "object-paths.json.gz"
        local_path = os.path.join(download_dir, "hf-objaverse-v1", object_paths_file)

        # download the object_paths file if it doesn't exist
        fs, path = fsspec.core.url_to_fs(local_path)
        if not fs.exists(path):
            hf_url = f"https://huggingface.co/datasets/allenai/objaverse/resolve/main/{object_paths_file}"
            fs.makedirs(os.path.dirname(path), exist_ok=True)

            # download the file to a tmp path to avoid partial downloads on interruption
            tmp_path = f"{path}.tmp"
            with fs.open(tmp_path, "wb") as f:
                with urllib.request.urlopen(hf_url) as response:
                    f.write(response.read())
            fs.rename(tmp_path, path)

        # read the object_paths
        with fs.open(path, "rb") as f:
            with gzip.GzipFile(fileobj=f) as gfile:
                content = gfile.read()
                object_paths = json.loads(content)

        return object_paths

    @classmethod
    def get_uids(cls, download_dir: str = "~/.objaverse") -> List[str]:
        """Load the uids from the dataset.

        Returns:
            A list of all the UIDs from the dataset.
        """
        return list(cls._get_object_paths(download_dir=download_dir).keys())

    @classmethod
    def _download_object(
        cls,
        file_identifier: str,
        hf_object_path: str,
        download_dir: Optional[str],
        expected_sha256: str,
        handle_found_object: Optional[Callable] = None,
        handle_modified_object: Optional[Callable] = None,
    ) -> Tuple[str, Optional[str]]:
        """Download the object for the given uid.

        Args:
            file_identifier: The file identifier of the object.
            hf_object_path: The path to the object in the Hugging Face repo. Here,
                hf_object_path is the part that comes after "main" in the Hugging Face
                repo url:
                https://huggingface.co/datasets/allenai/objaverse/resolve/main/{hf_object_path}
            download_dir: The base directory to download the object to. Supports all
                file systems supported by fsspec. Defaults to "~/.objaverse".
            expected_sha256 (str): The expected SHA256 of the contents of the downloade
                object.
            handle_found_object (Optional[Callable]): Called when an object is
                successfully found and downloaded. Here, the object has the same sha256
                as the one that was downloaded with Objaverse-XL. If None, the object
                will be downloaded, but nothing will be done with it. Args for the
                function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): GitHub URL of the 3D object.
                - sha256 (str): SHA256 of the contents of the 3D object.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used.
            handle_modified_object (Optional[Callable]): Called when a modified object
                is found and downloaded. Here, the object is successfully downloaded,
                but it has a different sha256 than the one that was downloaded with
                Objaverse-XL. This is not expected to happen very often, because the
                same commit hash is used for each repo. If None, the object will be
                downloaded, but nothing will be done with it. Args for the function
                include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): GitHub URL of the 3D object.
                - new_sha256 (str): SHA256 of the contents of the newly downloaded 3D
                    object.
                - old_sha256 (str): Expected SHA256 of the contents of the 3D object as
                    it was when it was downloaded with Objaverse-XL.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used.
            handle_missing_object (Optional[Callable]): Called when an object that is in
                Objaverse-XL is not found. Here, it is likely that the repository was
                deleted or renamed. If None, nothing will be done with the missing
                object. Args for the function include:
                - file_identifier (str): GitHub URL of the 3D object.
                - sha256 (str): SHA256 of the contents of the original 3D object.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used.


        Returns:
            A tuple of the uid and the path to where the downloaded object. If
            download_dir is None, the path will be None.
        """
        hf_url = f"https://huggingface.co/datasets/allenai/objaverse/resolve/main/{hf_object_path}"

        with tempfile.TemporaryDirectory() as temp_dir:
            # download the file locally
            temp_path = os.path.join(temp_dir, hf_object_path)
            os.makedirs(os.path.dirname(temp_path), exist_ok=True)
            temp_path_tmp = f"{temp_path}.tmp"
            with open(temp_path_tmp, "wb") as file:
                with urllib.request.urlopen(hf_url) as response:
                    file.write(response.read())
            os.rename(temp_path_tmp, temp_path)

            # get the sha256 of the downloaded file
            sha256 = get_file_hash(temp_path)

            if sha256 == expected_sha256:
                if handle_found_object is not None:
                    handle_found_object(
                        local_path=temp_path,
                        file_identifier=file_identifier,
                        sha256=sha256,
                        metadata={},
                    )
            else:
                if handle_modified_object is not None:
                    handle_modified_object(
                        local_path=temp_path,
                        file_identifier=file_identifier,
                        new_sha256=sha256,
                        old_sha256=expected_sha256,
                        metadata={},
                    )

            if download_dir is not None:
                filename = os.path.join(download_dir, "hf-objaverse-v1", hf_object_path)
                fs, path = fsspec.core.url_to_fs(filename)
                fs.makedirs(os.path.dirname(path), exist_ok=True)
                fs.put(temp_path, path)
            else:
                path = None

        return file_identifier, path

    @classmethod
    def _parallel_download_object(cls, args):
        # workaround since starmap doesn't work well with tqdm
        return cls._download_object(*args)

    @classmethod
    def _get_uid(cls, item: pd.Series) -> str:
        file_identifier = item["fileIdentifier"]
        return file_identifier.split("/")[-1]

    @classmethod
    def uid_to_file_identifier(cls, uid: str) -> str:
        """Convert the uid to the file identifier.

        Args:
            uid (str): The uid of the object.

        Returns:
            The file identifier of the object.
        """
        return f"https://sketchfab.com/3d-models/{uid}"

    @classmethod
    def file_identifier_to_uid(cls, file_identifier: str) -> str:
        """Convert the file identifier to the uid.

        Args:
            file_identifier (str): The file identifier of the object.

        Returns:
            The uid of the object.
        """
        return file_identifier.split("/")[-1]

    @classmethod
    def download_objects(
        cls,
        objects: pd.DataFrame,
        download_dir: Optional[str] = "~/.objaverse",
        processes: Optional[int] = None,
        handle_found_object: Optional[Callable] = None,
        handle_modified_object: Optional[Callable] = None,
        handle_missing_object: Optional[Callable] = None,
        **kwargs,
    ) -> Dict[str, str]:
        """Return the path to the object files for the given uids.

        If the object is not already downloaded, it will be downloaded.

        Args:
            objects (pd.DataFrame): Objects to download. Must have columns for
                the object "fileIdentifier" and "sha256". Use the `get_annotations`
                function to get the metadata.
            download_dir (Optional[str], optional): The base directory to download the
                object to. Supports all file systems supported by fsspec. If None, the
                objects will be removed after downloading. Defaults to "~/.objaverse".
            processes (Optional[int], optional): The number of processes to use to
                download the objects. If None, the number of processes will be set to
                the number of CPUs on the machine (multiprocessing.cpu_count()).
                Defaults to None.
            handle_found_object (Optional[Callable], optional): Called when an object is
                successfully found and downloaded. Here, the object has the same sha256
                as the one that was downloaded with Objaverse-XL. If None, the object
                will be downloaded, but nothing will be done with it. Args for the
                function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the 3D object.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object,
                    including the GitHub organization and repo names.
                Return is not used. Defaults to None.
            handle_modified_object (Optional[Callable], optional): Called when a
                modified object is found and downloaded. Here, the object is
                successfully downloaded, but it has a different sha256 than the one that
                was downloaded with Objaverse-XL. This is not expected to happen very
                often, because the same commit hash is used for each repo. If None, the
                object will be downloaded, but nothing will be done with it. Args for
                the function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - new_sha256 (str): SHA256 of the contents of the newly downloaded 3D
                    object.
                - old_sha256 (str): Expected SHA256 of the contents of the 3D object as
                    it was when it was downloaded with Objaverse-XL.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object, which is
                    particular to the souce.
                Return is not used. Defaults to None.
            handle_missing_object (Optional[Callable], optional): Called when an object
                that is in Objaverse-XL is not found. Here, it is likely that the
                repository was deleted or renamed. If None, nothing will be done with
                the missing object.
                Args for the function include:
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the original 3D object.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object, which is
                    particular to the source.
                Return is not used. Defaults to None.


        Returns:
            A dictionary mapping the object fileIdentifier to the local path of where
            the object downloaded.
        """
        hf_object_paths = cls._get_object_paths(
            download_dir=download_dir if download_dir is not None else "~/.objaverse"
        )
        if processes is None:
            processes = multiprocessing.cpu_count()

        # make a copy of the objects so we don't modify the original
        objects = objects.copy()
        objects["uid"] = objects.apply(cls._get_uid, axis=1)
        uids_to_sha256 = dict(zip(objects["uid"], objects["sha256"]))
        uids_set = set(uids_to_sha256.keys())

        # create a new df where the uids are the index
        objects_uid_index = objects.set_index("uid")

        out = {}
        objects_to_download = []
        if download_dir is None:
            for _, item in objects.iterrows():
                uid = item["uid"]
                if uid not in hf_object_paths:
                    logger.error(f"Could not find object with uid {uid}!")
                    if handle_missing_object is not None:
                        handle_missing_object(
                            file_identifier=item["fileIdentifier"],
                            sha256=item["sha256"],
                            metadata={},
                        )
                    continue
                objects_to_download.append(
                    (item["fileIdentifier"], hf_object_paths[uid], item["sha256"])
                )
        else:
            versioned_dirname = os.path.join(download_dir, "hf-objaverse-v1")
            fs, path = fsspec.core.url_to_fs(versioned_dirname)

            # Get the existing file paths. This is much faster than calling fs.exists() for each
            # file. `glob()` is like walk, but returns a list of files instead of the nested
            # directory structure. glob() is also faster than find() / walk() since it doesn't
            # need to traverse the entire directory structure.
            existing_file_paths = fs.glob(
                os.path.join(path, "glbs", "*", "*.glb"), refresh=True
            )
            existing_uids = {
                file.split("/")[-1].split(".")[0]
                for file in existing_file_paths
                if file.endswith(".glb")  # note partial files end with .glb.tmp
            }

            # add the existing downloaded uids to the return dict
            already_downloaded_uids = uids_set.intersection(existing_uids)
            for uid in already_downloaded_uids:
                hf_object_path = hf_object_paths[uid]
                fs_abs_object_path = os.path.join(versioned_dirname, hf_object_path)
                out[cls.uid_to_file_identifier(uid)] = fs_abs_object_path

            logger.info(
                f"Found {len(already_downloaded_uids)} objects already downloaded"
            )

            # get the uids that need to be downloaded
            remaining_uids = uids_set - existing_uids
            for uid in remaining_uids:
                item = objects_uid_index.loc[uid]
                if uid not in hf_object_paths:
                    logger.error(f"Could not find object with uid {uid}. Skipping it.")
                    if handle_missing_object is not None:
                        handle_missing_object(
                            file_identifier=item["fileIdentifier"],
                            sha256=item["sha256"],
                            metadata={},
                        )
                    continue
                objects_to_download.append(
                    (item["fileIdentifier"], hf_object_paths[uid], item["sha256"])
                )

            logger.info(
                f"Downloading {len(objects_to_download)} new objects across {processes} processes"
            )

        # check if all objects are already downloaded
        if len(objects_to_download) == 0:
            return out

        args = [
            (
                file_identifier,
                hf_object_path,
                download_dir,
                sha256,
                handle_found_object,
                handle_modified_object,
            )
            for file_identifier, hf_object_path, sha256 in objects_to_download
        ]

        # download the objects in parallel
        with Pool(processes) as pool:
            new_object_downloads = list(
                tqdm(
                    pool.imap_unordered(cls._parallel_download_object, args),
                    total=len(args),
                )
            )

        for file_identifier, local_path in new_object_downloads:
            out[file_identifier] = local_path

        return out

    @classmethod
    def get_lvis_annotations(
        cls,
        download_dir: str = "~/.objaverse",
    ) -> Dict[str, List[str]]:
        """Load the LVIS annotations.

        If the annotations are not already downloaded, they will be downloaded.

        Args:
            download_dir: The base directory to download the annotations to. Supports all
            file systems supported by fsspec. Defaults to "~/.objaverse".

        Returns:
            A dictionary mapping the LVIS category to the list of uids in that category.
        """
        hf_url = "https://huggingface.co/datasets/allenai/objaverse/resolve/main/lvis-annotations.json.gz"

        download_path = os.path.join(
            download_dir, "hf-objaverse-v1", "lvis-annotations.json.gz"
        )

        # use fsspec
        fs, path = fsspec.core.url_to_fs(download_path)
        if not fs.exists(path):
            # make dir if it doesn't exist
            fs.makedirs(os.path.dirname(path), exist_ok=True)

            # download the file
            with fs.open(path, "wb") as f:
                with urllib.request.urlopen(hf_url) as response:
                    f.write(response.read())

        # load the gzip file
        with fs.open(path, "rb") as f:
            with gzip.GzipFile(fileobj=f) as gfile:
                content = gfile.read()
                data = json.loads(content)

        return data

objaverse/xl/sketchfab.py
allenai/objaverse-xl
"""Script to download 3D objects from the Smithsonian Institution."""

import multiprocessing
import os
import tempfile
from multiprocessing import Pool
from typing import Callable, Dict, Optional, Tuple

import fsspec
import pandas as pd
import requests
from loguru import logger
from tqdm import tqdm

from objaverse.utils import get_file_hash, get_uid_from_str
from objaverse.xl.abstract import ObjaverseSource


class SmithsonianDownloader(ObjaverseSource):
    """Script to download objects from the Smithsonian Institute."""

    @classmethod
    def get_annotations(
        cls, download_dir: str = "~/.objaverse", refresh: bool = False
    ) -> pd.DataFrame:
        """Loads the Smithsonian Object Metadata dataset as a Pandas DataFrame.

        Args:
            download_dir (str, optional): Directory to download the parquet metadata
                file. Supports all file systems supported by fsspec. Defaults to
                "~/.objaverse".
            refresh (bool, optional): Whether to refresh the annotations by downloading
                them from the remote source. Defaults to False.

        Returns:
            pd.DataFrame: Smithsonian Object Metadata dataset as a Pandas DataFrame with
                columns for the object "title", "url", "quality", "file_type", "uid", and
                "license". The quality is always Medium and the file_type is always glb.
        """
        filename = os.path.join(download_dir, "smithsonian", "smithsonian.parquet")
        fs, path = fsspec.core.url_to_fs(filename)
        fs.makedirs(os.path.dirname(path), exist_ok=True)

        # download the parquet file if it doesn't exist
        if refresh or not fs.exists(path):
            url = "https://huggingface.co/datasets/allenai/objaverse-xl/resolve/main/smithsonian/smithsonian.parquet"
            logger.info(f"Downloading {url} to {filename}")
            response = requests.get(url)
            response.raise_for_status()
            with fs.open(path, "wb") as file:
                file.write(response.content)

        # load the parquet file with fsspec
        with fs.open(path) as f:
            df = pd.read_parquet(f)

        return df

    @classmethod
    def _download_smithsonian_object(
        cls,
        file_identifier: str,
        download_dir: Optional[str],
        expected_sha256: str,
        handle_found_object: Optional[Callable],
        handle_modified_object: Optional[Callable],
        handle_missing_object: Optional[Callable],
    ) -> Tuple[str, Optional[str]]:
        """Downloads a Smithsonian Object from a URL.

        Overwrites the file if it already exists and assumes this was previous checked.

        Args:
            file_identifier (str): URL to download the Smithsonian Object from.
            download_dir (Optional[str]): Directory to download the Smithsonian Object
                to. Supports all file systems supported by fsspec. If None, the
                Smithsonian Object will be deleted after it is downloaded and processed
                with the handler functions.
            expected_sha256 (str): The expected SHA256 of the contents of the downloaded
                object.
            handle_found_object (Optional[Callable]): Called when an object is
                successfully found and downloaded. Here, the object has the same sha256
                as the one that was downloaded with Objaverse-XL. If None, the object
                will be downloaded, but nothing will be done with it. Args for the
                function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the 3D object.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used.
            handle_modified_object (Optional[Callable]): Called when a modified object
                is found and downloaded. Here, the object is successfully downloaded,
                but it has a different sha256 than the one that was downloaded with
                Objaverse-XL. This is not expected to happen very often, because the
                same commit hash is used for each repo. If None, the object will be
                downloaded, but nothing will be done with it. Args for the function
                include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - new_sha256 (str): SHA256 of the contents of the newly downloaded 3D
                    object.
                - old_sha256 (str): Expected SHA256 of the contents of the 3D object as
                    it was when it was downloaded with Objaverse-XL.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used.
            handle_missing_object (Optional[Callable]): Called when an object that is in
                Objaverse-XL is not found. Here, it is likely that the repository was
                deleted or renamed. If None, nothing will be done with the missing
                object. Args for the function include:
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the original 3D object.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used.

        Returns:
            Tuple[str, Optional[str]]: Tuple of the URL and the path to the downloaded
                Smithsonian Object. If the Smithsonian Object was not downloaded, the path
                will be None.
        """
        uid = get_uid_from_str(file_identifier)

        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = os.path.join(temp_dir, f"{uid}.glb")
            temp_path_tmp = f"{temp_path}.tmp"

            response = requests.get(file_identifier)

            # check if the path is valid
            if response.status_code == 404:
                logger.warning(f"404 for {file_identifier}")
                if handle_missing_object is not None:
                    handle_missing_object(
                        file_identifier=file_identifier,
                        sha256=expected_sha256,
                        metadata={},
                    )
                return file_identifier, None

            with open(temp_path_tmp, "wb") as file:
                for chunk in response.iter_content(chunk_size=8192):
                    file.write(chunk)

            # rename to temp_path
            os.rename(temp_path_tmp, temp_path)

            # check the sha256
            sha256 = get_file_hash(temp_path)

            if sha256 == expected_sha256:
                if handle_found_object is not None:
                    handle_found_object(
                        local_path=temp_path,
                        file_identifier=file_identifier,
                        sha256=sha256,
                        metadata={},
                    )
            else:
                if handle_modified_object is not None:
                    handle_modified_object(
                        local_path=temp_path,
                        file_identifier=file_identifier,
                        new_sha256=sha256,
                        old_sha256=expected_sha256,
                        metadata={},
                    )

            if download_dir is not None:
                filename = os.path.join(
                    download_dir, "smithsonian", "objects", f"{uid}.glb"
                )
                fs, path = fsspec.core.url_to_fs(filename)
                fs.makedirs(os.path.dirname(path), exist_ok=True)
                fs.put(temp_path, path)
            else:
                path = None

        return file_identifier, path

    @classmethod
    def _parallel_download_object(cls, args):
        # workaround since starmap doesn't work well with tqdm
        return cls._download_smithsonian_object(*args)

    @classmethod
    def download_objects(
        cls,
        objects: pd.DataFrame,
        download_dir: Optional[str] = "~/.objaverse",
        processes: Optional[int] = None,
        handle_found_object: Optional[Callable] = None,
        handle_modified_object: Optional[Callable] = None,
        handle_missing_object: Optional[Callable] = None,
        **kwargs,
    ) -> Dict[str, str]:
        """Downloads all Smithsonian Objects.

        Args:
            objects (pd.DataFrmae): Objects to download. Must have columns for
                the object "fileIdentifier" and "sha256". Use the `get_annotations`
                function to get the metadata.
            download_dir (Optional[str], optional): Directory to download the
                Smithsonian Objects to. Supports all file systems supported by fsspec.
                If None, the Smithsonian Objects will be deleted after they are
                downloaded and processed with the handler functions. Defaults to
                "~/.objaverse".
            processes (Optional[int], optional): Number of processes to use for
                downloading the Smithsonian Objects. If None, the number of processes
                will be set to the number of CPUs on the machine
                (multiprocessing.cpu_count()). Defaults to None.
            handle_found_object (Optional[Callable], optional): Called when an object is
                successfully found and downloaded. Here, the object has the same sha256
                as the one that was downloaded with Objaverse-XL. If None, the object
                will be downloaded, but nothing will be done with it. Args for the
                function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the 3D object.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object,
                    including the GitHub organization and repo names.
                Return is not used. Defaults to None.
            handle_modified_object (Optional[Callable], optional): Called when a
                modified object is found and downloaded. Here, the object is
                successfully downloaded, but it has a different sha256 than the one that
                was downloaded with Objaverse-XL. This is not expected to happen very
                often, because the same commit hash is used for each repo. If None, the
                object will be downloaded, but nothing will be done with it. Args for
                the function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - new_sha256 (str): SHA256 of the contents of the newly downloaded 3D
                    object.
                - old_sha256 (str): Expected SHA256 of the contents of the 3D object as
                    it was when it was downloaded with Objaverse-XL.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object, which is
                    particular to the souce.
                Return is not used. Defaults to None.
            handle_missing_object (Optional[Callable], optional): Called when an object
                that is in Objaverse-XL is not found. Here, it is likely that the
                repository was deleted or renamed. If None, nothing will be done with
                the missing object.
                Args for the function include:
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the original 3D object.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object, which is
                    particular to the source.
                Return is not used. Defaults to None.

        Returns:
            Dict[str, str]: A dictionary mapping from the fileIdentifier to the
                download_path.
        """
        if processes is None:
            processes = multiprocessing.cpu_count()

        out = {}
        objects_to_download = []
        if download_dir is not None:
            objects_dir = os.path.join(download_dir, "smithsonian", "objects")
            fs, path = fsspec.core.url_to_fs(objects_dir)
            fs.makedirs(path, exist_ok=True)

            # get the existing glb files
            existing_glb_files = fs.glob(
                os.path.join(objects_dir, "*.glb"), refresh=True
            )
            existing_uids = set(
                os.path.basename(file).split(".")[0] for file in existing_glb_files
            )

            # find the urls that need to be downloaded
            already_downloaded_objects = set()
            for _, item in objects.iterrows():
                file_identifier = item["fileIdentifier"]
                uid = get_uid_from_str(file_identifier)
                if uid not in existing_uids:
                    objects_to_download.append(item)
                else:
                    already_downloaded_objects.add(file_identifier)
                out[file_identifier] = os.path.join(
                    os.path.expanduser(objects_dir), f"{uid}.glb"
                )
        else:
            existing_uids = set()
            objects_to_download = [item for _, item in objects.iterrows()]
            already_downloaded_objects = set()
            out = {}

        logger.info(
            f"Found {len(already_downloaded_objects)} Smithsonian Objects already downloaded"
        )
        logger.info(
            f"Downloading {len(objects_to_download)} Smithsonian Objects with {processes} processes"
        )

        if len(objects_to_download) == 0:
            return out

        args = [
            [
                item["fileIdentifier"],
                download_dir,
                item["sha256"],
                handle_found_object,
                handle_modified_object,
                handle_missing_object,
            ]
            for item in objects_to_download
        ]
        with Pool(processes=processes) as pool:
            results = list(
                tqdm(
                    pool.imap_unordered(cls._parallel_download_object, args),
                    total=len(objects_to_download),
                    desc="Downloading Smithsonian Objects",
                )
            )

        for file_identifier, download_path in results:
            if download_path is not None:
                out[file_identifier] = download_path

        return out

objaverse/xl/smithsonian.py
allenai/objaverse-xl
"""Script to download objects from Thingiverse."""

import multiprocessing
import os
import tempfile
import time
from multiprocessing import Pool
from typing import Callable, Dict, Optional, Tuple

import fsspec
import pandas as pd
import requests
from loguru import logger
from tqdm import tqdm

from objaverse.utils import get_file_hash
from objaverse.xl.abstract import ObjaverseSource


class ThingiverseDownloader(ObjaverseSource):
    """Script to download objects from Thingiverse."""

    @classmethod
    def get_annotations(
        cls, download_dir: str = "~/.objaverse", refresh: bool = False
    ) -> pd.DataFrame:
        """Load the annotations from the given directory.

        Args:
            download_dir (str, optional): The directory to load the annotations from.
                Supports all file systems supported by fsspec. Defaults to
                "~/.objaverse".
            refresh (bool, optional): Whether to refresh the annotations by downloading
                them from the remote source. Defaults to False.

        Returns:
            pd.DataFrame: The annotations, which includes the columns "thingId", "fileId",
                "filename", and "license".
        """
        remote_url = "https://huggingface.co/datasets/allenai/objaverse-xl/resolve/main/thingiverse/thingiverse.parquet"
        download_path = os.path.join(download_dir, "thingiverse", "thingiverse.parquet")
        fs, path = fsspec.core.url_to_fs(download_path)

        if refresh or not fs.exists(path):
            fs.makedirs(os.path.dirname(path), exist_ok=True)
            logger.info(f"Downloading {remote_url} to {download_path}")
            response = requests.get(remote_url)
            response.raise_for_status()
            with fs.open(path, "wb") as file:
                file.write(response.content)

        # read the file with pandas and fsspec
        with fs.open(download_path, "rb") as f:
            annotations_df = pd.read_parquet(f)

        return annotations_df

    @classmethod
    def _get_response_with_retries(
        cls, url: str, max_retries: int = 3, retry_delay: int = 5
    ) -> Optional[requests.models.Response]:
        """Get a response from a URL with retries.

        Args:
            url (str): The URL to get a response from.
            max_retries (int, optional): The maximum number of retries. Defaults to 3.
            retry_delay (int, optional): The delay between retries in seconds. Defaults to 5.

        Returns:
            Optional[requests.models.Response]: The response from the URL. If there was an error, returns None.
        """
        for i in range(max_retries):
            try:
                response = requests.get(url, stream=True)
                # if successful, break out of loop
                if response.status_code not in {200, 404}:
                    time.sleep(retry_delay)
                    continue
                break
            except ConnectionError:
                if i < max_retries - 1:  # i.e. not on the last try
                    time.sleep(retry_delay)
        else:
            return None

        return response

    @classmethod
    def _download_item(
        cls,
        thingi_file_id: str,
        thingi_thing_id: str,
        file_identifier: str,
        download_dir: Optional[str],
        expected_sha256: str,
        handle_found_object: Optional[Callable],
        handle_modified_object: Optional[Callable],
        handle_missing_object: Optional[Callable],
    ) -> Tuple[str, Optional[str]]:
        """Download the given item.

        Args:
            thingi_file_id (str): The Thingiverse file ID of the object.
            thingi_thing_id (str): The Thingiverse thing ID of the object.
            file_identifier (str): File identifier of the Thingiverse object.
            download_dir (Optional[str]): Directory to download the Smithsonian Object
                to. Supports all file systems supported by fsspec. If None, the
                Smithsonian Object will be deleted after it is downloaded and processed
                with the handler functions.
            expected_sha256 (str): The expected SHA256 of the contents of the downloaded
                object.
            handle_found_object (Optional[Callable]): Called when an object is
                successfully found and downloaded. Here, the object has the same sha256
                as the one that was downloaded with Objaverse-XL. If None, the object
                will be downloaded, but nothing will be done with it. Args for the
                function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the 3D object.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used.
            handle_modified_object (Optional[Callable]): Called when a modified object
                is found and downloaded. Here, the object is successfully downloaded,
                but it has a different sha256 than the one that was downloaded with
                Objaverse-XL. This is not expected to happen very often, because the
                same commit hash is used for each repo. If None, the object will be
                downloaded, but nothing will be done with it. Args for the function
                include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - new_sha256 (str): SHA256 of the contents of the newly downloaded 3D
                    object.
                - old_sha256 (str): Expected SHA256 of the contents of the 3D object as
                    it was when it was downloaded with Objaverse-XL.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used.
            handle_missing_object (Optional[Callable]): Called when an object that is in
                Objaverse-XL is not found. Here, it is likely that the repository was
                deleted or renamed. If None, nothing will be done with the missing
                object. Args for the function include:
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the original 3D object.
                - metadata (Dict[str, Any]): Metadata about the 3D object, including the
                    GitHub organization and repo names.
                Return is not used.


        Returns:
            Optional[str]: The path to the downloaded file. If there was an error or 404,
                returns None.
        """
        url = f"https://www.thingiverse.com/download:{thingi_file_id}"
        response = cls._get_response_with_retries(url)
        filename = f"thing-{thingi_thing_id}-file-{thingi_file_id}.stl"

        if response is None:
            logger.warning(
                f"Thingiverse file ID {thingi_file_id} could not get response from {url}"
            )
            # NOTE: the object is probably not missing, but the request failed
            return file_identifier, None

        # Check if the request was successful
        if response.status_code == 404:
            logger.warning(
                f"Thingiverse file ID {thingi_file_id} (404) could not find file"
            )
            if handle_missing_object is not None:
                handle_missing_object(
                    file_identifier=file_identifier, sha256=expected_sha256, metadata={}
                )
            return file_identifier, None

        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = os.path.join(temp_dir, filename)
            temp_path_tmp = temp_path + ".tmp"

            with open(temp_path_tmp, "wb") as file:
                for chunk in response.iter_content(chunk_size=8192):
                    file.write(chunk)

            # rename to temp_path
            os.rename(temp_path_tmp, temp_path)

            # check the sha256
            sha256 = get_file_hash(temp_path)

            if sha256 == expected_sha256:
                if handle_found_object is not None:
                    handle_found_object(
                        local_path=temp_path,
                        file_identifier=file_identifier,
                        sha256=sha256,
                        metadata={},
                    )
            else:
                if handle_modified_object is not None:
                    handle_modified_object(
                        local_path=temp_path,
                        file_identifier=file_identifier,
                        new_sha256=sha256,
                        old_sha256=expected_sha256,
                        metadata={},
                    )

            if download_dir is not None:
                filename = os.path.join(download_dir, filename)
                fs, path = fsspec.core.url_to_fs(filename)
                fs.makedirs(os.path.dirname(path), exist_ok=True)
                fs.put(temp_path, path)
            else:
                path = None

        return file_identifier, path

    @classmethod
    def _parallel_download_item(cls, args):
        return cls._download_item(*args)

    @classmethod
    def get_file_id_from_file_identifier(cls, file_identifier: str) -> str:
        """Get the thingiverse file ID from the Objaverse-XL file identifier.

        Args:
            file_identifier (str): The Objaverse-XL file identifier.

        Returns:
            str: The Thingiverse file ID.
        """
        return file_identifier.split("fileId=")[-1]

    @classmethod
    def get_thing_id_from_file_identifier(cls, file_identifier: str) -> str:
        """Get the thingiverse thing ID from the Objaverse-XL file identifier.

        Args:
            file_identifier (str): The Objaverse-XL file identifier.

        Returns:
            str: The Thingiverse thing ID.
        """
        return file_identifier.split("/")[-2].split(":")[1]

    @classmethod
    def download_objects(
        cls,
        objects: pd.DataFrame,
        download_dir: Optional[str] = "~/.objaverse",
        processes: Optional[int] = None,
        handle_found_object: Optional[Callable] = None,
        handle_modified_object: Optional[Callable] = None,
        handle_missing_object: Optional[Callable] = None,
        **kwargs,
    ) -> Dict[str, str]:
        """Download the objects from the given list of things and files.

        Args:
            objects (pd.DataFrame): Thingiverse objects to download. Must have columns
                for the object "fileIdentifier" and "sha256". Use the `get_annotations`
                function to get the metadata.
            download_dir (Optional[str], optional): The directory to save the files to.
                Supports all file systems supported by fsspec. If None, the objects will
                be deleted after they are downloaded. Defaults to "~/.objaverse".
            processes (int, optional): The number of processes to use. If None, maps to
                use all available CPUs using multiprocessing.cpu_count(). Defaults to
                None.
            handle_found_object (Optional[Callable], optional): Called when an object is
                successfully found and downloaded. Here, the object has the same sha256
                as the one that was downloaded with Objaverse-XL. If None, the object
                will be downloaded, but nothing will be done with it. Args for the
                function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the 3D object.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object,
                    including the GitHub organization and repo names.
                Return is not used. Defaults to None.
            handle_modified_object (Optional[Callable], optional): Called when a
                modified object is found and downloaded. Here, the object is
                successfully downloaded, but it has a different sha256 than the one that
                was downloaded with Objaverse-XL. This is not expected to happen very
                often, because the same commit hash is used for each repo. If None, the
                object will be downloaded, but nothing will be done with it. Args for
                the function include:
                - local_path (str): Local path to the downloaded 3D object.
                - file_identifier (str): File identifier of the 3D object.
                - new_sha256 (str): SHA256 of the contents of the newly downloaded 3D
                    object.
                - old_sha256 (str): Expected SHA256 of the contents of the 3D object as
                    it was when it was downloaded with Objaverse-XL.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object, which is
                    particular to the souce.
                Return is not used. Defaults to None.
            handle_missing_object (Optional[Callable], optional): Called when an object
                that is in Objaverse-XL is not found. Here, it is likely that the
                repository was deleted or renamed. If None, nothing will be done with
                the missing object.
                Args for the function include:
                - file_identifier (str): File identifier of the 3D object.
                - sha256 (str): SHA256 of the contents of the original 3D object.
                - metadata (Dict[Hashable, Any]): Metadata about the 3D object, which is
                    particular to the source.
                Return is not used. Defaults to None.

        Returns:
            Dict[str, str]: A dictionary mapping from the fileIdentifier to the path of
                the downloaded file.
        """
        if processes is None:
            processes = multiprocessing.cpu_count()

        objects = objects.copy()
        objects["thingiFileId"] = objects["fileIdentifier"].apply(
            cls.get_file_id_from_file_identifier
        )
        objects["thingiThingId"] = objects["fileIdentifier"].apply(
            cls.get_thing_id_from_file_identifier
        )

        # create the download directory
        out = {}
        if download_dir is not None:
            download_dir = os.path.join(download_dir, "thingiverse")
            fs, path = fsspec.core.url_to_fs(download_dir)
            fs.makedirs(path, exist_ok=True)

            # check to filter out files that already exist
            existing_files = fs.glob(os.path.join(download_dir, "*.stl"), refresh=True)
            existing_file_ids = {
                os.path.basename(file).split(".")[0].split("-")[-1]
                for file in existing_files
            }

            # filter out existing files
            items_to_download = []
            already_downloaded_count = 0
            for _, item in objects.iterrows():
                if item["thingiFileId"] in existing_file_ids:
                    already_downloaded_count += 1
                    out[item["fileIdentifier"]] = os.path.join(
                        os.path.expanduser(download_dir),
                        f"thing-{item['thingiThingId']}-file-{item['thingiFileId']}.stl",
                    )
                else:
                    items_to_download.append(item)

            logger.info(
                f"Found {already_downloaded_count} Thingiverse objects downloaded"
            )
        else:
            items_to_download = list(objects.to_dict(orient="records"))

        logger.info(
            f"Downloading {len(items_to_download)} Thingiverse objects with {processes=}"
        )
        if len(items_to_download) == 0:
            return out

        # download the files
        args = [
            (
                item["thingiFileId"],
                item["thingiThingId"],
                item["fileIdentifier"],
                download_dir,
                item["sha256"],
                handle_found_object,
                handle_modified_object,
                handle_missing_object,
            )
            for item in items_to_download
        ]

        with Pool(processes=processes) as pool:
            results = list(
                tqdm(
                    pool.imap_unordered(cls._parallel_download_item, args),
                    total=len(args),
                    desc="Downloading Thingiverse Objects",
                )
            )

        for file_identifier, download_path in results:
            if download_path is not None:
                out[file_identifier] = download_path

        return out

objaverse/xl/thingiverse.py
allenai/objaverse-xl
"""Blender script to render images of 3D models."""

import argparse
import json
import math
import os
import random
import sys
from typing import Any, Callable, Dict, Generator, List, Literal, Optional, Set, Tuple

import bpy
import numpy as np
from mathutils import Matrix, Vector

IMPORT_FUNCTIONS: Dict[str, Callable] = {
    "obj": bpy.ops.import_scene.obj,
    "glb": bpy.ops.import_scene.gltf,
    "gltf": bpy.ops.import_scene.gltf,
    "usd": bpy.ops.import_scene.usd,
    "fbx": bpy.ops.import_scene.fbx,
    "stl": bpy.ops.import_mesh.stl,
    "usda": bpy.ops.import_scene.usda,
    "dae": bpy.ops.wm.collada_import,
    "ply": bpy.ops.import_mesh.ply,
    "abc": bpy.ops.wm.alembic_import,
    "blend": bpy.ops.wm.append,
}


def reset_cameras() -> None:
    """Resets the cameras in the scene to a single default camera."""
    # Delete all existing cameras
    bpy.ops.object.select_all(action="DESELECT")
    bpy.ops.object.select_by_type(type="CAMERA")
    bpy.ops.object.delete()

    # Create a new camera with default properties
    bpy.ops.object.camera_add()

    # Rename the new camera to 'NewDefaultCamera'
    new_camera = bpy.context.active_object
    new_camera.name = "Camera"

    # Set the new camera as the active camera for the scene
    scene.camera = new_camera


def sample_point_on_sphere(radius: float) -> Tuple[float, float, float]:
    """Samples a point on a sphere with the given radius.

    Args:
        radius (float): Radius of the sphere.

    Returns:
        Tuple[float, float, float]: A point on the sphere.
    """
    theta = random.random() * 2 * math.pi
    phi = math.acos(2 * random.random() - 1)
    return (
        radius * math.sin(phi) * math.cos(theta),
        radius * math.sin(phi) * math.sin(theta),
        radius * math.cos(phi),
    )


def _sample_spherical(
    radius_min: float = 1.5,
    radius_max: float = 2.0,
    maxz: float = 1.6,
    minz: float = -0.75,
) -> np.ndarray:
    """Sample a random point in a spherical shell.

    Args:
        radius_min (float): Minimum radius of the spherical shell.
        radius_max (float): Maximum radius of the spherical shell.
        maxz (float): Maximum z value of the spherical shell.
        minz (float): Minimum z value of the spherical shell.

    Returns:
        np.ndarray: A random (x, y, z) point in the spherical shell.
    """
    correct = False
    vec = np.array([0, 0, 0])
    while not correct:
        vec = np.random.uniform(-1, 1, 3)
        #         vec[2] = np.abs(vec[2])
        radius = np.random.uniform(radius_min, radius_max, 1)
        vec = vec / np.linalg.norm(vec, axis=0) * radius[0]
        if maxz > vec[2] > minz:
            correct = True
    return vec


def randomize_camera(
    radius_min: float = 1.5,
    radius_max: float = 2.2,
    maxz: float = 2.2,
    minz: float = -2.2,
    only_northern_hemisphere: bool = False,
) -> bpy.types.Object:
    """Randomizes the camera location and rotation inside of a spherical shell.

    Args:
        radius_min (float, optional): Minimum radius of the spherical shell. Defaults to
            1.5.
        radius_max (float, optional): Maximum radius of the spherical shell. Defaults to
            2.0.
        maxz (float, optional): Maximum z value of the spherical shell. Defaults to 1.6.
        minz (float, optional): Minimum z value of the spherical shell. Defaults to
            -0.75.
        only_northern_hemisphere (bool, optional): Whether to only sample points in the
            northern hemisphere. Defaults to False.

    Returns:
        bpy.types.Object: The camera object.
    """

    x, y, z = _sample_spherical(
        radius_min=radius_min, radius_max=radius_max, maxz=maxz, minz=minz
    )
    camera = bpy.data.objects["Camera"]

    # only positive z
    if only_northern_hemisphere:
        z = abs(z)

    camera.location = Vector(np.array([x, y, z]))

    direction = -camera.location
    rot_quat = direction.to_track_quat("-Z", "Y")
    camera.rotation_euler = rot_quat.to_euler()

    return camera


def _set_camera_at_size(i: int, scale: float = 1.5) -> bpy.types.Object:
    """Debugging function to set the camera on the 6 faces of a cube.

    Args:
        i (int): Index of the face of the cube.
        scale (float, optional): Scale of the cube. Defaults to 1.5.

    Returns:
        bpy.types.Object: The camera object.
    """
    if i == 0:
        x, y, z = scale, 0, 0
    elif i == 1:
        x, y, z = -scale, 0, 0
    elif i == 2:
        x, y, z = 0, scale, 0
    elif i == 3:
        x, y, z = 0, -scale, 0
    elif i == 4:
        x, y, z = 0, 0, scale
    elif i == 5:
        x, y, z = 0, 0, -scale
    else:
        raise ValueError(f"Invalid index: i={i}, must be int in range [0, 5].")
    camera = bpy.data.objects["Camera"]
    camera.location = Vector(np.array([x, y, z]))
    direction = -camera.location
    rot_quat = direction.to_track_quat("-Z", "Y")
    camera.rotation_euler = rot_quat.to_euler()
    return camera


def _create_light(
    name: str,
    light_type: Literal["POINT", "SUN", "SPOT", "AREA"],
    location: Tuple[float, float, float],
    rotation: Tuple[float, float, float],
    energy: float,
    use_shadow: bool = False,
    specular_factor: float = 1.0,
):
    """Creates a light object.

    Args:
        name (str): Name of the light object.
        light_type (Literal["POINT", "SUN", "SPOT", "AREA"]): Type of the light.
        location (Tuple[float, float, float]): Location of the light.
        rotation (Tuple[float, float, float]): Rotation of the light.
        energy (float): Energy of the light.
        use_shadow (bool, optional): Whether to use shadows. Defaults to False.
        specular_factor (float, optional): Specular factor of the light. Defaults to 1.0.

    Returns:
        bpy.types.Object: The light object.
    """

    light_data = bpy.data.lights.new(name=name, type=light_type)
    light_object = bpy.data.objects.new(name, light_data)
    bpy.context.collection.objects.link(light_object)
    light_object.location = location
    light_object.rotation_euler = rotation
    light_data.use_shadow = use_shadow
    light_data.specular_factor = specular_factor
    light_data.energy = energy
    return light_object


def randomize_lighting() -> Dict[str, bpy.types.Object]:
    """Randomizes the lighting in the scene.

    Returns:
        Dict[str, bpy.types.Object]: Dictionary of the lights in the scene. The keys are
            "key_light", "fill_light", "rim_light", and "bottom_light".
    """

    # Clear existing lights
    bpy.ops.object.select_all(action="DESELECT")
    bpy.ops.object.select_by_type(type="LIGHT")
    bpy.ops.object.delete()

    # Create key light
    key_light = _create_light(
        name="Key_Light",
        light_type="SUN",
        location=(0, 0, 0),
        rotation=(0.785398, 0, -0.785398),
        energy=random.choice([3, 4, 5]),
    )

    # Create fill light
    fill_light = _create_light(
        name="Fill_Light",
        light_type="SUN",
        location=(0, 0, 0),
        rotation=(0.785398, 0, 2.35619),
        energy=random.choice([2, 3, 4]),
    )

    # Create rim light
    rim_light = _create_light(
        name="Rim_Light",
        light_type="SUN",
        location=(0, 0, 0),
        rotation=(-0.785398, 0, -3.92699),
        energy=random.choice([3, 4, 5]),
    )

    # Create bottom light
    bottom_light = _create_light(
        name="Bottom_Light",
        light_type="SUN",
        location=(0, 0, 0),
        rotation=(3.14159, 0, 0),
        energy=random.choice([1, 2, 3]),
    )

    return dict(
        key_light=key_light,
        fill_light=fill_light,
        rim_light=rim_light,
        bottom_light=bottom_light,
    )


def reset_scene() -> None:
    """Resets the scene to a clean state.

    Returns:
        None
    """
    # delete everything that isn't part of a camera or a light
    for obj in bpy.data.objects:
        if obj.type not in {"CAMERA", "LIGHT"}:
            bpy.data.objects.remove(obj, do_unlink=True)

    # delete all the materials
    for material in bpy.data.materials:
        bpy.data.materials.remove(material, do_unlink=True)

    # delete all the textures
    for texture in bpy.data.textures:
        bpy.data.textures.remove(texture, do_unlink=True)

    # delete all the images
    for image in bpy.data.images:
        bpy.data.images.remove(image, do_unlink=True)


def load_object(object_path: str) -> None:
    """Loads a model with a supported file extension into the scene.

    Args:
        object_path (str): Path to the model file.

    Raises:
        ValueError: If the file extension is not supported.

    Returns:
        None
    """
    file_extension = object_path.split(".")[-1].lower()
    if file_extension is None:
        raise ValueError(f"Unsupported file type: {object_path}")

    if file_extension == "usdz":
        # install usdz io package
        dirname = os.path.dirname(os.path.realpath(__file__))
        usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
        bpy.ops.preferences.addon_install(filepath=usdz_package)
        # enable it
        addon_name = "io_scene_usdz"
        bpy.ops.preferences.addon_enable(module=addon_name)
        # import the usdz
        from io_scene_usdz.import_usdz import import_usdz

        import_usdz(context, filepath=object_path, materials=True, animations=True)
        return None

    # load from existing import functions
    import_function = IMPORT_FUNCTIONS[file_extension]

    if file_extension == "blend":
        import_function(directory=object_path, link=False)
    elif file_extension in {"glb", "gltf"}:
        import_function(filepath=object_path, merge_vertices=True)
    else:
        import_function(filepath=object_path)


def scene_bbox(
    single_obj: Optional[bpy.types.Object] = None, ignore_matrix: bool = False
) -> Tuple[Vector, Vector]:
    """Returns the bounding box of the scene.

    Taken from Shap-E rendering script
    (https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)

    Args:
        single_obj (Optional[bpy.types.Object], optional): If not None, only computes
            the bounding box for the given object. Defaults to None.
        ignore_matrix (bool, optional): Whether to ignore the object's matrix. Defaults
            to False.

    Raises:
        RuntimeError: If there are no objects in the scene.

    Returns:
        Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
    """
    bbox_min = (math.inf,) * 3
    bbox_max = (-math.inf,) * 3
    found = False
    for obj in get_scene_meshes() if single_obj is None else [single_obj]:
        found = True
        for coord in obj.bound_box:
            coord = Vector(coord)
            if not ignore_matrix:
                coord = obj.matrix_world @ coord
            bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
            bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))

    if not found:
        raise RuntimeError("no objects in scene to compute bounding box for")

    return Vector(bbox_min), Vector(bbox_max)


def get_scene_root_objects() -> Generator[bpy.types.Object, None, None]:
    """Returns all root objects in the scene.

    Yields:
        Generator[bpy.types.Object, None, None]: Generator of all root objects in the
            scene.
    """
    for obj in bpy.context.scene.objects.values():
        if not obj.parent:
            yield obj


def get_scene_meshes() -> Generator[bpy.types.Object, None, None]:
    """Returns all meshes in the scene.

    Yields:
        Generator[bpy.types.Object, None, None]: Generator of all meshes in the scene.
    """
    for obj in bpy.context.scene.objects.values():
        if isinstance(obj.data, (bpy.types.Mesh)):
            yield obj


def get_3x4_RT_matrix_from_blender(cam: bpy.types.Object) -> Matrix:
    """Returns the 3x4 RT matrix from the given camera.

    Taken from Zero123, which in turn was taken from
    https://github.com/panmari/stanford-shapenet-renderer/blob/master/render_blender.py

    Args:
        cam (bpy.types.Object): The camera object.

    Returns:
        Matrix: The 3x4 RT matrix from the given camera.
    """
    # Use matrix_world instead to account for all constraints
    location, rotation = cam.matrix_world.decompose()[0:2]
    R_world2bcam = rotation.to_matrix().transposed()

    # Use location from matrix_world to account for constraints:
    T_world2bcam = -1 * R_world2bcam @ location

    # put into 3x4 matrix
    RT = Matrix(
        (
            R_world2bcam[0][:] + (T_world2bcam[0],),
            R_world2bcam[1][:] + (T_world2bcam[1],),
            R_world2bcam[2][:] + (T_world2bcam[2],),
        )
    )
    return RT


def delete_invisible_objects() -> None:
    """Deletes all invisible objects in the scene.

    Returns:
        None
    """
    bpy.ops.object.select_all(action="DESELECT")
    for obj in scene.objects:
        if obj.hide_viewport or obj.hide_render:
            obj.hide_viewport = False
            obj.hide_render = False
            obj.hide_select = False
            obj.select_set(True)
    bpy.ops.object.delete()

    # Delete invisible collections
    invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
    for col in invisible_collections:
        bpy.data.collections.remove(col)


def normalize_scene() -> None:
    """Normalizes the scene by scaling and translating it to fit in a unit cube centered
    at the origin.

    Mostly taken from the Point-E / Shap-E rendering script
    (https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
    but fix for multiple root objects: (see bug report here:
    https://github.com/openai/shap-e/pull/60).

    Returns:
        None
    """
    if len(list(get_scene_root_objects())) > 1:
        # create an empty object to be used as a parent for all root objects
        parent_empty = bpy.data.objects.new("ParentEmpty", None)
        bpy.context.scene.collection.objects.link(parent_empty)

        # parent all root objects to the empty object
        for obj in get_scene_root_objects():
            if obj != parent_empty:
                obj.parent = parent_empty

    bbox_min, bbox_max = scene_bbox()
    scale = 1 / max(bbox_max - bbox_min)
    for obj in get_scene_root_objects():
        obj.scale = obj.scale * scale

    # Apply scale to matrix_world.
    bpy.context.view_layer.update()
    bbox_min, bbox_max = scene_bbox()
    offset = -(bbox_min + bbox_max) / 2
    for obj in get_scene_root_objects():
        obj.matrix_world.translation += offset
    bpy.ops.object.select_all(action="DESELECT")

    # unparent the camera
    bpy.data.objects["Camera"].parent = None


def delete_missing_textures() -> Dict[str, Any]:
    """Deletes all missing textures in the scene.

    Returns:
        Dict[str, Any]: Dictionary with keys "count", "files", and "file_path_to_color".
            "count" is the number of missing textures, "files" is a list of the missing
            texture file paths, and "file_path_to_color" is a dictionary mapping the
            missing texture file paths to a random color.
    """
    missing_file_count = 0
    out_files = []
    file_path_to_color = {}

    # Check all materials in the scene
    for material in bpy.data.materials:
        if material.use_nodes:
            for node in material.node_tree.nodes:
                if node.type == "TEX_IMAGE":
                    image = node.image
                    if image is not None:
                        file_path = bpy.path.abspath(image.filepath)
                        if file_path == "":
                            # means it's embedded
                            continue

                        if not os.path.exists(file_path):
                            # Find the connected Principled BSDF node
                            connected_node = node.outputs[0].links[0].to_node

                            if connected_node.type == "BSDF_PRINCIPLED":
                                if file_path not in file_path_to_color:
                                    # Set a random color for the unique missing file path
                                    random_color = [random.random() for _ in range(3)]
                                    file_path_to_color[file_path] = random_color + [1]

                                connected_node.inputs[
                                    "Base Color"
                                ].default_value = file_path_to_color[file_path]

                            # Delete the TEX_IMAGE node
                            material.node_tree.nodes.remove(node)
                            missing_file_count += 1
                            out_files.append(image.filepath)
    return {
        "count": missing_file_count,
        "files": out_files,
        "file_path_to_color": file_path_to_color,
    }


def _get_random_color() -> Tuple[float, float, float, float]:
    """Generates a random RGB-A color.

    The alpha value is always 1.

    Returns:
        Tuple[float, float, float, float]: A random RGB-A color. Each value is in the
        range [0, 1].
    """
    return (random.random(), random.random(), random.random(), 1)


def _apply_color_to_object(
    obj: bpy.types.Object, color: Tuple[float, float, float, float]
) -> None:
    """Applies the given color to the object.

    Args:
        obj (bpy.types.Object): The object to apply the color to.
        color (Tuple[float, float, float, float]): The color to apply to the object.

    Returns:
        None
    """
    mat = bpy.data.materials.new(name=f"RandomMaterial_{obj.name}")
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    principled_bsdf = nodes.get("Principled BSDF")
    if principled_bsdf:
        principled_bsdf.inputs["Base Color"].default_value = color
    obj.data.materials.append(mat)


def apply_single_random_color_to_all_objects() -> Tuple[float, float, float, float]:
    """Applies a single random color to all objects in the scene.

    Returns:
        Tuple[float, float, float, float]: The random color that was applied to all
        objects.
    """
    rand_color = _get_random_color()
    for obj in bpy.context.scene.objects:
        if obj.type == "MESH":
            _apply_color_to_object(obj, rand_color)
    return rand_color


class MetadataExtractor:
    """Class to extract metadata from a Blender scene."""

    def __init__(
        self, object_path: str, scene: bpy.types.Scene, bdata: bpy.types.BlendData
    ) -> None:
        """Initializes the MetadataExtractor.

        Args:
            object_path (str): Path to the object file.
            scene (bpy.types.Scene): The current scene object from `bpy.context.scene`.
            bdata (bpy.types.BlendData): The current blender data from `bpy.data`.

        Returns:
            None
        """
        self.object_path = object_path
        self.scene = scene
        self.bdata = bdata

    def get_poly_count(self) -> int:
        """Returns the total number of polygons in the scene."""
        total_poly_count = 0
        for obj in self.scene.objects:
            if obj.type == "MESH":
                total_poly_count += len(obj.data.polygons)
        return total_poly_count

    def get_vertex_count(self) -> int:
        """Returns the total number of vertices in the scene."""
        total_vertex_count = 0
        for obj in self.scene.objects:
            if obj.type == "MESH":
                total_vertex_count += len(obj.data.vertices)
        return total_vertex_count

    def get_edge_count(self) -> int:
        """Returns the total number of edges in the scene."""
        total_edge_count = 0
        for obj in self.scene.objects:
            if obj.type == "MESH":
                total_edge_count += len(obj.data.edges)
        return total_edge_count

    def get_lamp_count(self) -> int:
        """Returns the number of lamps in the scene."""
        return sum(1 for obj in self.scene.objects if obj.type == "LIGHT")

    def get_mesh_count(self) -> int:
        """Returns the number of meshes in the scene."""
        return sum(1 for obj in self.scene.objects if obj.type == "MESH")

    def get_material_count(self) -> int:
        """Returns the number of materials in the scene."""
        return len(self.bdata.materials)

    def get_object_count(self) -> int:
        """Returns the number of objects in the scene."""
        return len(self.bdata.objects)

    def get_animation_count(self) -> int:
        """Returns the number of animations in the scene."""
        return len(self.bdata.actions)

    def get_linked_files(self) -> List[str]:
        """Returns the filepaths of all linked files."""
        image_filepaths = self._get_image_filepaths()
        material_filepaths = self._get_material_filepaths()
        linked_libraries_filepaths = self._get_linked_libraries_filepaths()

        all_filepaths = (
            image_filepaths | material_filepaths | linked_libraries_filepaths
        )
        if "" in all_filepaths:
            all_filepaths.remove("")
        return list(all_filepaths)

    def _get_image_filepaths(self) -> Set[str]:
        """Returns the filepaths of all images used in the scene."""
        filepaths = set()
        for image in self.bdata.images:
            if image.source == "FILE":
                filepaths.add(bpy.path.abspath(image.filepath))
        return filepaths

    def _get_material_filepaths(self) -> Set[str]:
        """Returns the filepaths of all images used in materials."""
        filepaths = set()
        for material in self.bdata.materials:
            if material.use_nodes:
                for node in material.node_tree.nodes:
                    if node.type == "TEX_IMAGE":
                        image = node.image
                        if image is not None:
                            filepaths.add(bpy.path.abspath(image.filepath))
        return filepaths

    def _get_linked_libraries_filepaths(self) -> Set[str]:
        """Returns the filepaths of all linked libraries."""
        filepaths = set()
        for library in self.bdata.libraries:
            filepaths.add(bpy.path.abspath(library.filepath))
        return filepaths

    def get_scene_size(self) -> Dict[str, list]:
        """Returns the size of the scene bounds in meters."""
        bbox_min, bbox_max = scene_bbox()
        return {"bbox_max": list(bbox_max), "bbox_min": list(bbox_min)}

    def get_shape_key_count(self) -> int:
        """Returns the number of shape keys in the scene."""
        total_shape_key_count = 0
        for obj in self.scene.objects:
            if obj.type == "MESH":
                shape_keys = obj.data.shape_keys
                if shape_keys is not None:
                    total_shape_key_count += (
                        len(shape_keys.key_blocks) - 1
                    )  # Subtract 1 to exclude the Basis shape key
        return total_shape_key_count

    def get_armature_count(self) -> int:
        """Returns the number of armatures in the scene."""
        total_armature_count = 0
        for obj in self.scene.objects:
            if obj.type == "ARMATURE":
                total_armature_count += 1
        return total_armature_count

    def read_file_size(self) -> int:
        """Returns the size of the file in bytes."""
        return os.path.getsize(self.object_path)

    def get_metadata(self) -> Dict[str, Any]:
        """Returns the metadata of the scene.

        Returns:
            Dict[str, Any]: Dictionary of the metadata with keys for "file_size",
            "poly_count", "vert_count", "edge_count", "material_count", "object_count",
            "lamp_count", "mesh_count", "animation_count", "linked_files", "scene_size",
            "shape_key_count", and "armature_count".
        """
        return {
            "file_size": self.read_file_size(),
            "poly_count": self.get_poly_count(),
            "vert_count": self.get_vertex_count(),
            "edge_count": self.get_edge_count(),
            "material_count": self.get_material_count(),
            "object_count": self.get_object_count(),
            "lamp_count": self.get_lamp_count(),
            "mesh_count": self.get_mesh_count(),
            "animation_count": self.get_animation_count(),
            "linked_files": self.get_linked_files(),
            "scene_size": self.get_scene_size(),
            "shape_key_count": self.get_shape_key_count(),
            "armature_count": self.get_armature_count(),
        }


def render_object(
    object_file: str,
    num_renders: int,
    only_northern_hemisphere: bool,
    output_dir: str,
) -> None:
    """Saves rendered images with its camera matrix and metadata of the object.

    Args:
        object_file (str): Path to the object file.
        num_renders (int): Number of renders to save of the object.
        only_northern_hemisphere (bool): Whether to only render sides of the object that
            are in the northern hemisphere. This is useful for rendering objects that
            are photogrammetrically scanned, as the bottom of the object often has
            holes.
        output_dir (str): Path to the directory where the rendered images and metadata
            will be saved.

    Returns:
        None
    """
    os.makedirs(output_dir, exist_ok=True)

    # load the object
    if object_file.endswith(".blend"):
        bpy.ops.object.mode_set(mode="OBJECT")
        reset_cameras()
        delete_invisible_objects()
    else:
        reset_scene()
        load_object(object_file)

    # Set up cameras
    cam = scene.objects["Camera"]
    cam.data.lens = 35
    cam.data.sensor_width = 32

    # Set up camera constraints
    cam_constraint = cam.constraints.new(type="TRACK_TO")
    cam_constraint.track_axis = "TRACK_NEGATIVE_Z"
    cam_constraint.up_axis = "UP_Y"
    empty = bpy.data.objects.new("Empty", None)
    scene.collection.objects.link(empty)
    cam_constraint.target = empty

    # Extract the metadata. This must be done before normalizing the scene to get
    # accurate bounding box information.
    metadata_extractor = MetadataExtractor(
        object_path=object_file, scene=scene, bdata=bpy.data
    )
    metadata = metadata_extractor.get_metadata()

    # delete all objects that are not meshes
    if object_file.lower().endswith(".usdz"):
        # don't delete missing textures on usdz files, lots of them are embedded
        missing_textures = None
    else:
        missing_textures = delete_missing_textures()
    metadata["missing_textures"] = missing_textures

    # possibly apply a random color to all objects
    if object_file.endswith(".stl") or object_file.endswith(".ply"):
        assert len(bpy.context.selected_objects) == 1
        rand_color = apply_single_random_color_to_all_objects()
        metadata["random_color"] = rand_color
    else:
        metadata["random_color"] = None

    # save metadata
    metadata_path = os.path.join(output_dir, "metadata.json")
    os.makedirs(os.path.dirname(metadata_path), exist_ok=True)
    with open(metadata_path, "w", encoding="utf-8") as f:
        json.dump(metadata, f, sort_keys=True, indent=2)

    # normalize the scene
    normalize_scene()

    # randomize the lighting
    randomize_lighting()

    # render the images
    for i in range(num_renders):
        # set camera
        camera = randomize_camera(
            only_northern_hemisphere=only_northern_hemisphere,
        )

        # render the image
        render_path = os.path.join(output_dir, f"{i:03d}.png")
        scene.render.filepath = render_path
        bpy.ops.render.render(write_still=True)

        # save camera RT matrix
        rt_matrix = get_3x4_RT_matrix_from_blender(camera)
        rt_matrix_path = os.path.join(output_dir, f"{i:03d}.npy")
        np.save(rt_matrix_path, rt_matrix)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--object_path",
        type=str,
        required=True,
        help="Path to the object file",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        required=True,
        help="Path to the directory where the rendered images and metadata will be saved.",
    )
    parser.add_argument(
        "--engine",
        type=str,
        default="BLENDER_EEVEE",
        choices=["CYCLES", "BLENDER_EEVEE"],
    )
    parser.add_argument(
        "--only_northern_hemisphere",
        action="store_true",
        help="Only render the northern hemisphere of the object.",
        default=False,
    )
    parser.add_argument(
        "--num_renders",
        type=int,
        default=12,
        help="Number of renders to save of the object.",
    )
    argv = sys.argv[sys.argv.index("--") + 1 :]
    args = parser.parse_args(argv)

    context = bpy.context
    scene = context.scene
    render = scene.render

    # Set render settings
    render.engine = args.engine
    render.image_settings.file_format = "PNG"
    render.image_settings.color_mode = "RGBA"
    render.resolution_x = 512
    render.resolution_y = 512
    render.resolution_percentage = 100

    # Set cycles settings
    scene.cycles.device = "GPU"
    scene.cycles.samples = 128
    scene.cycles.diffuse_bounces = 1
    scene.cycles.glossy_bounces = 1
    scene.cycles.transparent_max_bounces = 3
    scene.cycles.transmission_bounces = 3
    scene.cycles.filter_width = 0.01
    scene.cycles.use_denoising = True
    scene.render.film_transparent = True
    bpy.context.preferences.addons["cycles"].preferences.get_devices()
    bpy.context.preferences.addons[
        "cycles"
    ].preferences.compute_device_type = "CUDA"  # or "OPENCL"

    # Render the images
    render_object(
        object_file=args.object_path,
        num_renders=args.num_renders,
        only_northern_hemisphere=args.only_northern_hemisphere,
        output_dir=args.output_dir,
    )

scripts/rendering/blender_script.py
allenai/objaverse-xl
import glob
import json
import multiprocessing
import os
import platform
import random
import subprocess
import tempfile
import time
import zipfile
from functools import partial
from typing import Any, Dict, List, Literal, Optional, Union

import fire
import fsspec
import GPUtil
import pandas as pd
from loguru import logger

import objaverse.xl as oxl
from objaverse.utils import get_uid_from_str


def log_processed_object(csv_filename: str, *args) -> None:
    """Log when an object is done being used.

    Args:
        csv_filename (str): Name of the CSV file to save the logs to.
        *args: Arguments to save to the CSV file.

    Returns:
        None
    """
    args = ",".join([str(arg) for arg in args])
    # log that this object was rendered successfully
    # saving locally to avoid excessive writes to the cloud
    dirname = os.path.expanduser(f"~/.objaverse/logs/")
    os.makedirs(dirname, exist_ok=True)
    with open(os.path.join(dirname, csv_filename), "a", encoding="utf-8") as f:
        f.write(f"{time.time()},{args}\n")


def zipdir(path: str, ziph: zipfile.ZipFile) -> None:
    """Zip up a directory with an arcname structure.

    Args:
        path (str): Path to the directory to zip.
        ziph (zipfile.ZipFile): ZipFile handler object to write to.

    Returns:
        None
    """
    # ziph is zipfile handle
    for root, dirs, files in os.walk(path):
        for file in files:
            # this ensures the structure inside the zip starts at folder/
            arcname = os.path.join(os.path.basename(root), file)
            ziph.write(os.path.join(root, file), arcname=arcname)


def handle_found_object(
    local_path: str,
    file_identifier: str,
    sha256: str,
    metadata: Dict[str, Any],
    num_renders: int,
    render_dir: str,
    only_northern_hemisphere: bool,
    gpu_devices: Union[int, List[int]],
    render_timeout: int,
    successful_log_file: Optional[str] = "handle-found-object-successful.csv",
    failed_log_file: Optional[str] = "handle-found-object-failed.csv",
) -> bool:
    """Called when an object is successfully found and downloaded.

    Here, the object has the same sha256 as the one that was downloaded with
    Objaverse-XL. If None, the object will be downloaded, but nothing will be done with
    it.

    Args:
        local_path (str): Local path to the downloaded 3D object.
        file_identifier (str): File identifier of the 3D object.
        sha256 (str): SHA256 of the contents of the 3D object.
        metadata (Dict[str, Any]): Metadata about the 3D object, such as the GitHub
            organization and repo names.
        num_renders (int): Number of renders to save of the object.
        render_dir (str): Directory where the objects will be rendered.
        only_northern_hemisphere (bool): Only render the northern hemisphere of the
            object.
        gpu_devices (Union[int, List[int]]): GPU device(s) to use for rendering. If
            an int, the GPU device will be randomly selected from 0 to gpu_devices - 1.
            If a list, the GPU device will be randomly selected from the list.
            If 0, the CPU will be used for rendering.
        render_timeout (int): Number of seconds to wait for the rendering job to
            complete.
        successful_log_file (str): Name of the log file to save successful renders to.
        failed_log_file (str): Name of the log file to save failed renders to.

    Returns: True if the object was rendered successfully, False otherwise.
    """
    save_uid = get_uid_from_str(file_identifier)
    args = f"--object_path '{local_path}' --num_renders {num_renders}"

    # get the GPU to use for rendering
    using_gpu: bool = True
    gpu_i = 0
    if isinstance(gpu_devices, int) and gpu_devices > 0:
        num_gpus = gpu_devices
        gpu_i = random.randint(0, num_gpus - 1)
    elif isinstance(gpu_devices, list):
        gpu_i = random.choice(gpu_devices)
    elif isinstance(gpu_devices, int) and gpu_devices == 0:
        using_gpu = False
    else:
        raise ValueError(
            f"gpu_devices must be an int > 0, 0, or a list of ints. Got {gpu_devices}."
        )

    with tempfile.TemporaryDirectory() as temp_dir:
        # get the target directory for the rendering job
        target_directory = os.path.join(temp_dir, save_uid)
        os.makedirs(target_directory, exist_ok=True)
        args += f" --output_dir {target_directory}"

        # check for Linux / Ubuntu or MacOS
        if platform.system() == "Linux" and using_gpu:
            args += " --engine BLENDER_EEVEE"
        elif platform.system() == "Darwin" or (
            platform.system() == "Linux" and not using_gpu
        ):
            # As far as I know, MacOS does not support BLENER_EEVEE, which uses GPU
            # rendering. Generally, I'd only recommend using MacOS for debugging and
            # small rendering jobs, since CYCLES is much slower than BLENDER_EEVEE.
            args += " --engine CYCLES"
        else:
            raise NotImplementedError(f"Platform {platform.system()} is not supported.")

        # check if we should only render the northern hemisphere
        if only_northern_hemisphere:
            args += " --only_northern_hemisphere"

        # get the command to run
        command = f"blender-3.2.2-linux-x64/blender --background --python blender_script.py -- {args}"
        if using_gpu:
            command = f"export DISPLAY=:0.{gpu_i} && {command}"

        # render the object (put in dev null)
        subprocess.run(
            ["bash", "-c", command],
            timeout=render_timeout,
            check=False,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )

        # check that the renders were saved successfully
        png_files = glob.glob(os.path.join(target_directory, "*.png"))
        metadata_files = glob.glob(os.path.join(target_directory, "*.json"))
        npy_files = glob.glob(os.path.join(target_directory, "*.npy"))
        if (
            (len(png_files) != num_renders)
            or (len(npy_files) != num_renders)
            or (len(metadata_files) != 1)
        ):
            logger.error(
                f"Found object {file_identifier} was not rendered successfully!"
            )
            if failed_log_file is not None:
                log_processed_object(
                    failed_log_file,
                    file_identifier,
                    sha256,
                )
            return False

        # update the metadata
        metadata_path = os.path.join(target_directory, "metadata.json")
        with open(metadata_path, "r", encoding="utf-8") as f:
            metadata_file = json.load(f)
        metadata_file["sha256"] = sha256
        metadata_file["file_identifier"] = file_identifier
        metadata_file["save_uid"] = save_uid
        metadata_file["metadata"] = metadata
        with open(metadata_path, "w", encoding="utf-8") as f:
            json.dump(metadata_file, f, indent=2, sort_keys=True)

        # Make a zip of the target_directory.
        # Keeps the {save_uid} directory structure when unzipped
        with zipfile.ZipFile(
            f"{target_directory}.zip", "w", zipfile.ZIP_DEFLATED
        ) as ziph:
            zipdir(target_directory, ziph)

        # move the zip to the render_dir
        fs, path = fsspec.core.url_to_fs(render_dir)

        # move the zip to the render_dir
        fs.makedirs(os.path.join(path, "renders"), exist_ok=True)
        fs.put(
            os.path.join(f"{target_directory}.zip"),
            os.path.join(path, "renders", f"{save_uid}.zip"),
        )

        # log that this object was rendered successfully
        if successful_log_file is not None:
            log_processed_object(successful_log_file, file_identifier, sha256)

        return True


def handle_new_object(
    local_path: str,
    file_identifier: str,
    sha256: str,
    metadata: Dict[str, Any],
    log_file: str = "handle-new-object.csv",
) -> None:
    """Called when a new object is found.

    Here, the object is not used in Objaverse-XL, but is still downloaded with the
    repository. The object may have not been used because it does not successfully
    import into Blender. If None, the object will be downloaded, but nothing will be
    done with it.

    Args:
        local_path (str): Local path to the downloaded 3D object.
        file_identifier (str): The file identifier of the new 3D object.
        sha256 (str): SHA256 of the contents of the 3D object.
        metadata (Dict[str, Any]): Metadata about the 3D object, including the GitHub
            organization and repo names.
        log_file (str): Name of the log file to save the handle_new_object logs to.

    Returns:
        None
    """
    # log the new object
    log_processed_object(log_file, file_identifier, sha256)


def handle_modified_object(
    local_path: str,
    file_identifier: str,
    new_sha256: str,
    old_sha256: str,
    metadata: Dict[str, Any],
    num_renders: int,
    render_dir: str,
    only_northern_hemisphere: bool,
    gpu_devices: Union[int, List[int]],
    render_timeout: int,
) -> None:
    """Called when a modified object is found and downloaded.

    Here, the object is successfully downloaded, but it has a different sha256 than the
    one that was downloaded with Objaverse-XL. This is not expected to happen very
    often, because the same commit hash is used for each repo. If None, the object will
    be downloaded, but nothing will be done with it.

    Args:
        local_path (str): Local path to the downloaded 3D object.
        file_identifier (str): File identifier of the 3D object.
        new_sha256 (str): SHA256 of the contents of the newly downloaded 3D object.
        old_sha256 (str): Expected SHA256 of the contents of the 3D object as it was
            when it was downloaded with Objaverse-XL.
        metadata (Dict[str, Any]): Metadata about the 3D object, such as the GitHub
            organization and repo names.
        num_renders (int): Number of renders to save of the object.
        render_dir (str): Directory where the objects will be rendered.
        only_northern_hemisphere (bool): Only render the northern hemisphere of the
            object.
        gpu_devices (Union[int, List[int]]): GPU device(s) to use for rendering. If
            an int, the GPU device will be randomly selected from 0 to gpu_devices - 1.
            If a list, the GPU device will be randomly selected from the list.
            If 0, the CPU will be used for rendering.
        render_timeout (int): Number of seconds to wait for the rendering job to
            complete.

    Returns:
        None
    """
    success = handle_found_object(
        local_path=local_path,
        file_identifier=file_identifier,
        sha256=new_sha256,
        metadata=metadata,
        num_renders=num_renders,
        render_dir=render_dir,
        only_northern_hemisphere=only_northern_hemisphere,
        gpu_devices=gpu_devices,
        render_timeout=render_timeout,
        successful_log_file=None,
        failed_log_file=None,
    )

    if success:
        log_processed_object(
            "handle-modified-object-successful.csv",
            file_identifier,
            old_sha256,
            new_sha256,
        )
    else:
        log_processed_object(
            "handle-modified-object-failed.csv",
            file_identifier,
            old_sha256,
            new_sha256,
        )


def handle_missing_object(
    file_identifier: str,
    sha256: str,
    metadata: Dict[str, Any],
    log_file: str = "handle-missing-object.csv",
) -> None:
    """Called when an object that is in Objaverse-XL is not found.

    Here, it is likely that the repository was deleted or renamed. If None, nothing
    will be done with the missing object.

    Args:
        file_identifier (str): File identifier of the 3D object.
        sha256 (str): SHA256 of the contents of the original 3D object.
        metadata (Dict[str, Any]): Metadata about the 3D object, such as the GitHub
            organization and repo names.
        log_file (str): Name of the log file to save missing renders to.

    Returns:
        None
    """
    # log the missing object
    log_processed_object(log_file, file_identifier, sha256)


def get_example_objects() -> pd.DataFrame:
    """Returns a DataFrame of example objects to use for debugging."""
    return pd.read_json("example-objects.json", orient="records")


def render_objects(
    render_dir: str = "~/.objaverse",
    download_dir: Optional[str] = None,
    num_renders: int = 12,
    processes: Optional[int] = None,
    save_repo_format: Optional[Literal["zip", "tar", "tar.gz", "files"]] = None,
    only_northern_hemisphere: bool = False,
    render_timeout: int = 300,
    gpu_devices: Optional[Union[int, List[int]]] = None,
) -> None:
    """Renders objects in the Objaverse-XL dataset with Blender

    Args:
        render_dir (str, optional): Directory where the objects will be rendered.
        download_dir (Optional[str], optional): Directory where the objects will be
            downloaded. If None, the objects will not be downloaded. Defaults to None.
        num_renders (int, optional): Number of renders to save of the object. Defaults
            to 12.
        processes (Optional[int], optional): Number of processes to use for downloading
            the objects. If None, defaults to multiprocessing.cpu_count() * 3. Defaults
            to None.
        save_repo_format (Optional[Literal["zip", "tar", "tar.gz", "files"]], optional):
            If not None, the GitHub repo will be deleted after rendering each object
            from it.
        only_northern_hemisphere (bool, optional): Only render the northern hemisphere
            of the object. Useful for rendering objects that are obtained from
            photogrammetry, since the southern hemisphere is often has holes. Defaults
            to False.
        render_timeout (int, optional): Number of seconds to wait for the rendering job
            to complete. Defaults to 300.
        gpu_devices (Optional[Union[int, List[int]]], optional): GPU device(s) to use
            for rendering. If an int, the GPU device will be randomly selected from 0 to
            gpu_devices - 1. If a list, the GPU device will be randomly selected from
            the list. If 0, the CPU will be used for rendering. If None, all available
            GPUs will be used. Defaults to None.

    Returns:
        None
    """
    if platform.system() not in ["Linux", "Darwin"]:
        raise NotImplementedError(
            f"Platform {platform.system()} is not supported. Use Linux or MacOS."
        )
    if download_dir is None and save_repo_format is not None:
        raise ValueError(
            f"If {save_repo_format=} is not None, {download_dir=} must be specified."
        )
    if download_dir is not None and save_repo_format is None:
        logger.warning(
            f"GitHub repos will not save. While {download_dir=} is specified, {save_repo_format=} None."
        )

    # get the gpu devices to use
    parsed_gpu_devices: Union[int, List[int]] = 0
    if gpu_devices is None:
        parsed_gpu_devices = len(GPUtil.getGPUs())
    logger.info(f"Using {parsed_gpu_devices} GPU devices for rendering.")

    if processes is None:
        processes = multiprocessing.cpu_count() * 3

    # get the objects to render
    objects = get_example_objects()
    objects.iloc[0]["fileIdentifier"]
    objects = objects.copy()
    logger.info(f"Provided {len(objects)} objects to render.")

    # get the already rendered objects
    fs, path = fsspec.core.url_to_fs(render_dir)
    try:
        zip_files = fs.glob(os.path.join(path, "renders", "*.zip"), refresh=True)
    except TypeError:
        # s3fs may not support refresh depending on the version
        zip_files = fs.glob(os.path.join(path, "renders", "*.zip"))

    saved_ids = set(zip_file.split("/")[-1].split(".")[0] for zip_file in zip_files)
    logger.info(f"Found {len(saved_ids)} objects already rendered.")

    # filter out the already rendered objects
    objects["saveUid"] = objects["fileIdentifier"].apply(get_uid_from_str)
    objects = objects[~objects["saveUid"].isin(saved_ids)]
    objects = objects.reset_index(drop=True)
    logger.info(f"Rendering {len(objects)} new objects.")

    # shuffle the objects
    objects = objects.sample(frac=1).reset_index(drop=True)

    oxl.download_objects(
        objects=objects,
        processes=processes,
        save_repo_format=save_repo_format,
        download_dir=download_dir,
        handle_found_object=partial(
            handle_found_object,
            render_dir=render_dir,
            num_renders=num_renders,
            only_northern_hemisphere=only_northern_hemisphere,
            gpu_devices=parsed_gpu_devices,
            render_timeout=render_timeout,
        ),
        handle_new_object=handle_new_object,
        handle_modified_object=partial(
            handle_modified_object,
            render_dir=render_dir,
            num_renders=num_renders,
            only_northern_hemisphere=only_northern_hemisphere,
            gpu_devices=parsed_gpu_devices,
            render_timeout=render_timeout,
        ),
        handle_missing_object=handle_missing_object,
    )


if __name__ == "__main__":
    fire.Fire(render_objects)

scripts/rendering/main.py
allenai/objaverse-xl
# Taken from https://github.com/allenai/ai2thor/blob/main/scripts/ai2thor-xorg
# Starts an x-server to support running Blender on a headless machine with
# dedicated NVIDIA GPUs

import argparse

#!/usr/bin/env python3
import os
import platform
import re
import shlex
import signal
import subprocess
import sys
import time

# Turning off automatic black formatting for this script as it breaks quotes.
# fmt: off
from typing import List

PID_FILE = "/var/run/ai2thor-xorg.pid"
CONFIG_FILE = "/tmp/ai2thor-xorg.conf"

DEFAULT_HEIGHT = 768
DEFAULT_WIDTH = 1024


def process_alive(pid):
    """
    Use kill(0) to determine if pid is alive
    :param pid: process id
    :rtype: bool
    """
    try:
        os.kill(pid, 0)
    except OSError:
        return False

    return True


def find_devices(excluded_device_ids):
    devices = []
    id_counter = 0
    for r in pci_records():
        if r.get("Vendor", "") == "NVIDIA Corporation" and r["Class"] in [
            "VGA compatible controller",
            "3D controller",
        ]:
            bus_id = "PCI:" + ":".join(
                map(lambda x: str(int(x, 16)), re.split(r"[:\.]", r["Slot"]))
            )

            if id_counter not in excluded_device_ids:
                devices.append(bus_id)

            id_counter += 1

    if not devices:
        print("Error: ai2thor-xorg requires at least one NVIDIA device")
        sys.exit(1)

    return devices

def active_display_bus_ids():
    # this determines whether a monitor is connected to the GPU
    # if one is, the following Option is added for the Screen "UseDisplayDevice" "None"
    command = "nvidia-smi --query-gpu=pci.bus_id,display_active --format=csv,noheader"
    active_bus_ids = set()
    result = subprocess.run(command, shell=True, stdout=subprocess.PIPE)
    if result.returncode == 0:
        for line in result.stdout.decode().strip().split("\n"):
            nvidia_bus_id, display_status = re.split(r",\s?", line.strip())
            bus_id = "PCI:" + ":".join(
                map(lambda x: str(int(x, 16)), re.split(r"[:\.]", nvidia_bus_id)[1:])
            )
            if display_status.lower() == "enabled":
                active_bus_ids.add(bus_id)

    return active_bus_ids

def pci_records():
    records = []
    command = shlex.split("lspci -vmm")
    output = subprocess.check_output(command).decode()

    for devices in output.strip().split("\n\n"):
        record = {}
        records.append(record)
        for row in devices.split("\n"):
            key, value = row.split("\t")
            record[key.split(":")[0]] = value

    return records


def read_pid():
    if os.path.isfile(PID_FILE):
        with open(PID_FILE) as f:
            return int(f.read())
    else:
        return None


def start(display: str, excluded_device_ids: List[int], width: int, height: int):
    pid = read_pid()

    if pid and process_alive(pid):
        print("Error: ai2thor-xorg is already running with pid: %s" % pid)
        sys.exit(1)

    with open(CONFIG_FILE, "w") as f:
        f.write(generate_xorg_conf(excluded_device_ids, width=width, height=height))

    log_file = "/var/log/ai2thor-xorg.%s.log" % display
    error_log_file = "/var/log/ai2thor-xorg-error.%s.log" % display
    command = shlex.split(
        "Xorg -quiet -maxclients 1024 -noreset +extension GLX +extension RANDR +extension RENDER -logfile %s -config %s :%s"
        % (log_file, CONFIG_FILE, display)
    )

    pid = None
    with open(error_log_file, "w") as error_log_f:
        proc = subprocess.Popen(command, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=error_log_f)
        pid = proc.pid
        try:
            proc.wait(timeout=0.25)
        except subprocess.TimeoutExpired:
            pass

    if pid and process_alive(pid):
        with open(PID_FILE, "w") as f:
            f.write(str(proc.pid))
    else:
        print("Error: error with command '%s'" % " ".join(command))
        with open(error_log_file, "r") as f:
            print(f.read())


def print_config(excluded_device_ids: List[int], width: int, height: int):
    print(generate_xorg_conf(excluded_device_ids, width=width, height=height))


def stop():
    pid = read_pid()
    if pid and process_alive(pid):
        os.kill(pid, signal.SIGTERM)

        for i in range(10):
            time.sleep(0.2)
            if not process_alive(pid):
                os.unlink(PID_FILE)
                break


def generate_xorg_conf(
        excluded_device_ids: List[int], width: int, height: int
):
    devices = find_devices(excluded_device_ids)
    active_display_devices = active_display_bus_ids()

    xorg_conf = []

    device_section = """
Section "Device"
    Identifier     "Device{device_id}"
    Driver         "nvidia"
    VendorName     "NVIDIA Corporation"
    BusID          "{bus_id}"
EndSection
"""
    server_layout_section = """
Section "ServerLayout"
    Identifier     "Layout0"
    {screen_records}
EndSection
"""
    screen_section = """
Section "Screen"
    Identifier     "Screen{screen_id}"
    Device         "Device{device_id}"
    DefaultDepth    24
    Option         "AllowEmptyInitialConfiguration" "True"
    Option         "Interactive" "False"
    {extra_options}
    SubSection     "Display"
        Depth       24
        Virtual {width} {height}
    EndSubSection
EndSection
"""
    screen_records = []
    for i, bus_id in enumerate(devices):
        extra_options = ""
        if bus_id in active_display_devices:
            # See https://github.com/allenai/ai2thor/pull/990
            # when a monitor is connected, this option must be used otherwise
            # Xorg will fail to start
            extra_options = 'Option         "UseDisplayDevice" "None"'
        xorg_conf.append(device_section.format(device_id=i, bus_id=bus_id))
        xorg_conf.append(screen_section.format(device_id=i, screen_id=i, width=width, height=height, extra_options=extra_options))
        screen_records.append(
            'Screen {screen_id} "Screen{screen_id}" 0 0'.format(screen_id=i)
        )

    xorg_conf.append(
        server_layout_section.format(screen_records="\n    ".join(screen_records))
    )

    output = "\n".join(xorg_conf)
    return output


# fmt: on

if __name__ == "__main__":
    if os.geteuid() != 0:
        path = os.path.abspath(__file__)
        print("Executing ai2thor-xorg with sudo")
        args = ["--", path] + sys.argv[1:]
        os.execvp("sudo", args)

    if platform.system() != "Linux":
        print("Error: Can only run ai2thor-xorg on linux")
        sys.exit(1)

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--exclude-device",
        help="exclude a specific GPU device",
        action="append",
        type=int,
        default=[],
    )
    parser.add_argument(
        "--width",
        help="width of the screen to start (should be greater than the maximum"
        f" width of any ai2thor instance you will start) [default: {DEFAULT_WIDTH}]",
        type=int,
        default=DEFAULT_WIDTH,
    )
    parser.add_argument(
        "--height",
        help="height of the screen to start (should be greater than the maximum"
        f" height of any ai2thor instance you will start) [default: {DEFAULT_HEIGHT}]",
        type=int,
        default=DEFAULT_HEIGHT,
    )
    parser.add_argument(
        "command",
        help="command to be executed",
        choices=["start", "stop", "print-config"],
    )
    parser.add_argument(
        "display", help="display to be used", nargs="?", type=int, default=0
    )
    args = parser.parse_args()
    if args.command == "start":
        start(
            display=args.display,
            excluded_device_ids=args.exclude_device,
            height=args.height,
            width=args.width,
        )
    elif args.command == "stop":
        stop()
    elif args.command == "print-config":
        print_config(
            excluded_device_ids=args.exclude_device,
            width=args.width,
            height=args.height,
        )

scripts/rendering/start_x_server.py
allenai/objaverse-xl
from setuptools import find_packages, setup

# Read the README content
with open("README.md", "r", encoding="utf-8") as fh:
    long_description = fh.read()

# Read requirements from requirements.txt
with open("requirements.txt", "r", encoding="utf-8") as f:
    requirements = [line.strip() for line in f.readlines()]

setup(
    name="objaverse",
    version="0.1.5",
    author="Allen Institute for AI",
    author_email="mattd@allenai.org",
    description="Objaverse is an open dataset with over 10 million 3D objects",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://huggingface.co/datasets/allenai/objaverse-xl",
    packages=find_packages(),
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: Apache Software License",
    ],
    install_requires=requirements,
    extras_require={"s3": ["s3fs==0.4.2"]},
)

setup.py
allenai/objaverse-xl
import os
import shutil

import fsspec
import pandas as pd

from objaverse.xl.github import GitHubDownloader
from objaverse.xl.sketchfab import SketchfabDownloader
from objaverse.xl.thingiverse import ThingiverseDownloader
from objaverse.xl.smithsonian import SmithsonianDownloader


def test_github_process_repo():
    github_downloader = GitHubDownloader()
    download_dir = "~/.objaverse-tests"
    base_download_dir = os.path.join(download_dir, "github")
    fs, path = fsspec.core.url_to_fs(base_download_dir)
    fs.makedirs(path, exist_ok=True)

    new_objects = []
    handle_new_object = (
        lambda local_path, file_identifier, sha256, metadata: new_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    for save_repo_format in ["tar", "tar.gz", "zip", "files"]:
        shutil.rmtree(os.path.join(path, "repos"), ignore_errors=True)
        out = github_downloader._process_repo(
            repo_id="mattdeitke/objaverse-xl-test-files",
            fs=fs,
            base_dir=path,
            save_repo_format=save_repo_format,  # type: ignore
            expected_objects=dict(),
            handle_found_object=None,
            handle_modified_object=None,
            handle_missing_object=None,
            handle_new_object=handle_new_object,
            commit_hash="6928b08a2501aa7a4a4aabac1f888b66e7782056",
        )

        # test that the sha256's are correct
        assert len(out) == 0
        sha256s = [x["sha256"] for x in new_objects]
        for sha256 in [
            "d2b9a5d7c47dc93526082c9b630157ab6bce4fd8669610d942176f4a36444e71",
            "04e6377317d6818e32c5cbd1951e76deb3641bbf4f6db6933046221d5fbf1c5c",
            "7037575f47816118e5a34e7c0da9927e1be7be3f5b4adfac337710822eb50fa9",
        ]:
            assert sha256 in sha256s, f"{sha256=} not in {sha256s=}"
        github_urls = [x["file_identifier"] for x in new_objects]
        for github_url in [
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.fbx",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.glb",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.obj",
        ]:
            assert github_url in github_urls, f"{github_url=} not in {github_urls=}"

        # test that the files are correct
        if save_repo_format != "files":
            assert fs.exists(
                os.path.join(
                    path,
                    "repos",
                    "mattdeitke",
                    f"objaverse-xl-test-files.{save_repo_format}",
                )
            )
        else:
            assert fs.exists(
                os.path.join(
                    base_download_dir, "repos", "mattdeitke", "objaverse-xl-test-files"
                )
            )


def test_github_handle_new_object():
    github_downloader = GitHubDownloader()
    found_objects = []
    handle_found_object = (
        lambda local_path, file_identifier, sha256, metadata: found_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    missing_objects = []
    handle_missing_object = (
        lambda file_identifier, sha256, metadata: missing_objects.append(
            dict(
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    new_objects = []
    handle_new_object = (
        lambda local_path, file_identifier, sha256, metadata: new_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    modified_objects = []
    handle_modified_object = lambda local_path, file_identifier, new_sha256, old_sha256, metadata: modified_objects.append(
        dict(
            local_path=local_path,
            file_identifier=file_identifier,
            new_sha256=new_sha256,
            old_sha256=old_sha256,
            metadata=metadata,
        )
    )

    download_dir = "~/.objaverse-tests"
    base_download_dir = os.path.join(download_dir, "github")
    fs, path = fsspec.core.url_to_fs(base_download_dir)
    fs.makedirs(path, exist_ok=True)

    shutil.rmtree(os.path.join(path, "repos"), ignore_errors=True)
    out = github_downloader._process_repo(
        repo_id="mattdeitke/objaverse-xl-test-files",
        fs=fs,
        base_dir=path,
        save_repo_format=None,
        expected_objects=dict(),
        handle_found_object=handle_found_object,
        handle_modified_object=handle_modified_object,
        handle_missing_object=handle_missing_object,
        handle_new_object=handle_new_object,
        commit_hash="6928b08a2501aa7a4a4aabac1f888b66e7782056",
    )

    assert len(out) == 0
    assert len(new_objects) == 3
    assert len(found_objects) == 0
    assert len(modified_objects) == 0
    assert len(missing_objects) == 0


def test_github_handle_found_object():
    github_downloader = GitHubDownloader()
    found_objects = []
    handle_found_object = (
        lambda local_path, file_identifier, sha256, metadata: found_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    missing_objects = []
    handle_missing_object = (
        lambda file_identifier, sha256, metadata: missing_objects.append(
            dict(
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    new_objects = []
    handle_new_object = (
        lambda local_path, file_identifier, sha256, metadata: new_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    modified_objects = []
    handle_modified_object = lambda local_path, file_identifier, new_sha256, old_sha256, metadata: modified_objects.append(
        dict(
            local_path=local_path,
            file_identifier=file_identifier,
            new_sha256=new_sha256,
            old_sha256=old_sha256,
            metadata=metadata,
        )
    )

    download_dir = "~/.objaverse-tests"
    base_download_dir = os.path.join(download_dir, "github")
    fs, path = fsspec.core.url_to_fs(base_download_dir)
    fs.makedirs(path, exist_ok=True)

    shutil.rmtree(os.path.join(path, "repos"), ignore_errors=True)
    out = github_downloader._process_repo(
        repo_id="mattdeitke/objaverse-xl-test-files",
        fs=fs,
        base_dir=path,
        save_repo_format=None,
        expected_objects={
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.fbx": "7037575f47816118e5a34e7c0da9927e1be7be3f5b4adfac337710822eb50fa9",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.glb": "04e6377317d6818e32c5cbd1951e76deb3641bbf4f6db6933046221d5fbf1c5c",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.obj": "d2b9a5d7c47dc93526082c9b630157ab6bce4fd8669610d942176f4a36444e71",
        },
        handle_found_object=handle_found_object,
        handle_modified_object=handle_modified_object,
        handle_missing_object=handle_missing_object,
        handle_new_object=handle_new_object,
        commit_hash="6928b08a2501aa7a4a4aabac1f888b66e7782056",
    )

    assert len(out) == 0
    assert len(found_objects) == 3
    assert len(missing_objects) == 0
    assert len(new_objects) == 0
    assert len(modified_objects) == 0


def test_github_handle_modified_object():
    github_downloader = GitHubDownloader()
    found_objects = []
    handle_found_object = (
        lambda local_path, file_identifier, sha256, metadata: found_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    missing_objects = []
    handle_missing_object = (
        lambda file_identifier, sha256, metadata: missing_objects.append(
            dict(
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    new_objects = []
    handle_new_object = (
        lambda local_path, file_identifier, sha256, metadata: new_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    modified_objects = []
    handle_modified_object = lambda local_path, file_identifier, new_sha256, old_sha256, metadata: modified_objects.append(
        dict(
            local_path=local_path,
            file_identifier=file_identifier,
            new_sha256=new_sha256,
            old_sha256=old_sha256,
            metadata=metadata,
        )
    )

    download_dir = "~/.objaverse-tests"
    base_download_dir = os.path.join(download_dir, "github")
    fs, path = fsspec.core.url_to_fs(base_download_dir)
    fs.makedirs(path, exist_ok=True)

    shutil.rmtree(os.path.join(path, "repos"), ignore_errors=True)
    out = github_downloader._process_repo(
        repo_id="mattdeitke/objaverse-xl-test-files",
        fs=fs,
        base_dir=path,
        save_repo_format=None,
        expected_objects={
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.fbx": "7037575f47816118e5a34e7c0da9927e1be7be3f5b4adfac337710822eb50fa9<modified>",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.glb": "04e6377317d6818e32c5cbd1951e76deb3641bbf4f6db6933046221d5fbf1c5c",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.obj": "d2b9a5d7c47dc93526082c9b630157ab6bce4fd8669610d942176f4a36444e71",
        },
        handle_found_object=handle_found_object,
        handle_modified_object=handle_modified_object,
        handle_missing_object=handle_missing_object,
        handle_new_object=handle_new_object,
        commit_hash="6928b08a2501aa7a4a4aabac1f888b66e7782056",
    )

    assert len(out) == 0
    assert len(found_objects) == 2
    assert len(missing_objects) == 0
    assert len(new_objects) == 0
    assert len(modified_objects) == 1


def test_github_handle_missing_object():
    github_downloader = GitHubDownloader()
    found_objects = []
    handle_found_object = (
        lambda local_path, file_identifier, sha256, metadata: found_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    missing_objects = []
    handle_missing_object = (
        lambda file_identifier, sha256, metadata: missing_objects.append(
            dict(
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    new_objects = []
    handle_new_object = (
        lambda local_path, file_identifier, sha256, metadata: new_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    modified_objects = []
    handle_modified_object = lambda local_path, file_identifier, new_sha256, old_sha256, metadata: modified_objects.append(
        dict(
            local_path=local_path,
            file_identifier=file_identifier,
            new_sha256=new_sha256,
            old_sha256=old_sha256,
            metadata=metadata,
        )
    )

    download_dir = "~/.objaverse-tests"
    base_download_dir = os.path.join(download_dir, "github")
    fs, path = fsspec.core.url_to_fs(base_download_dir)
    fs.makedirs(path, exist_ok=True)

    shutil.rmtree(os.path.join(path, "repos"), ignore_errors=True)
    out = github_downloader._process_repo(
        repo_id="mattdeitke/objaverse-xl-test-files",
        fs=fs,
        base_dir=path,
        save_repo_format=None,
        expected_objects={
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.fbx": "7037575f47816118e5a34e7c0da9927e1be7be3f5b4adfac337710822eb50fa9",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example-2.fbx": "<fake-missing-object>",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.glb": "04e6377317d6818e32c5cbd1951e76deb3641bbf4f6db6933046221d5fbf1c5c",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.obj": "d2b9a5d7c47dc93526082c9b630157ab6bce4fd8669610d942176f4a36444e71",
        },
        handle_found_object=handle_found_object,
        handle_modified_object=handle_modified_object,
        handle_missing_object=handle_missing_object,
        handle_new_object=handle_new_object,
        commit_hash="6928b08a2501aa7a4a4aabac1f888b66e7782056",
    )

    assert len(out) == 0
    assert len(found_objects) == 3
    assert len(missing_objects) == 1
    assert len(new_objects) == 0
    assert len(modified_objects) == 0


def test_github_handle_missing_object_2():
    github_downloader = GitHubDownloader()
    found_objects = []
    handle_found_object = (
        lambda local_path, file_identifier, sha256, metadata: found_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    missing_objects = []
    handle_missing_object = (
        lambda file_identifier, sha256, metadata: missing_objects.append(
            dict(
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    new_objects = []
    handle_new_object = (
        lambda local_path, file_identifier, sha256, metadata: new_objects.append(
            dict(
                local_path=local_path,
                file_identifier=file_identifier,
                sha256=sha256,
                metadata=metadata,
            )
        )
    )

    modified_objects = []
    handle_modified_object = lambda local_path, file_identifier, new_sha256, old_sha256, metadata: modified_objects.append(
        dict(
            local_path=local_path,
            file_identifier=file_identifier,
            new_sha256=new_sha256,
            old_sha256=old_sha256,
            metadata=metadata,
        )
    )

    download_dir = "~/.objaverse-tests"
    base_download_dir = os.path.join(download_dir, "github")
    fs, path = fsspec.core.url_to_fs(base_download_dir)
    fs.makedirs(path, exist_ok=True)

    shutil.rmtree(os.path.join(path, "repos"), ignore_errors=True)
    out = github_downloader._process_repo(
        repo_id="mattdeitke/objaverse-xl-test-files-does-not-exist",
        fs=fs,
        base_dir=path,
        save_repo_format=None,
        expected_objects={
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.fbx": "7037575f47816118e5a34e7c0da9927e1be7be3f5b4adfac337710822eb50fa9",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example-2.fbx": "<fake-missing-object>",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.glb": "04e6377317d6818e32c5cbd1951e76deb3641bbf4f6db6933046221d5fbf1c5c",
            "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.obj": "d2b9a5d7c47dc93526082c9b630157ab6bce4fd8669610d942176f4a36444e71",
        },
        handle_found_object=handle_found_object,
        handle_modified_object=handle_modified_object,
        handle_missing_object=handle_missing_object,
        handle_new_object=handle_new_object,
        commit_hash="6928b08a2501aa7a4a4aabac1f888b66e7782056",
    )

    assert len(out) == 0
    assert len(found_objects) == 0
    assert len(missing_objects) == 4
    assert len(new_objects) == 0
    assert len(modified_objects) == 0


def test_github_download_cache():
    github_downloader = GitHubDownloader()
    objects = pd.DataFrame(
        [
            {
                "fileIdentifier": "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.fbx",
                "license": None,
                "sha256": "7037575f47816118e5a34e7c0da9927e1be7be3f5b4adfac337710822eb50fa9",
            },
            {
                "fileIdentifier": "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.glb",
                "license": None,
                "sha256": "04e6377317d6818e32c5cbd1951e76deb3641bbf4f6db6933046221d5fbf1c5c",
            },
            {
                "fileIdentifier": "https://github.com/mattdeitke/objaverse-xl-test-files/blob/6928b08a2501aa7a4a4aabac1f888b66e7782056/example.obj",
                "license": None,
                "sha256": "d2b9a5d7c47dc93526082c9b630157ab6bce4fd8669610d942176f4a36444e71",
            },
        ]
    )

    # remove the repos directory
    for save_repo_format in ["tar", "tar.gz", "zip", "files"]:
        repos_dir = "~/.objaverse-tests/github/repos"
        shutil.rmtree(os.path.expanduser(repos_dir), ignore_errors=True)

        out = github_downloader.download_objects(
            objects=objects,
            processes=1,
            download_dir="~/.objaverse-tests",
            save_repo_format=save_repo_format,  # type: ignore
        )
        assert len(out) == 3

        out = github_downloader.download_objects(
            objects=objects,
            processes=1,
            download_dir="~/.objaverse-tests",
            save_repo_format=save_repo_format,  # type: ignore
        )
        assert len(out) == 0


def test_annotations():
    downloaders = [
        GitHubDownloader(),
        SketchfabDownloader(),
        SmithsonianDownloader(),
        ThingiverseDownloader(),
    ]

    for downloader in downloaders:
        annotations_df = downloader.get_annotations()

        # make sure the columns are
        assert set(annotations_df.columns) == set(
            ["fileIdentifier", "source", "license", "fileType", "sha256", "metadata"]
        )


def test_download_objects():
    downloaders = [
        GitHubDownloader(),
        SketchfabDownloader(),
        SmithsonianDownloader(),
        ThingiverseDownloader(),
    ]

    download_dir = "~/.objaverse-tests"

    for downloader in downloaders:
        shutil.rmtree(os.path.expanduser(download_dir), ignore_errors=True)

        annotations_df = downloader.get_annotations()

        test_objects = annotations_df.head(n=2)

        out = downloader.download_objects(
            objects=test_objects,
            download_dir=download_dir,
            processes=2,
        )
        assert isinstance(out, dict), f"{out=}"

tests/test_api.py
allenai/objaverse-xl