This repository is developed by Tae Ha "Jeff" Park at Space Rendezvous Laboratory (SLAB) of Stanford University.

• [2024/03/18] The udpated dataset (v1.1) is now available at the Stanford Digital Repository. This version corrects bugs with the pose labels in the initial deposit, so the users are encouraged to download the updated version. Please see UPDATES.md inside the new dataset for a detailed update log.

This is the official PyTorch implementation of the paper titled Rapid Abstraction of Spacecraft 3D Structure from Single 2D Image.

This paper presents a Convolutional Neural Network (CNN) to simultaneously abstract the 3D structure of the target space resident object and estimate its pose from a single 2D image. Specifically, the CNN predicts from a single image of the target a unit-size assembly of superquadric primitives which can individually describe a wide range of simple 3D shapes (e.g., cuboid, ellipsoid) using only a few parameters. The proposed training pipeline employs various types of supervision in both 2D and 3D spaces to fit an assembly of superquadrics to man-made satellite structures. In order to avoid numerical instability encountered when evaluating superquadrics, this work proposes a novel, numerically stable algorithm based on dual superquadrics to evaluate a point on the surface of and inside a superquadric for all shape parameters. Furthermore, in order to train the CNN, this work also introduces SPE3R, a novel dataset comprising 64 different satellite models and 1,000 images, binary masks and pose labels for each model. The experimental studies reveal that the proposed CNN can be trained to reconstruct accurate superquadric assemblies when tested on unseen images of known models and capture high-level structures of the unknown models most of the time despite having been trained on an extremely small dataset.

This repository uses poetry to manage virtual environment and library dependencies. It is developed and tested on Ubuntu 22.04 and trained on a single NVIDIA GeForce RTX 4090 24GB GPU.

1. Install poetry and create a virtual environment for python 3.10 as shown in the pyproject.toml file.

2. Install dependencies specified in pyproject.toml by running

   poetry install

3. Separately install kaleido via

   pip install -U kaleido

   which is necessary to use plotly for visualization. It seems there is no poetry support for it yet.

4. Compile the extension module via

   python setup.py build_ext --inplace

5. [Optional] Get pre-trained model from here. It is trained with M = 8 primitives for both RGH and grayscale image inputs.

First, obtain the SPE3R dataset from Stanford Digital Repository and place it at ROOT. You can use tools/get_spe3r.sh to unzip all files

sh tools/get_spe3r.sh /path/to/dataset

Note that this dataset is slightly different from the one used in the experiments reported in the paper. There is a small difference in model composition for the validation set.

tools/preprocess.py will create 100,000 occupancy labels (occupancy_points.npz) and surface points (surface_points.npz) for each model. They will also save images of these points for visual verification.

python tools/preprocess.py --cfg experiments/config.yaml

You will want to set the following before running the script.

You can also change LOSS.RECON_TYPE, LOSS.POSE_TYPE, LOSS.REG_TYPE to change the main supervised losses, pose losses, and regularizations to be used in the training.

The training can then be done via

python tools/train.py --cfg experiments/config.yaml

Below command will save the output superquadric mesh and input image to SAVE_DIR. Unless --modelidx and --imageidx are explicited provided, it will evaluate on a random image of a random model for the given data split.

SPLIT=validation
SAVE_DIR=figures
NUM_PRIM=8
PRETRAIN=output/model_m8.pth.tar

python tools/evaluate.py --cfg experiments/config.yaml --split ${SPLIT} --save_dir ${SAVE_DIR} \
        MODEL.NUM_MAX_PRIMITIVES ${NUM_PRIM} MODEL.PRETRAIN_FILE ${PRETRAIN}

This repository is released under the MIT License (see LICENSE.md for more details).

@inbook{park_2024_scitech_spe3r,
    author = {Park, Tae Ha and D'Amico, Simone},
    title = {Rapid Abstraction of Spacecraft 3D Structure from Single 2D Image},
    booktitle = {AIAA SCITECH 2024 Forum},
    chapter = {},
    pages = {},
    doi = {10.2514/6.2024-2768},
}