src/task/evaluation/table1_unified_sliding.py

# -*- coding: utf-8 -*-
"""
Unified evaluation script supporting VAE and RVQ models with sliding window strategy.
"""
import os, argparse, tqdm, torch, gc
import numpy as np
from sata.mypath import *
from sata.mydataset import PairedDataset, get_mi_src_tgt_all_graph, PairedGraph_collate_fn
from sata.skel_pose_graph import SkelPoseGraph
from sata.mymodel import out_post_fwd

from sata.metric import compute_metric
from sata.conversions.graph_to_motion import hatD_recon_motion, gt_recon_motion
from fairmotion.data import bvh


def load_model_by_type(model_type, model_epoch, device):
    """
    Dynamically load the corresponding model based on model type.
    
    Args:
        model_type: "vae" or "rvq"
        model_epoch: model checkpoint path
        device: device
    
    Returns:
        model, cfg, ms_dict
    """
    if model_type == "vae":
        from sata.test import prepare_model_test
        print(f"[Model] Loading VAE model from: {model_epoch}")
    elif model_type == "rvq":
        from sata.test_vq import prepare_model_test
        print(f"[Model] Loading RVQ model from: {model_epoch}")
    else:
        raise ValueError(f"Unknown model_type: {model_type}. Must be 'vae' or 'rvq'")
    
    model, cfg, ms_dict = prepare_model_test(model_epoch, device)
    return model, cfg, ms_dict


def extract_window_from_dataset(dataset, mi, src_ri, tgt_ri, start_frame, window_size, total_frames):
    """
    Extract a window from the dataset directly from its internal structure (faster than loading frame by frame).
    
    Args:
        dataset: PairedDataset instance
        mi: motion index
        src_ri: source rig index
        tgt_ri: target rig index
        start_frame: window start frame
        window_size: window size
        total_frames: total number of frames
    
    Returns:
        (src_window, tgt_window), actual_window_size
    """
    end_frame = min(start_frame + window_size, total_frames)
    actual_window_size = end_frame - start_frame
    
    # get file indices
    src_fi = dataset.mi_ri_2_fi[mi][src_ri]
    tgt_fi = dataset.mi_ri_2_fi[mi][tgt_ri]
    
    # get skeleton data (shared across all frames)
    src_skel = dataset.skel_list[src_fi]
    tgt_skel = dataset.skel_list[tgt_fi]
    
    # get pose data for window frames (from pose_list)
    src_start_idx = dataset.start_frames[src_fi] + start_frame
    tgt_start_idx = dataset.start_frames[tgt_fi] + start_frame
    
    src_pose_window = dataset.pose_list[src_start_idx:src_start_idx + actual_window_size]
    tgt_pose_window = dataset.pose_list[tgt_start_idx:tgt_start_idx + actual_window_size]
    
    # build graph list for window
    src_graphs = [SkelPoseGraph(src_skel, pose) for pose in src_pose_window]
    tgt_graphs = [SkelPoseGraph(tgt_skel, pose) for pose in tgt_pose_window]
    
    # collate into batch
    window_data = list(zip(src_graphs, tgt_graphs))
    src_window, tgt_window = PairedGraph_collate_fn(window_data)
    
    return (src_window, tgt_window), actual_window_size


def process_with_sliding_window(model, model_type, ds, mi, src_ri, tgt_ri, total_frames, 
                                window_size, overlap, device, debug=False):
    """
    Process a motion with sliding windows, supporting overlap.
    
    Strategy: in overlapping regions, keep only the prediction from the earlier window.
    - First window: keep all frames
    - Subsequent windows: discard the first `overlap` frames (the overlapping part), keep only the non-overlapping part
    
    Args:
        model: trained model
        model_type: "vae" or "rvq"
        ds: PairedDataset instance
        mi: motion index
        src_ri: source rig index
        tgt_ri: target rig index
        total_frames: total number of frames
        window_size: window size
        overlap: overlap size between consecutive windows
        device: device
        debug: print debug information
    
    Returns:
        z_full: concatenated z across all windows (overlaps discarded)
        hatD_full: concatenated hatD across all windows (overlaps discarded)
        tgt_batch_full: full target batch for GT computation
    """
    stride = window_size - overlap
    
    # compute number of windows
    if total_frames <= window_size:
        num_windows = 1
    else:
        num_windows = (total_frames - window_size + stride - 1) // stride + 1
    
    if debug:
        print(f"\n[DEBUG] Processing motion {mi}:")
        print(f"  Total frames: {total_frames}")
        print(f"  Window size: {window_size}, Overlap: {overlap}, Stride: {stride}")
        print(f"  Number of windows: {num_windows}")
        print(f"  Model type: {model_type}")
        print(f"  Strategy: Discard overlap from later windows, keep only from first window")
    
    # determine number of nodes per frame from dataset structure
    # important: use tgt_skel node count since hatD is output for the target skeleton
    tgt_fi = ds.mi_ri_2_fi[mi][tgt_ri]
    tgt_skel = ds.skel_list[tgt_fi]
    num_nodes_per_frame = tgt_skel.lo.shape[0]  # number of joints in target skeleton
    
    if debug:
        print(f"  Detected {num_nodes_per_frame} nodes per frame")
    
    # collect non-overlapping parts from each window
    z_parts = []
    hatD_parts = []
    
    # process each window
    for window_idx in range(num_windows):
        start_frame = window_idx * stride
        
        # skip if start_frame exceeds total_frames
        if start_frame >= total_frames:
            if debug:
                print(f"  Window {window_idx}: start_frame {start_frame} >= {total_frames}, skipping")
            break
            
        # extract window from dataset
        (src_window, tgt_window), window_frames = extract_window_from_dataset(
            ds, mi, src_ri, tgt_ri, start_frame, window_size, total_frames
        )
        
        # move to device
        src_window = src_window.to(device)
        tgt_window = tgt_window.to(device)
        
        # call model according to type (handles different return value counts)
        with torch.no_grad():
            if model_type == "vae":
                # VAE model returns 2 values
                z_win, hatD_win = model(src_window, tgt_window, window_frames)
            elif model_type == "rvq":
                # RVQ model returns 4 values
                z_win, hatD_win, _, _ = model(src_window, tgt_window, window_frames)
            else:
                raise ValueError(f"Unknown model_type: {model_type}")
        
        # hatD shape is [T*N, D] where T=frames, N=nodes_per_frame
        hatD_dim = hatD_win.shape[1]
        
        # reshape hatD from [T*N, D] to [T, N, D]
        hatD_win_reshaped = hatD_win.view(window_frames, num_nodes_per_frame, hatD_dim)
        
        # determine which frames to keep from this window
        if window_idx == 0:
            # first window: keep all frames
            keep_start_idx = 0
            keep_end_idx = window_frames
            global_start = start_frame
            global_end = start_frame + window_frames
        else:
            # subsequent windows: discard overlapping part, keep only non-overlapping part
            keep_start_idx = overlap
            keep_end_idx = window_frames
            global_start = start_frame + overlap
            global_end = start_frame + window_frames
        
        # extract non-overlapping part
        z_keep = z_win[keep_start_idx:keep_end_idx]
        hatD_keep = hatD_win_reshaped[keep_start_idx:keep_end_idx]
        
        if debug:
            print(f"  Window {window_idx}: frames [{start_frame}, {start_frame + window_frames})")
            print(f"    Window internal frames: {window_frames}")
            print(f"    Keep indices: [{keep_start_idx}, {keep_end_idx}) -> global frames [{global_start}, {global_end})")
            print(f"    z_keep shape: {z_keep.shape}, hatD_keep shape: {hatD_keep.shape}")
        
        # append to lists
        z_parts.append(z_keep)
        hatD_parts.append(hatD_keep)
        
        # cleanup
        del src_window, tgt_window, z_win, hatD_win, hatD_win_reshaped
    
    # concatenate all parts
    z_full = torch.cat(z_parts, dim=0)  # [T, D]
    hatD_full_3d = torch.cat(hatD_parts, dim=0)  # [T, N, D]
    
    # reshape hatD back to [T*N, D] format (expected by parse_hatD)
    actual_frames = z_full.shape[0]
    hatD_full = hatD_full_3d.view(actual_frames * num_nodes_per_frame, -1)
    
    # verify frame count is correct
    if actual_frames != total_frames:
        if debug:
            print(f"  WARNING: Expected {total_frames} frames, got {actual_frames} frames")
    
    # check for NaN or Inf
    if torch.isnan(z_full).any() or torch.isinf(z_full).any():
        raise ValueError("NaN or Inf detected in z_full")
    if torch.isnan(hatD_full).any() or torch.isinf(hatD_full).any():
        raise ValueError("NaN or Inf detected in hatD_full")
    
    if debug:
        print(f"  Final shapes: z_full {z_full.shape}, hatD_full {hatD_full.shape}")
        print(f"  Total frames collected: {actual_frames} (expected: {total_frames})")
    
    # load the full target batch for GT computation
    (_, tgt_batch_full), _ = get_mi_src_tgt_all_graph(ds, mi, src_ri, tgt_ri, device)
    tgt_batch_full = tgt_batch_full.to(device)
    
    return z_full, hatD_full, tgt_batch_full


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_type", type=str, required=True, choices=["vae", "rvq"],
                        help="Model type: vae or rvq (REQUIRED)")
    parser.add_argument("--model_epoch", type=str, default="ckpt0")
    parser.add_argument("--device", type=str, default="cuda:0")
    parser.add_argument("--data_dir", type=str, default="test/motion/processed")
    parser.add_argument("--num_samples", type=int, default=-1,
                        help="Number of samples to evaluate (-1 for all)")
    parser.add_argument("--window_size", type=int, default=64,
                        help="Sliding window size (number of frames)")
    parser.add_argument("--overlap", type=int, default=16,
                        help="Overlap size between consecutive windows")
    parser.add_argument("--visualize", action="store_true",
                        help="Enable BVH visualization (only when num_samples < 3)")
    args = parser.parse_args()
    
    print("="*80)
    print("Unified Table1 Evaluation (with sliding window)")
    print("="*80)
    print(f"Model type: {args.model_type.upper()}")
    print(f"Model: {args.model_epoch}")
    print(f"Device: {args.device}")
    print(f"Data dir: {args.data_dir}")
    print(f"Window size: {args.window_size}")
    print(f"Overlap: {args.overlap}")
    print(f"Stride: {args.window_size - args.overlap}")
    print(f"Visualization: {args.visualize}")
    print("="*80)

    # load model by type
    model, cfg, ms_dict = load_model_by_type(args.model_type, args.model_epoch, args.device)

    # dataset
    ds = PairedDataset(min_motion_lens=8)
    data_dir = os.path.join(DATA_DIR, args.data_dir)
    ds.load_data_dir_pairs(data_dir)
    
    # determine test range based on num_samples
    total_samples = len(ds.mi_ri_2_fi)
    if args.num_samples > 0 and args.num_samples < total_samples:
        N = args.num_samples
    else:
        N = total_samples
    
    print(f"\nEvaluating {N} samples out of {total_samples} total samples\n")

    metric_key = ["qR", "ra_xz", "pa", "slide", "pen"]
    metric = {key: 0.0 for key in metric_key}
    
    model.eval()  # set model to evaluation mode
    
    # create output directory if visualization is enabled
    if args.visualize and N <= 3:
        vis_dir = os.path.join(os.path.dirname(__file__), "visual_result", f"table1_unified_{args.model_type}")
        os.makedirs(vis_dir, exist_ok=True)
        print(f"\nVisualization output directory: {vis_dir}\n")
    
    for mi in tqdm.tqdm(range(N)):
        # get total frame count for this motion
        fi = ds.mi_ri_2_fi[mi][0]  # use first rig to get frame count
        total_frames = ds.frame_cnts[fi]
        
        # enable debug for first sample or when visualizing
        debug_mode = (mi == 0) or (args.visualize and mi < 3)
        z, hatD, tgt_batch = process_with_sliding_window(
            model, args.model_type, ds, mi, 0, 0, total_frames, 
            args.window_size, args.overlap, args.device,
            debug=debug_mode
        )
        consq_n = total_frames
        
        # ========== BVH Visualization ==========
        if args.visualize and N <= 3:
            print(f"\n[Visualization] Generating BVH for motion {mi} (frames: {total_frames})...")
            try:
                # load source batch for visualization
                (src_batch, _), _ = get_mi_src_tgt_all_graph(ds, mi, 0, 0, args.device)
                src_batch = src_batch.to(args.device)
                
                # reconstruct source motion
                src_motion_list, _ = gt_recon_motion(src_batch, consq_n)
                
                # reconstruct motion from prediction
                out_motion_list, _ = hatD_recon_motion(
                    hatD, tgt_batch, cfg["representation"]["out"], ms_dict, consq_n
                )
                
                # reconstruct GT motion for comparison
                gt_motion_list, _ = gt_recon_motion(tgt_batch, consq_n)
                
                # save BVH files
                mi_vis_dir = os.path.join(vis_dir, f"motion_{mi}")
                os.makedirs(mi_vis_dir, exist_ok=True)
                
                src_path = os.path.join(mi_vis_dir, "src_motion.bvh")
                pred_path = os.path.join(mi_vis_dir, "predicted.bvh")
                gt_path = os.path.join(mi_vis_dir, "ground_truth.bvh")
                
                bvh.save(src_motion_list[0], src_path, rot_order="XYZ")
                bvh.save(out_motion_list[0], pred_path, rot_order="XYZ")
                bvh.save(gt_motion_list[0], gt_path, rot_order="XYZ")
                
                print(f"  ✓ Source motion: {src_path} ({src_motion_list[0].num_frames()} frames)")
                print(f"  ✓ Predicted motion: {pred_path} ({out_motion_list[0].num_frames()} frames)")
                print(f"  ✓ Ground truth: {gt_path} ({gt_motion_list[0].num_frames()} frames)")
                
            except Exception as e:
                print(f"  ✗ Visualization failed: {e}")
                import traceback
                traceback.print_exc()
        
        # compute output and GT using out_post_fwd
        out, gt = out_post_fwd(
            {"hatD": hatD, "z": z},
            tgt_batch,
            ms_dict,
            cfg["representation"]["out"],
            consq_n,
        )
        
        # compute metrics
        mi_metric = compute_metric(metric_key, out, gt)
        for key in metric_key:
            metric[key] += mi_metric[key].detach().item()

        # del tgt_batch, z, hatD, out, gt
        # gc.collect()
        # torch.cuda.empty_cache()

    # print results
    print("\n" + "="*80)
    print("Results:")
    print("="*80)
    for key in metric_key:
        print(f"{key}: {metric[key]/N:.4f}", end="\t")
    print("\n" + "="*80)
    
    # save results to file
    metric_fp = os.path.join(RESULT_DIR, args.model_epoch.split("/")[0], f"table1_unified_{args.model_type}_sliding.txt")
    os.makedirs(os.path.dirname(metric_fp), exist_ok=True)
    
    if not os.path.exists(metric_fp):
        with open(metric_fp, "w") as f:
            line = "model_epoch,\tmodel_type,\twindow_size,\toverlap,\t"
            for key in metric_key:
                line += f"{key},\t"
            f.write(line + "\n")
    
    with open(metric_fp, "a") as f:
        line = f"{args.model_epoch},\t{args.model_type},\t{args.window_size},\t{args.overlap},\t"
        for key in metric_key:
            line += f"{metric[key]/N:.4f},\t"
        f.write(line + "\n")
    
    print(f"\nResults saved to: {metric_fp}")