Example #1
0
def extract_skeleton(scene,
                     frame_ids=None,
                     frame_multiplier=1.,
                     time_multiplier=1.):
    joints = {
        ob.name.split('.')[1]: ob
        for ob in bpy.data.objects
        if ob.name.startswith('Output') and ob.name.endswith('Sphere')
    }
    print("joints: %s" % joints)
    assert len(joints) == Joint.get_num_joints(), \
        "No: %s != %s" % (len(joints), Joint.get_num_joints())
    if not frame_ids:
        frame_ids = range(scene.frame_start, scene.frame_end + 1)

    skeleton = Skeleton()
    for frame_id in frame_ids:
        o_frame_id = int(round(frame_id * frame_multiplier))
        if skeleton.has_pose(o_frame_id):
            print("skipping %s, because already set" % frame_id)
            continue
        print("frame_id: %s, o_frame_id: %s" % (frame_id, o_frame_id))
        scene.frame_set(frame_id)
        bpy.context.scene.update()
        pose = np.zeros(shape=(Skeleton.DIM, len(joints)))
        for joint, ob in joints.items():
            pos = ob.matrix_world.col[3]
            joint_id = Joint.from_string(joint)
            pose[:, joint_id] = from_blender(pos)
        assert not skeleton.has_pose(o_frame_id), "Already has %s" % frame_id
        skeleton.set_pose(frame_id=o_frame_id,
                          pose=pose,
                          time=o_frame_id * time_multiplier)

    # scenelet = Scenelet(skeleton=skeleton)
    # scenelet.save(path=path_out, save_obj=False)

    return skeleton
Example #2
0
def extract_annotated_scenelet(
        scene,
        prefix_obj='obb',
        frame_ids=None,
        frame_multiplier=1.,
        time_multiplier=1.,
        f_ob_is_joint=lambda ob: ob.name.startswith(
            'Output') and ob.name.endswith('Sphere'),
        f_joint_name_from_ob=lambda ob: ob.name.split('.')[1]):
    """
    
    Args:
        scene (bpy.types.Scene):
            The current scene (e.g. bpy.context.scene).
        prefix_obj (str):
            Start of object names that we want to include in the scenelet
            as oriented bounding boxes.
        frame_ids (List[int]):
            A subset of frame IDs to export.
        frame_multiplier (float):
            Scaling for frame IDs. The result will be rounded and truncated.
            output.frame_id := int(round(frame_id * frame_multiplier))
        time_multipler (float):
            Scaling for times associated with frame_ids.
            output.time := int(round(frame_id * frame_multiplier)) 
            * time_multiplier.
        f_ob_is_joint (Callable[[bpy.types.Object], bool]]):
            Decides if a Blender object is a joint.
        f_joint_name_from_ob (Callable[[bpy.types.Object], str]):
            Gets the joint name from the Blender object name.
    """
    # joints = {
    #     ob.name.split('.')[1]: ob
    #     for ob in bpy.data.objects
    #     if ob.name.startswith('Output') and ob.name.endswith('Sphere')}
    joints = {
        f_joint_name_from_ob(ob): ob
        for ob in bpy.data.objects if f_ob_is_joint(ob)
    }
    print("joints: %s" % joints)
    skeleton = Skeleton()
    if len(joints):
        assert len(joints) == 16, "No: %s" % len(joints)
        if not frame_ids:
            frame_ids = range(scene.frame_start, scene.frame_end + 1)
        for frame_id in frame_ids:
            o_frame_id = int(round(frame_id * frame_multiplier))
            if skeleton.has_pose(o_frame_id):
                print("skipping %s" % frame_id)
                continue
            print("frame_id: %s" % frame_id)
            scene.frame_set(frame_id)
            bpy.context.scene.update()
            # bpy.ops.anim.change_frame(frame_id)
            pose = np.zeros(shape=(3, len(joints)))
            for joint, ob in joints.items():
                pos = ob.matrix_world.col[3]
                print("pos[%s]: %s" % (ob.name, pos))
                joint_id = Joint.from_string(joint)
                print("joint %s is %s" % (joint, Joint(joint_id)))
                pose[:, joint_id] = from_blender(pos)
            print("o_frame: %s from %s" % (o_frame_id, frame_id))
            assert not skeleton.has_pose(o_frame_id), \
                "Already has %s" % frame_id
            skeleton.set_pose(frame_id=o_frame_id,
                              pose=pose,
                              time=o_frame_id * time_multiplier)
    objs_bl = {}
    for obj in bpy.data.objects:
        if obj.name.startswith(prefix_obj) and not obj.hide:
            obj_id = int(obj.name.split('_')[1])
            try:
                objs_bl[obj_id].append(obj)
            except KeyError:
                objs_bl[obj_id] = [obj]

    print("objs: %s" % objs_bl)
    scenelet = Scenelet(skeleton=skeleton)
    print("scenelet: %s" % scenelet)
    for obj_id, parts_bl in objs_bl.items():
        name_category = None
        scene_obj = None
        for part_id, part_bl in enumerate(parts_bl):
            transl, rot, scale = part_bl.matrix_world.decompose()
            rot = rot.to_matrix()
            if any(comp < 0. for comp in scale):
                scale *= -1.
                rot *= -1.
            assert not any(comp < 0. for comp in scale), "No: %s" % scale

            matrix_world = part_bl.matrix_world.copy()

            # need to save full scale, not only half axes
            for c in range(3):
                for r in range(3):
                    matrix_world[r][c] *= 2.
            name_parts = part_bl.name.split('_')
            if name_category is None:
                name_category = name_parts[2]
                scene_obj = SceneObj(label=name_category)
            else:
                assert name_category == name_parts[2], \
                    "No: %s %s" % (name_category, name_parts[2])
            name_part = name_parts[3]
            print("part: %s" % name_part)
            part = SceneObjPart(name_part)
            part.obb = Obb(centroid=np.array(
                from_blender([transl[0], transl[1], transl[2]])),
                           axes=np.array([[rot[0][0], rot[0][1], rot[0][2]],
                                          [-rot[2][0], -rot[2][1], -rot[2][2]],
                                          [rot[1][0], rot[1][1], rot[1][2]]]),
                           scales=np.array(
                               [scale[0] * 2., scale[1] * 2., scale[2] * 2.]))
            # if 'table' in name_category:
            #     print(part.obb.axes)
            #     raise RuntimeError("stop")
            print("obb: %s" % part.obb.to_json(0))
            scene_obj.add_part(part_id, part)
        scenelet.add_object(obj_id, scene_obj, clone=False)
    return scenelet