Function viz_epoch

mGPT/render/visualize.py:352–616

Generate & viz samples

(model,
              dataset,
              epoch,
              params,
              folder,
              module=None,
              writer=None,
              exps='')

def viz_epoch(model,
              dataset,
              epoch,
              params,
              folder,
              module=None,
              writer=None,
              exps=''):
    """ Generate & viz samples """
    module = model if module is None else module

    # visualize with joints3D
    model.outputxyz = True

    print(f"Visualization of the epoch {epoch}")

    noise_same_action = params["noise_same_action"]
    noise_diff_action = params["noise_diff_action"]
    duration_mode = params["duration_mode"]
    reconstruction_mode = params["reconstruction_mode"]
    decoder_test = params["decoder_test"]

    fact = params["fact_latent"]
    figname = params["figname"].format(epoch)

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

    num_classes = params["num_classes"]
    # nats = min(num_classes, nats)

    # define some classes
    classes = torch.randperm(num_classes)[:nats]
    # duplicate same classes when sampling too much
    if nats > num_classes:
        classes = classes.expand(nats)

    meandurations = torch.from_numpy(
        np.array([
            round(dataset.get_mean_length_label(cl.item())) for cl in classes
        ]))

    if duration_mode == "interpolate" or decoder_test == "diffduration":
        points, step = np.linspace(-nspa, nspa, nspa, retstep=True)
        # points = np.round(10*points/step).astype(int)
        points = np.array([5, 10, 16, 30, 60, 80]).astype(int)
        # gendurations = meandurations.repeat((nspa, 1)) + points[:, None]
        gendurations = torch.from_numpy(points[:, None]).expand(
            (nspa, 1)).repeat((1, nats))
    else:
        gendurations = meandurations.repeat((nspa, 1))
    print("Duration time: ")
    print(gendurations[:, 0])

    # extract the real samples
    # real_samples, real_theta, mask_real, real_lengths, imgs, paths
    batch = dataset.get_label_sample_batch(classes.numpy())