vis_collage.py

import numpy as np
import os
import imageio.v2 as imageio
import cv2

vid_id = 0
model_name = "v02"
base_folder = "vis_%s" % model_name
save_base_folder = "collage_%s" % model_name
os.makedirs(save_base_folder, exist_ok=True)

T = 40
seg_h, seg_w = 600, 600
rgb_h, rgb_w = 187, 400
assert(seg_w * 2 == rgb_w * 3)

for i in range(80):  # Assuming you have folders sample_vis_000 to sample_vis_079
    folder_name = os.path.join(base_folder, "sample_vis_%03d" % i)
    savefolder_name = os.path.join(save_base_folder, "sample_vis_%03d" % i)
    os.makedirs(savefolder_name, exist_ok=True)
    
    # Check if the collage already exists
    collage_file = os.path.join(savefolder_name, f"collage_000.png")
    if os.path.exists(collage_file):
        print(f"Collage for {folder_name} already exists, skipping...")
        continue
    
    # Initialize collage
    collage_t = np.zeros((seg_h + rgb_h * 2, seg_w * 2, 3))
    
    for t in range(T):
        print(f"Processing {folder_name} - Frame {t}")

        seg_g_file = os.path.join(folder_name, f"seg_gt_{t:03d}.png")
        seg_e_file = os.path.join(folder_name, f"seg_et_{t:03d}.png")
        cam_files = [
            os.path.join(folder_name, f"cam{i}_rgb_{t:03d}.png") for i in range(6)
        ]

        # Check if all required files exist for current frame
        if not (os.path.exists(seg_g_file) and os.path.exists(seg_e_file) and all(os.path.exists(cam_file) for cam_file in cam_files)):
            continue

        # Read and process images
        seg_g = imageio.imread(seg_g_file)
        seg_g = cv2.resize(seg_g.astype(np.uint8), (seg_w, seg_h))
        seg_g = np.flip(seg_g, 0)

        seg_e = imageio.imread(seg_e_file)
        seg_e = cv2.resize(seg_e.astype(np.uint8), (seg_w, seg_h))
        seg_e = np.flip(seg_e, 0)

        cam_images = [imageio.imread(cam_file) for cam_file in cam_files]
        cam_images_resized = [
            cv2.resize(cam_image.astype(np.uint8), (rgb_w, rgb_h)) for cam_image in cam_images
        ]

        # Create collage
        collage_t[:seg_h, :seg_w] = seg_e
        collage_t[:seg_h, seg_w:] = seg_g

        for j, cam_image in enumerate(cam_images_resized):
            row_index = seg_h + (j // 3) * rgb_h
            col_index = (j % 3) * rgb_w
            collage_t[row_index:row_index + rgb_h, col_index:col_index + rgb_w] = cam_image

        # Save collage
        collage_t_name = os.path.join(savefolder_name, f"collage_{t:03d}.png")
        imageio.imsave(collage_t_name, collage_t.astype(np.uint8))

print("Processing complete.")