pabloruizponce / in2IN

[CVPRW 2024] Official Implementation of "in2IN: Leveraging individual Information to Generate Human INteractions".
https://pabloruizponce.github.io/in2IN/
Other
41 stars 1 forks source link

About visualization #2

Open wangwenlonggg opened 3 months ago

wangwenlonggg commented 3 months ago

Could you provide the code for visualizing the SMPL meshes?

pabloruizponce commented 3 months ago

Hi, unfortunately, the SMPL visualization code won't be released soon. The code is very messy and it will be difficult to integrate in the actual codebase. However, almost all the code was borrowed from the MotionGPT. Additionally, I included this function in their codebase to render interaction sequences. Hope it helps

def render_interaction(npydata1,
                       npydata2,
                       frames_folder,
                       *,
                       mode,
                       model_path,
                       faces_path,
                       gt=False,
                       exact_frame=None,
                       num=8,
                       downsample=True,
                       canonicalize=True,
                       always_on_floor=False,
                       denoising=True,
                       oldrender=True,
                       res="high",
                       init=True,
                       accelerator='gpu',
                       device=[0]):
    if init:
        # Setup the scene (lights / render engine / resolution etc)
        setup_scene(res=res,
                    denoising=denoising,
                    oldrender=oldrender,
                    accelerator=accelerator,
                    device=device)

    is_mesh, is_smplx, jointstype = style_detect(npydata1)

    if not is_mesh:
        npydata1 = npydata1 * smplh_to_mmm_scaling_factor
        npydata2 = npydata2 * smplh_to_mmm_scaling_factor

    if is_smplx:
        smplx_model_male = smplx.create(model_path,
                                        model_type='smplx',
                                        gender='male',
                                        ext='npz',
                                        num_betas=10,
                                        flat_hand_mean=True,
                                        use_pca=False)
        faces_path = smplx_model_male.faces

    # Put everything in this folder
    if mode == "video":
        if always_on_floor:
            frames_folder += "_of"
        os.makedirs(frames_folder, exist_ok=True)
        # if it is a mesh, it is already downsampled
        if downsample and not is_mesh:
            npydata1 = npydata1[::8]
            npydata2 = npydata2[::8]
    elif mode == "sequence":
        img_name, ext = os.path.splitext(frames_folder)
        if always_on_floor:
            img_name += "_of"
        img_path = f"{img_name}{ext}"

    elif mode == "frame":
        img_name, ext = os.path.splitext(frames_folder)
        if always_on_floor:
            img_name += "_of"
        img_path = f"{img_name}_{exact_frame}{ext}"

    if mode == "sequence":
        perc = 0.2
        npydata1 = prune_begin_end(npydata1, perc)
        npydata2 = prune_begin_end(npydata2, perc)

    if is_mesh:
        from .meshes import Meshes
        data1 = Meshes(npydata1,
                      gt=gt,
                      mode=mode,
                      faces_path=faces_path,
                      canonicalize=canonicalize,
                      always_on_floor=always_on_floor,
                      is_smplx=is_smplx,
                      color="red")

        data2 = Meshes(npydata2,
                        gt=gt,
                        mode=mode,
                        faces_path=faces_path,
                        canonicalize=canonicalize,
                        always_on_floor=always_on_floor,
                        is_smplx=is_smplx,
                        color="green")

    else:
        from .joints import Joints
        data1 = Joints(npydata1,
                      gt=gt,
                      mode=mode,
                      canonicalize=canonicalize,
                      always_on_floor=always_on_floor,
                      jointstype=jointstype)

        data2 = Joints(npydata2,
                        gt=gt,
                        mode=mode,
                        canonicalize=canonicalize,
                        always_on_floor=always_on_floor,
                        jointstype=jointstype)

    # Number of frames possible to render
    nframes = len(data1)

    # Show the trajectory
    show_traj(data1.trajectory)

    # Initialize the camera
    initial_root = (data1.get_mean_root() + data2.get_mean_root()) / 2
    camera = Camera(first_root=initial_root, mode=mode, is_mesh=is_mesh)

    frameidx = get_frameidx(mode=mode,
                            nframes=nframes,
                            exact_frame=exact_frame,
                            frames_to_keep=num)

    nframes_to_render = len(frameidx)

    # Center the camera to the middle
    if mode == "sequence":
        mean_root = (data1.get_mean_root() + data2.get_mean_root()) / 2
        camera.update(mean_root)

    imported_obj_names = []
    for index, frameidx in enumerate(frameidx):
        if mode == "sequence":
            frac = index / (nframes_to_render - 1)
            mat1 = data1.get_sequence_mat1(frac)
            mat2 = data1.get_sequence_mat2(frac)
        else:
            mat1 = data1.mat
            mat2 = data2.mat
            new_root = (data1.get_root(frameidx) + data2.get_root(frameidx)) / 2

        islast = index == (nframes_to_render - 1)

        objname1 = data1.load_in_blender(frameidx, mat1)
        objname2 = data2.load_in_blender(frameidx, mat2)

        # Retrieve the objects
        obj1 = mathutils.Vector(data1.get_root(0))
        obj2 = mathutils.Vector(data2.get_root(0))

        # Calculate the midpoint between the two objects
        midpoint = (obj1 + obj2) / 2

        # Desired distance from the midpoint to the camera
        distance = 10.0

        # Determine the lateral direction
        lateral_direction = mathutils.Vector((0.0, 1.0, 0.0))

        # Calculate the camera position laterally using the lateral direction
        camera_position = midpoint + lateral_direction.normalized() * distance
        camera.camera.location = camera_position

        # Point the camera towards the midpoint
        # Calculate the direction vector from the camera to the midpoint
        direction = midpoint - camera.camera.location
        # Calculate the rotation quaternion to align the camera's -Z axis with the direction vector
        rot_quat = direction.to_track_quat('-Z', 'Y')
        # Apply the rotation
        camera.camera.rotation_euler = rot_quat.to_euler()

        name = f"{str(index).zfill(4)}"

        if mode == "video":
            path = os.path.join(frames_folder, f"frame_{name}.png")
        else:
            path = img_path

        if mode == "sequence":
            imported_obj_names.extend(objname1)
            imported_obj_names.extend(objname2)
        elif mode == "frame":
            camera.update(data1.get_root(frameidx))

        if mode != "sequence" or islast:
            render_current_frame(path)
            delete_objs(objname1)
            delete_objs(objname2)

    bpy.ops.wm.save_as_mainfile(filepath=frames_folder.replace('.png','.blend').replace('_frames','.blend'))

    # remove every object created
    delete_objs(imported_obj_names)
    delete_objs(["Plane", "myCurve", "Cylinder"])

    if mode == "video":
        return frames_folder
    else:
        return img_path