export_scene.py

#!/usr/bin/env python

import json

import imgviz
import numpy as np
import path
import trimesh
import trimesh.transformations as tf

import morefusion


models = morefusion.datasets.YCBVideoModels()
dataset = morefusion.datasets.YCBVideoDataset("val")

example = dataset[0]

meta = example["meta"]

here = path.Path(__file__).abspath().parent
out_dir = here / "scene"
out_dir.makedirs_p()

data = {"poses": []}

scene = trimesh.Scene()
for i, cls_id in enumerate(meta["cls_indexes"]):
    pose = meta["poses"][:, :, i]
    cad = models.get_cad(cls_id)
    transform = np.r_[pose, [[0, 0, 0, 1]]]
    scene.add_geometry(cad, transform=transform)

    (out_dir / "cad").makedirs_p()
    class_name = models.class_names[cls_id]
    cad_file = models.get_cad_file(cls_id)
    for file in cad_file.parent.listdir():
        if file.basename() in ["textured.obj", "textured.mtl", "texture_map.png"]:
            (out_dir / "cad" / class_name).makedirs_p()
            file.copy2(out_dir / "cad" / class_name / file.basename())
    data["poses"].append(
        {
            "class_id": int(cls_id),
            "class_name": str(class_name),
            "position": list(tf.translation_from_matrix(transform)),
            "quaternion": list(
                tf.quaternion_from_matrix(transform)[[1, 2, 3, 0]]
            ),
        }
    )

H, W = example["color"].shape[:2]
K = meta["intrinsic_matrix"]
camera = trimesh.scene.Camera(resolution=(W, H), focal=(K[0, 0], -K[1, 1]))

scene.add_geometry(trimesh.creation.camera_marker(camera))
scene.camera_transform = morefusion.extra.trimesh.to_opengl_transform()
# scene.show()

with open(out_dir / "meta.json", "w") as f:
    json.dump(data, f, indent=2)

imgviz.io.imsave("scene/rgb.jpg", example["color"])