テキストから3Dモデルを生成してみた / 開発者向けブログ・イベント

こんにちは。GMOインターネットグループ株式会社の新里です。

最近はDALL・E 2、Midjourney、Stable Diffusionなどを使って画像を自動生成するのはよく見かけますね。他にもOpenAIが出したChatGPTでテキストで対話するモデルだったり、機械学習は本当に日進月歩で進んでいるように感じます。ここでは、OpenAIが2022年末に公開したPoint-Eを使って、テキストから3Dモデルを生成してみます。

テキストから3Dを生成～3Dプリンタで印刷

入力としてテキストから3Dモデルを生成するものは色々ありました。
Point-E （OpenAI）、Dream Fusion（Google）、Magic3D（NVIDIA）などですね。GoogleはImagen（Text to Image）のモデルは公開していないので、Stable Diffusionを利用したStable Fusionというのも公開されていました。

ここではPoint-Eを使って3Dモデルを生成、さらにそのモデルを3Dプリンタで印刷する所までやってみます。

Point-Eのコード

お手軽にGoogle Colab上でPoint-Eを動かしてみます。githubにサンプルコードがあるので、そのままGoogle Colabに持ってくれば動きました。テキストから点群を作ってくれるのは text2pointcloud.ipynb ですね。あと、メッシュにする pointcloud2mesh.ipynb を使えばPLYファイルを生成・OBJなどに変換して3Dプリンタで印刷できそうです。

どちらのコードも非常にシンプルでマージしてしまえば、サクッと動かせそうです。

!pip install -U scikit-image
!pip install git+https://github.com/openai/point-e
!nvidia-smi

import torch
from PIL import Image
from tqdm.auto import tqdm
import matplotlib.pyplot as plt
from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config
from point_e.diffusion.sampler import PointCloudSampler
from point_e.models.configs import MODEL_CONFIGS, model_from_config
from point_e.util.plotting import plot_point_cloud
from point_e.models.download import load_checkpoint
from point_e.util.pc_to_mesh import marching_cubes_mesh
from point_e.util.point_cloud import PointCloud

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

# モデル群をダウンロード
base_name = 'base40M-textvec'
base_model = model_from_config(MODEL_CONFIGS[base_name], device)
base_model.eval()
base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[base_name])
upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device)
upsampler_model.eval()
upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample'])
base_model.load_state_dict(load_checkpoint(base_name, device))
upsampler_model.load_state_dict(load_checkpoint('upsample', device))

# 生成する点群の設定
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, 0.0],
    model_kwargs_key_filter=('texts', ''),
)

# "一人がけのソファで、ゆったりとした背もたれ" というプロンプト
prompt = 'One-person sofas with spacious backrests.'

# 3Dモデルの生成と表示
samples = None
for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(texts=[prompt]))):
    samples = x
pc = sampler.output_to_point_clouds(samples)[0]
fig = plot_point_cloud(pc, grid_size=3, fixed_bounds=((-0.75, -0.75, -0.75),(0.75, 0.75, 0.75)))

# ここから点群をメッシュ化
name = 'sdf'
model = model_from_config(MODEL_CONFIGS[name], device)
model.eval()
model.load_state_dict(load_checkpoint(name, device))

mesh = marching_cubes_mesh(
    pc=pc,
    model=model,
    batch_size=4096,
    grid_size=128,
    progress=True,
)

# PLYとして保存
with open('model.ply', 'wb') as f:
    mesh.write_ply(f)