Image to point cloud with Point-E

2 min read -

Arnaud Hillen - January 18th, 2023 -

Jump to section

Just before Christmas, OpenAI published its Point-E model. During the holidays, we took it for a spin and experimented with how easy it would be to integrate the model into a Segments.ai workflow. Read with me, or give it a run for its money yourself.

Colab notebook:

? Point-E

Point-E is a deep learning model created by OpenAI that transforms a text caption into a colored point cloud. More specifically, Point-E consists of three steps, each handled by a dedicated ML model:

Generate an image conditioned on a text caption
Create a point cloud (1024 points) conditioned on the image
Upsample the point cloud (to 4096 points) conditioned on the image and low-resolution point cloud

In this experiment we skip the first step and instead create a point cloud based on an image. This process typically results in higher-quality point clouds.

Let’s zoom in on step 2. Point-E uses a so-called diffusion model to generate point clouds. Intuitively, this model was trained to gradually remove noise from a point cloud. By initially giving it an input that is pure noise and repeatedly feeding its outputs to its inputs, we eventually end up with a clean point cloud. So the model takes in three inputs:

A noisy point cloud
A time step to keep track of how far we are in the diffusion process
A vector representation of an image, on which the denoising process is conditioned. In this case, the vector representation is a CLIP embedding. CLIP is an OpenAI model trained on text-image pairs and used here to create a meaningful vector representation of an image.

The authors created several million image-point cloud pairs to train the model by taking 3D renders (i.e., a 2D image from a 3D model) of a Blender model.

Input and output of the Point-E diffusion model (image-to-point cloud step). A CLIP embedding, timestep and point cloud with noise are input. A denoised point cloud is the output. Image taken from Point-E paper.

Create a point cloud

Copy to Clipboard

from PIL import Image
import torch
from tqdm.auto import tqdm

from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config
from point_e.diffusion.sampler import PointCloudSampler
from point_e.models.download import load_checkpoint
from point_e.models.configs import MODEL_CONFIGS, model_from_config
from point_e.util.plotting import plot_point_cloud

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("Creating base model")
base_name = "base300M"  # Use base1B for better results
base_model = model_from_config(MODEL_CONFIGS[base_name], device)
base_model.eval()
base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[base_name])

print("Creating upsample model")
upsampler_model = model_from_config(MODEL_CONFIGS["upsample"], device)
upsampler_model.eval()
upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS["upsample"])

print("Downloading base checkpoint")
base_model.load_state_dict(load_checkpoint(base_name, device))

print("Downloading upsampler checkpoint")
upsampler_model.load_state_dict(load_checkpoint("upsample", device))

# Combine the image-to-point cloud and upsampler model
sampler = PointCloudSampler(
    device=device,
    models=[base_model, upsampler_model],
    diffusions=[base_diffusion, upsampler_diffusion],
    num_points=[1024, 4096 - 1024],
    aux_channels=["R", "G", "B"],
    guidance_scale=[3.0, 3.0],
)

Choose an image and paste its path here.

Copy to Clipboard

Christmas tree image with gifts in front of it

Christmas tree – Image taken from freepik

Create a point cloud.

Copy to Clipboard

Let’s look at the point cloud before we upload it to Segments.ai.

Copy to Clipboard

Upload a point cloud to Segments.ai

To be able to store, manage and manipulate our point cloud. Let’s upload it to Segments.ai. We will upload it to a point cloud segmentation dataset so that you can label individual points.

Create an account

If you don’t yet have an account on Segments.ai, you can create a free account for data labeling here (don’t worry, we don’t ask for your credit card ?).

Copy to Clipboard

Create a dataset

Install the Segments.ai Python SDK.

Copy to Clipboard

Copy your API key from the settings page and create a dataset.

Copy to Clipboard

Upload a point cloud

To upload our point cloud, we use a function upload_pcd_to_segments (you can check out the colab notebook if you’re interested in implementing this function ?). It takes in the point cloud (positions and colors), sample name, and dataset name and uploads the point cloud to the dataset.

Copy to Clipboard

That’s it. Now you can use your point cloud on Segments.ai. Curious what point clouds you will make! If you have any questions, feel free to reach out at arnaud@segments.ai or support@segments.ai.