export.py

#!/usr/bin/env python

import json

import numpy as np
import path
import trimesh

import morefusion



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

example = dataset[0]

meta = example["meta"]

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

data = {"intrinsic_matrix": meta["intrinsic_matrix"].tolist(), "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)

    (here / "spam" / "cad").makedirs_p()
    class_name = models.class_names[cls_id]
    cad_file = models.get_cad_file(cls_id)
    cad_file.copy2(here / "spam" / "cad" / class_name + ".obj")
    data["poses"].append(dict(class_id=int(cls_id), class_name=str(class_name), T_cad2cam=transform.tolist()))

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(here / "spam/meta.json", "w") as f:
    json.dump(data, f, indent=2)