<jupyter_start><jupyter_text>DIMO viewer<jupyter_code>ds_path = Path('/home/peter/Projects/Baekeland/keylines/data/raw/keylines/dimo_blocks_small/dimo_small')
dimo_loader = DimoLoader()
dimo_ds = dimo_loader.load(ds_path, cameras=['real_jaigo', 'real_d415', 'real_l515', 'real_mvbluefox', 'sim_jaigo'])<jupyter_output><empty_output><jupyter_text>Rendering setup<jupyter_code>def t_from_xyz(x=0., y=0., z=0.):
    t = np.eye(4, 4)
    t[:3, 3] = [x, y, z]
    return t

# OpenCV (dimo) -> OpenGL (pyrender) conversion
def transform_viewpoint(viewpoint):
    rotate_x_180 = np.eye(4, 4)
    rotate_x_180[:3, :3] = Rotation.from_euler('x', [180.], degrees=True).as_matrix()
    return viewpoint @ rotate_x_180

def setup_renderer(im_size):
    im_w = im_size[0]
    im_h = im_size[1]
    p_renderer = pyrender.OffscreenRenderer(im_w, im_h)
    p_scene = pyrender.Scene(bg_color=[0., 0., 0., 0.])
    return p_renderer, p_scene

def update_camera(camera, fx, fy, cx, cy):
    camera.fx = fx
    camera.fy = fy
    camera.cx = cx
    camera.cy = cy

def add_lights(p_scene):
    p_scene.ambient_light = [0.5, 0.5, 0.5]
    light = pyrender.DirectionalLight(color=[1., 1., 1.], intensity=5.)
    light = p_scene.add(light, pose=transform_viewpoint(t_from_xyz(z=500.)))

def prepare_image(image, models, p_scene):
    p_scene.clear()
    add_lights(p_scene)
    K = image['camera']['K']
    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
    p_camera = p_scene.add(pyrender.IntrinsicsCamera(fx, fy, cx, cy, 0.05, 1000))
    p_scene.set_pose(p_camera, transform_viewpoint(t_from_xyz()))
    for obj in image['objects']:
        mesh = models[obj['id']].copy()
        mesh.visual.vertex_colors = trimesh.visual.color.random_color()
        o = p_scene.add(pyrender.Mesh.from_trimesh(mesh))
        obj_pose = obj['model_2cam']
        o.translation =  obj_pose[:3, 3].copy()
        o.rotation = Rotation.from_matrix(obj_pose[:3, :3]).as_quat()
        
def load_image(path):
    image = Image.open(path)
    image = image.convert('RGBA')
    return image
def load_models(models):
    return {model['id']: trimesh.load(model['cad']) for model in models}
cameras = [c for c in dimo_ds.keys() if c != 'models']
models = load_models(dimo_ds['models'])

try:
    p_renderer.delete()
except:
    pass

def show_camera(camera):
    scenes = dimo_ds[camera]
    first_img = load_image(scenes[0]['images'][0]['path'])
    
    p_renderer, p_scene = setup_renderer(first_img.size)
    
    def show_scene(i):
        scene = scenes[i]
        
        def show_image(j):
            prepare_image(scene['images'][j], models, p_scene)
            img = load_image(scene['images'][j]['path'])
            print(scene['images'][j]['scene_info'])
            rgba, _ = p_renderer.render(p_scene, flags=pyrender.RenderFlags.RGBA) 
            overlay = Image.fromarray(rgba)
            composite = Image.alpha_composite(img, overlay)
            blended = Image.blend(img, composite, alpha=0.4)
            blended.thumbnail((600, 600))
            display(blended)
            
        interact(show_image, j=range(len(scene['images'])))
    
    interact(show_scene, i=range(len(scenes)))

interact(show_camera, camera=cameras);<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15413070283600494, 'fraction_numerical': 0.02404438964241677, 'mean_word_length': 3.5768688293370947, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '444847', 'license': 'CC0-1.0', 'n_tokens_mistral': 1272, 'n_tokens_neox': 1227, 'n_words': 210, 'path': 'scripts/dimo_viewer.ipynb', 'repo_name': 'pderoovere/dimo', 'size': '3244'}