Function viz_dataset

mGPT/render/visualize.py:619–684

Generate & viz samples

(dataset, params, folder)

def viz_dataset(dataset, params, folder):
    """ Generate & viz samples """
    print("Visualization of the dataset")

    nspa = params["num_samples_per_action"]
    nats = params["num_actions_to_sample"]

    num_classes = params["num_classes"]

    figname = "{}_{}_numframes_{}_sampling_{}_step_{}".format(
        params["dataset"], params["pose_rep"], params["num_frames"],
        params["sampling"], params["sampling_step"])

    # define some classes
    classes = torch.randperm(num_classes)[:nats]

    allclasses = classes.repeat(nspa, 1).reshape(nspa * nats)
    # extract the real samples
    real_samples, mask_real, real_lengths = dataset.get_label_sample_batch(
        allclasses.numpy())
    # to visualize directly

    # Visualizaion of real samples
    visualization = {
        "x": real_samples,
        "y": allclasses,
        "mask": mask_real,
        'lengths': real_lengths,
        "output": real_samples
    }

    from mGPT.models.rotation2xyz import Rotation2xyz

    device = params["device"]
    rot2xyz = Rotation2xyz(device=device)

    rot2xyz_params = {
        "pose_rep": params["pose_rep"],
        "glob_rot": params["glob_rot"],
        "glob": params["glob"],
        "jointstype": params["jointstype"],
        "translation": params["translation"]
    }

    output = visualization["output"]
    visualization["output_xyz"] = rot2xyz(output.to(device),
                                          visualization["mask"].to(device),
                                          **rot2xyz_params)

    for key, val in visualization.items():
        if len(visualization[key].shape) == 1:
            visualization[key] = val.reshape(nspa, nats)
        else:
            visualization[key] = val.reshape(nspa, nats, *val.shape[1:])

    finalpath = os.path.join(folder, figname + ".gif")
    tmp_path = os.path.join(folder, f"subfigures_{figname}")
    os.makedirs(tmp_path, exist_ok=True)