Questions about visualizing predicted image.

[w4nanch1]

Hi,

Thank you for sharing your source code publicly. I appreciate your contribution. However, I encountered an issue while trying to reproduce the future frame prediction results as described in the paper. The predicted images do not seem to perform well, as illustrated as following.

To visualize the predictions, I added the following code in the rollout function of eval_utils_calvin.py. Some parts of the code are omitted for brevity:

for step in range(EP_LEN):
    action, img_pred = model.step(obs, lang_annotation, step)
    os.makedirs('./saved_img/', exist_ok=True)
    if step % 50 == 0:
        img1 = unpatchify(img_pred[:, 0, ...], patch_size=args.patch_size, img_size=(args.calvin_input_image_size, args.calvin_input_image_size))
        img2 = unpatchify(img_pred[:, 1, ...], patch_size=args.patch_size, img_size=(args.calvin_input_image_size, args.calvin_input_image_size))
        for i in range(img1.shape[0]):
            torchvision.utils.save_image(img1[i, ...], f"./saved_img/img_{step}_t{i}_1.png")
            torchvision.utils.save_image(img2[i, ...], f"./saved_img/img_{step}_t{i}_2.png")
    if len(planned_actions) == 0: ...

Additionally, here is the unpatchify function I used:

def unpatchify(x, patch_size, img_size):
    N, L, _ = x.shape
    H, W = img_size
    h = H // patch_size
    w = W // patch_size

    x = x.view(N, h, w, 3, patch_size, patch_size)
    x = x.permute(0, 3, 1, 4, 2, 5).contiguous()
    x = x.view(N, 3, H, W)
    return x 

I used the pre-trained weights located in the folder finetune_bs=640_lr1e-4_atten_goal_state4_atten_only_obs_sv10_abc_reset_act_obs_ep5_abc, specifically 19.pth.

Here is a part of my eval.sh script:

calvin_dataset_path="calvin/dataset/task_ABC_D"
calvin_conf_path="calvin/calvin_models/conf"
vit_checkpoint_path="checkpoints/vit_mae/mae_pretrain_vit_base.pth" # downloaded from https://drive.google.com/file/d/1bSsvRI4mDM3Gg51C6xO0l9CbojYw3OEt/view?usp=sharing
save_checkpoint_path="checkpoints/"
### NEED TO CHANGE the checkpoint path ###
resume_from_checkpoint="checkpoints/calvin/19.pth"

Could you help me understand why the predicted images are not performing well? Am I missing something in the evaluation setup?
Thank you again for your time and assistance!

[yangsizhe]

Hi, thank you for your interest in our work.

The reason for the degraded image quality is that we applied the same normalization as described in [1] to improve the performance of the main task (we clarified this matter in the paper). For manipulation tasks, the quality of the predicted images is not critical. What matters is that the images adequately capture environmental changes and provide sufficient guidance for action prediction.

If you aim to achieve the performance described in the paper, you should remove the normalization for the image labels during training.

[1] He, Kaiming, et al. "Masked autoencoders are scalable vision learners." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[w4nanch1]

That's make sense. Thanks a lot!

[wxy301]

Hello, I am also trying to predict image visualizations and saw that you have successfully implemented this part before. Could you please share your modified eval_utils_calvin.py file? I would really appreciate it. Thank you very much!

[w4nanch1]

Hi,

I can't find the first version of my code cause there are lots of modification on my code :( But I can share that you have to delete the normalize_patchfied_image function in the train_utils.py file because as the author mentioned, The reason for the degraded image quality is that we applied the same normalization .

Then when you are testing the model, you can decode the image in this way (I'm just using the very simple method to visualize):

        reconstructed = unpatchify_multi_view(img_pred)
        save_path = './saved_img/' + args.run_name
        os.makedirs(save_path, exist_ok=True) 
        if step % 10 == 0:
            view1 = reconstructed[0, 0] 
            view2 = reconstructed[0, 1]

            view1_np = view1.permute(1, 2, 0).cpu().numpy() 
            view2_np = view2.permute(1, 2, 0).cpu().numpy()

            view1_np = (view1_np - view1_np.min()) / (view1_np.max() - view1_np.min()) * 255
            view2_np = (view2_np - view2_np.min()) / (view2_np.max() - view2_np.min()) * 255

            view1_np = view1_np.astype(np.uint8)
            view2_np = view2_np.astype(np.uint8)

            cv2.imwrite(os.path.join(save_path, f"img_{step}.png"), cv2.cvtColor(view1_np, cv2.COLOR_RGB2BGR))
            cv2.imwrite(os.path.join(save_path, f"wrist_{step}.png"), cv2.cvtColor(view2_np, cv2.COLOR_RGB2BGR))

The unpatchify_multi_view function is as shown below:

def unpatchify_multi_view(patched, patch_size=16, img_size=224):
    """
    Unpatchify tensor from (seq, views, num_patches, patch_dim) back to (1, seq, views, channels, H, W)
    """
    S, V, N, D = patched.shape  # (10, 2, 196, 768)
    H = W = img_size
    P = patch_size
    C = D // (P * P)
    patched = patched.reshape(S, V, int(H / P), int(W / P), P, P, C)
    patched = patched.permute(0, 1, 6, 2, 4, 3, 5).reshape(S, V, C, H, W)
    return patched.unsqueeze(0)