examples/custom_vis_torch_sample.py

import dataclasses
import os
import cv2
import numpy as np
import tqdm
import matplotlib
import io
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import tyro

import logging
import torch
import safetensors.torch
from openpi_value.training import config as _config
from openpi_value.shared import download
import openpi_value.training.data_loader as _data
from openpi_value.models_pytorch.pi0_pytorch import PI0Pytorch
import openpi_value.models.tokenizer as _tokenizer
from types import SimpleNamespace
from openpi_value.shared import image_tools

all_pred = []
all_tgt  = []


def write_episode_video(
    episode_id: int,
    value_list: list[float],
    img_list: list[list[np.ndarray]],  # now list of image lists
    output_dir: str,
    fig_w: float = 12.0,
    fig_h: float = 4.0,
    dpi: int = 100,
    fps: int = 30,
    metric_only: bool = False,  # * Skip visualizations
):
    """Writes a video for a single episode visualizing predicted values and image frames."""
    if not value_list:
        return 
    
    n_frames = len(value_list)
    tgt_progress = np.linspace(0.0, 1.0, n_frames)
    val_pred     = np.array(value_list, dtype=np.float32)

    global all_pred
    global all_tgt

    all_pred.extend(val_pred.tolist())
    all_tgt.extend(tgt_progress.tolist())

    if metric_only:
        return

    os.makedirs(output_dir, exist_ok=True)
    out_path = os.path.join(output_dir, f"episode_{episode_id:03d}.mp4")

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")

    width_px = int(fig_w * dpi)
    height_px = int(fig_h * dpi)
    video_writer = cv2.VideoWriter(out_path, fourcc, fps, (width_px, height_px))

    logging.info(f"Writing video for episode {episode_id} with {n_frames} frames...")

    for idx in tqdm.tqdm(range(n_frames), desc=f"Episode {episode_id} Video"):
        fig, axes = plt.subplots(1, 1 + len(img_list[idx]), figsize=(fig_w, fig_h), dpi=dpi)

        ax_plot = axes[0]
        # Plotting
        x = np.arange(idx + 1)
        y = np.array(value_list[: idx + 1], dtype=np.float32)
        ax_plot.plot(x, y, linewidth=2, color="tab:blue")
        ax_plot.set_xlim(0, n_frames)
        ax_plot.set_ylim(0.0, 1.0)
        ax_plot.set_xlabel("Frame")
        ax_plot.set_ylabel("Predicted Value")
        ax_plot.set_title("Value Prediction Over Time")
        ax_plot.grid(True)

        # Images: e.g., base, wrist_left, wrist_right
        views = img_list[idx]
        titles = ["Base Frame", "Wrist Left", "Wrist Right"]
        
        titles = titles[:len(views)]  # Adjust titles to match number of views

        for ax_img, view, title in zip(axes[1:], views, titles):
            # ax_img.imshow(cv2.cvtColor(view, cv2.COLOR_BGR2RGB))
            ax_img.imshow(view)
            
            ax_img.set_title(title)
            ax_img.axis("off")

        plt.tight_layout()

        buf = io.BytesIO()
        fig.savefig(buf, format="png", dpi=dpi, bbox_inches="tight")
        buf.seek(0)
        png_bytes = np.frombuffer(buf.getvalue(), dtype=np.uint8)
        buf.close()
        plt.close(fig)

        img_bgr = cv2.imdecode(png_bytes, cv2.IMREAD_COLOR)
        h_b, w_b = img_bgr.shape[:2]
        if (w_b, h_b) != (width_px, height_px):
            img_bgr = cv2.resize(img_bgr, (width_px, height_px), interpolation=cv2.INTER_LINEAR)
        video_writer.write(img_bgr)

    video_writer.release()

    # Encode with H.264 for size/speed
    new_out_path = out_path.replace(".mp4", "_new.mp4")
    os.system(f"ffmpeg -y -i {out_path} -c:v libx264 -crf 18 -preset veryfast {new_out_path} > /dev/null 2>&1")
    logging.info(f"=> Episode {episode_id} generated to: {new_out_path}")
    os.remove(out_path)

def process_and_convert(img: np.ndarray, target_size: int) -> torch.Tensor:
    rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    tensor = torch.from_numpy(rgb_img).float() / 255.0
    tensor = tensor * 2.0 - 1.0  # Normalize to [-1, 1]

    tensor = image_tools.resize_with_pad_torch(tensor, 224, 224)  # * Keep ratio

    # HWC to CHW
    tensor = tensor.permute(2, 0, 1) 

    # Batch Dimension: Change C, H, W to 1, C, H, W (the final shape)
    tensor = tensor.unsqueeze(0)
    
    # Final shape check
    return tensor


def main(
    config_name: str,
    ckpt_dir: str,
    split: str = "val_tasks",
    metric_only: bool = False,
    output_video_dir: str = "./visualizations",
    # headview_only: bool = False,
    # n_batches: int = 20000,
    # fps_downsample: int = 1,
):
    """Main function to run value prediction and visualization."""
    
    global all_pred
    global all_tgt

    # --- Config and Checkpoint Setup ---
    config = _config.get_config(config_name)
    checkpoint_dir = download.maybe_download(ckpt_dir)


    # =============== Load Model =============================
    new_model = config.model.__class__(**{**config.model.__dict__,
                                            'p_mask_ego_state': 1,
                                        })
    
    config = dataclasses.replace(config, model=new_model)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = PI0Pytorch(new_model).to(device)
    
    model.sample_values = torch.compile(model.sample_values, mode="reduce-overhead")
    model.eval()  # Set model to evaluation mode
    model_path = os.path.join(checkpoint_dir, "model.safetensors")
    logging.info(f"Loading weights from: {model_path}")
        
    safetensors.torch.load_model(model, model_path, strict=True)

    logging.info(f"Loaded PyTorch weights successfully.")

    tokenizer = _tokenizer.PaligemmaTokenizer(new_model.max_token_len)
    # =============== Load Model finished ======================
    
    # --- Modify Config for Inference ---
    assert split in ['all', 'val_tasks', 'heldout_tasks']
    config = dataclasses.replace(
        config,
        batch_size=1, 
        is_train=False,
        num_workers=8, 
        split=split, 
        preceding_skipping_ratio=0,
    )
    # --- Data Loader ---

    # --- Helper for Image Processing ---
    def process_view_torch(torch_arr):
        """Converts a single image tensor [-1, 1] to a numpy array [0, 255]."""
        arr = torch_arr.cpu().numpy()
        img = ((arr + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
        # Assuming CHW -> HWC
        if img.shape[0] == 3:
            img = np.transpose(img, (1, 2, 0))
        return img

    value_frames = []  # type: list[float]
    raw_frames = []    # type: list[list[np.ndarray]]
    
    fps = 30

    # Setup output directory
    output_dir = os.path.join(output_video_dir, config_name, os.path.basename(checkpoint_dir) + f"_{split}")
    print(f">>>>>> Generating episode videos to: {output_dir} ...")
    os.makedirs(output_dir, exist_ok=True)



    # # * Load one video and draw visualization based on that
    image_folder_root = "path/to/your/images"  # * Change this to your image folder root
    left_folder = os.path.join(image_folder_root, 'hand_left')
    right_folder = os.path.join(image_folder_root, 'hand_right')
    top_folder = os.path.join(image_folder_root, 'top_head')

    left_images = sorted([os.path.join(left_folder, f) for f in os.listdir(left_folder) if f.endswith('.jpg')])
    right_images = sorted([os.path.join(right_folder, f) for f in os.listdir(right_folder) if f.endswith('.jpg')])
    top_images = sorted([os.path.join(top_folder, f) for f in os.listdir(top_folder) if f.endswith('.jpg')])

    n_frames = min(len(left_images), len(right_images), len(top_images))

    left_img_0 = cv2.imread(left_images[0])
    right_img_0 = cv2.imread(right_images[0])
    top_img_0 = cv2.imread(top_images[0])

    top_torch_0 = process_and_convert(top_img_0, 224)
    left_torch_0 = process_and_convert(left_img_0, 224)
    right_torch_0 = process_and_convert(right_img_0, 224)

    for i in range(1, n_frames):
        left_img = cv2.imread(left_images[i])
        right_img = cv2.imread(right_images[i])
        top_img = cv2.imread(top_images[i])
      
        top_torch = process_and_convert(top_img, 224)
        left_torch = process_and_convert(left_img, 224)
        right_torch = process_and_convert(right_img, 224)
        
        prompt = "Insert the memory stick."
        observation = {
            "state": torch.zeros((1, 32), dtype=torch.float32).to(device),
            "images": {
                "base_-100_rgb": top_torch_0.to(device),
                "left_wrist_-100_rgb": left_torch_0.to(device),
                "right_wrist_-100_rgb": right_torch_0.to(device),

                "base_0_rgb": top_torch.to(device),
                "left_wrist_0_rgb": left_torch.to(device),
                "right_wrist_0_rgb": right_torch.to(device),
            },
            "image_masks":{}  # * Set empty
        }

        #  * ['base_-100_rgb', 'left_wrist_-100_rgb', 'right_wrist_-100_rgb', 'base_0_rgb', 'left_wrist_0_rgb', 'right_wrist_0_rgb']


        tokens, token_masks = tokenizer.tokenize(prompt, state=None)
        tokens = tokens[np.newaxis, :]  # Add batch dim  
        token_masks = token_masks[np.newaxis, :]  # Add batch dim

        observation = {**observation, 
                       "tokenized_prompt": torch.from_numpy(tokens).to(device), 
                       "tokenized_prompt_mask": torch.from_numpy(token_masks).to(device)}

        observation = SimpleNamespace(**observation)

        with torch.no_grad():
            # Predict value
            val_arr = model.sample_values(device, observation)  # Shape=(1, 1)
            val = float(val_arr[0, 0].item())

            value_frames.append(val)

        base_torch = observation.images["base_0_rgb"][0]
        img_base = process_view_torch(base_torch)
        current_views = [img_base] # Start with base view
        raw_frames.append(current_views)

    logging.info(f"Writing final episode {0}.")
    write_episode_video(
        episode_id=0,
        value_list=value_frames,
        img_list=raw_frames,
        output_dir=output_dir,
        fig_w=12.0,
        fig_h=4.0,
        dpi=100,
        fps=fps,
        metric_only=metric_only,
    )





if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    tyro.cli(main)