update evaluation and fix psnr

README.md

@@ -139,6 +139,7 @@ python calculate_fid.py --videos_dir outputs/2025-11-30/00-02-42/videos/test_vis
 | FID    | 15.13  |
 
 *Note: FID is computed over 5000 frames.*
+
 *Note: Previous versions incorrectly used `data_range=2.0` for PSNR calculation, but the decoded video data is in the range [0, 1], so `data_range=1.0` should be used. This bug inflated PSNR values by approximately 6 dB. We have now corrected this in the latest version by clipping predictions to [0, 1] before metric computation. The relative performance comparisons and conclusions remain unchanged.*
 
 ---

calculate_fid.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python3
+"""
+Calculate FID (Fréchet Inception Distance) between predicted and ground truth videos.
+
+Usage:
+    python calculate_fid.py --videos_dir /path/to/videos
+    python calculate_fid.py --videos_dir /path/to/videos --batch_size 32
+"""
+
+import torch
+import numpy as np
+from pathlib import Path
+from tqdm import tqdm
+import argparse
+import cv2
+from torchmetrics.image.fid import FrechetInceptionDistance
+
+
+def load_video_frames(video_path, max_frames=None):
+    """
+    Load frames from a video file.
+    
+    Args:
+        video_path: Path to the video file
+        max_frames: Maximum number of frames to load (None = all frames)
+        
+    Returns:
+        torch.Tensor: Video frames with shape (T, C, H, W) in range [0, 255]
+    """
+    cap = cv2.VideoCapture(str(video_path))
+    frames = []
+    frame_count = 0
+    
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        # Convert BGR to RGB
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frames.append(frame)
+        frame_count += 1
+        
+        if max_frames and frame_count >= max_frames:
+            break
+    
+    cap.release()
+    
+    if len(frames) == 0:
+        raise ValueError(f"No frames loaded from {video_path}")
+    
+    # Convert to tensor: (T, H, W, C) -> (T, C, H, W)
+    frames = np.stack(frames, axis=0)
+    frames = torch.from_numpy(frames).permute(0, 3, 1, 2)
+    
+    return frames
+
+
+def load_videos_from_directory(video_dir, max_frames_per_video=None, max_videos=None):
+    """
+    Load all videos from a directory.
+    
+    Args:
+        video_dir: Directory containing .mp4 files
+        max_frames_per_video: Maximum frames to load per video
+        max_videos: Maximum number of videos to load
+        
+    Returns:
+        torch.Tensor: All frames concatenated with shape (N, C, H, W)
+    """
+    video_dir = Path(video_dir)
+    video_paths = sorted(list(video_dir.glob("**/*.mp4")))
+    
+    if max_videos:
+        video_paths = video_paths[:max_videos]
+    
+    all_frames = []
+    
+    print(f"Loading videos from {video_dir}")
+    print(f"Found {len(video_paths)} videos")
+    
+    for video_path in tqdm(video_paths, desc="Loading videos"):
+        try:
+            frames = load_video_frames(video_path, max_frames=max_frames_per_video)
+            all_frames.append(frames)
+        except Exception as e:
+            print(f"\nWarning: Failed to load {video_path.name}: {e}")
+            continue
+    
+    if len(all_frames) == 0:
+        raise ValueError(f"No videos loaded from {video_dir}")
+    
+    # Concatenate all frames: (N_videos, T, C, H, W) -> (N_total_frames, C, H, W)
+    all_frames = torch.cat(all_frames, dim=0)
+    
+    print(f"Loaded {all_frames.shape[0]} frames total")
+    print(f"Frame shape: {all_frames.shape[1:]}")
+    
+    return all_frames
+
+
+def calculate_fid(pred_dir, gt_dir, batch_size=32, device='cuda', 
+                  max_frames_per_video=None, max_videos=None):
+    """
+    Calculate FID between predicted and ground truth videos.
+    
+    Args:
+        pred_dir: Directory containing predicted videos
+        gt_dir: Directory containing ground truth videos
+        batch_size: Batch size for FID calculation
+        device: Device to use ('cuda' or 'cpu')
+        max_frames_per_video: Maximum frames to load per video
+        max_videos: Maximum number of videos to load from each directory
+        
+    Returns:
+        float: FID score
+    """
+    print("="*60)
+    print("FID Calculation")
+    print("="*60)
+    print(f"Pred directory: {pred_dir}")
+    print(f"GT directory: {gt_dir}")
+    print(f"Device: {device}")
+    print(f"Batch size: {batch_size}")
+    print("="*60 + "\n")
+    
+    # Check if directories exist
+    pred_dir = Path(pred_dir)
+    gt_dir = Path(gt_dir)
+    
+    if not pred_dir.exists():
+        raise ValueError(f"Pred directory does not exist: {pred_dir}")
+    if not gt_dir.exists():
+        raise ValueError(f"GT directory does not exist: {gt_dir}")
+    
+    # Load videos
+    print("\n[1/3] Loading predicted videos...")
+    pred_frames = load_videos_from_directory(
+        pred_dir, 
+        max_frames_per_video=max_frames_per_video,
+        max_videos=max_videos
+    )
+    
+    print("\n[2/3] Loading ground truth videos...")
+    gt_frames = load_videos_from_directory(
+        gt_dir,
+        max_frames_per_video=max_frames_per_video,
+        max_videos=max_videos
+    )
+    
+    # Initialize FID model
+    print("\n[3/3] Calculating FID...")
+    fid_model = FrechetInceptionDistance(normalize=True).to(device)
+    
+    # Process pred frames in batches
+    print("Processing predicted frames...")
+    num_pred_frames = pred_frames.shape[0]
+    for i in tqdm(range(0, num_pred_frames, batch_size)):
+        batch = pred_frames[i:i+batch_size]
+        batch = batch.to(device)
+        fid_model.update(batch, real=False)
+    
+    # Process gt frames in batches
+    print("Processing ground truth frames...")
+    num_gt_frames = gt_frames.shape[0]
+    for i in tqdm(range(0, num_gt_frames, batch_size)):
+        batch = gt_frames[i:i+batch_size]
+        batch = batch.to(device)
+        fid_model.update(batch, real=True)
+    
+    # Compute FID
+    fid_score = fid_model.compute().item()
+    
+    return fid_score
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Calculate FID between predicted and ground truth videos"
+    )
+    parser.add_argument(
+        "--videos_dir",
+        type=str,
+        default="/mnt/worldmem_valid/outputs/2025-12-01/08-09-46/videos/test_vis",
+        help="Base directory containing 'pred' and 'gt' subdirectories"
+    )
+    parser.add_argument(
+        "--pred_dir",
+        type=str,
+        default=None,
+        help="Override pred directory (default: {videos_dir}/pred)"
+    )
+    parser.add_argument(
+        "--gt_dir",
+        type=str,
+        default=None,
+        help="Override gt directory (default: {videos_dir}/gt)"
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=32,
+        help="Batch size for FID calculation (default: 32)"
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Device to use (default: cuda if available)"
+    )
+    parser.add_argument(
+        "--max_frames_per_video",
+        type=int,
+        default=None,
+        help="Maximum frames to load per video (default: None, load all)"
+    )
+    parser.add_argument(
+        "--max_videos",
+        type=int,
+        default=50,
+        help="Maximum number of videos to load (default: None, load all)"
+    )
+    
+    args = parser.parse_args()
+    
+    # Determine pred and gt directories
+    videos_dir = Path(args.videos_dir)
+    
+    if args.pred_dir:
+        pred_dir = Path(args.pred_dir)
+    else:
+        pred_dir = videos_dir / "pred"
+    
+    if args.gt_dir:
+        gt_dir = Path(args.gt_dir)
+    else:
+        gt_dir = videos_dir / "gt"
+    
+    # Calculate FID
+    try:
+        fid_score = calculate_fid(
+            pred_dir=pred_dir,
+            gt_dir=gt_dir,
+            batch_size=args.batch_size,
+            device=args.device,
+            max_frames_per_video=args.max_frames_per_video,
+            max_videos=args.max_videos
+        )
+        
+        # Print results
+        print("\n" + "="*60)
+        print("RESULTS")
+        print("="*60)
+        print(f"FID Score: {fid_score:.4f}")
+        print("="*60)
+        
+        # Save results to file
+        output_file = videos_dir / "fid_results.txt"
+        with open(output_file, 'w') as f:
+            f.write(f"FID Score: {fid_score:.4f}\n")
+            f.write(f"Pred directory: {pred_dir}\n")
+            f.write(f"GT directory: {gt_dir}\n")
+        
+        print(f"\nResults saved to: {output_file}")
+        
+    except Exception as e:
+        print(f"\n✗ Error: {e}")
+        import traceback
+        traceback.print_exc()
+        return 1
+    
+    return 0
+
+
+if __name__ == "__main__":
+    exit(main())
+

compute_video_psnr.py

@@ -1,195 +0,0 @@
-#!/usr/bin/env python3
-"""
-Compute PSNR for each video pair in gt and pred directories.
-"""
-import os
-import numpy as np
-from pathlib import Path
-from tqdm import tqdm
-import pandas as pd
-import imageio
-
-
-def load_video(video_path):
-    """
-    Load a video file and return frames as a numpy array.
-    
-    Args:
-        video_path: Path to the video file
-        
-    Returns:
-        np.ndarray: Video frames of shape (T, C, H, W) normalized to [0, 1]
-    """
-    reader = imageio.get_reader(str(video_path))
-    frames = []
-    for frame in reader:
-        # frame is already RGB
-        frame = frame.astype(np.float32) / 255.0
-        frames.append(frame)
-    reader.close()
-    
-    if len(frames) == 0:
-        raise ValueError(f"No frames loaded from {video_path}")
-    
-    # Convert to numpy array: (T, H, W, C) -> (T, C, H, W)
-    frames = np.stack(frames, axis=0)
-    frames = np.transpose(frames, (0, 3, 1, 2))
-    
-    return frames
-
-
-def calculate_psnr(pred, gt, data_range=1.0):
-    """
-    Calculate PSNR between two images/videos.
-    
-    Args:
-        pred: Predicted frames
-        gt: Ground truth frames
-        data_range: Data range (default: 1.0)
-        
-    Returns:
-        float: PSNR value in dB
-    """
-    mse = np.mean((pred - gt) ** 2)
-    if mse == 0:
-        return float('inf')
-    return 10 * np.log10((data_range ** 2) / mse)
-
-
-def compute_psnr(pred_frames, gt_frames, data_range=1.0):
-    """
-    Compute PSNR between predicted and ground truth frames.
-    
-    Args:
-        pred_frames: Predicted frames array (T, C, H, W)
-        gt_frames: Ground truth frames array (T, C, H, W)
-        data_range: Data range of the frames (default: 1.0 for [0, 1] range)
-        
-    Returns:
-        dict: Dictionary containing overall PSNR and frame-wise PSNR
-    """
-    # Ensure same shape
-    assert pred_frames.shape == gt_frames.shape, \
-        f"Shape mismatch: pred {pred_frames.shape} vs gt {gt_frames.shape}"
-    
-    T, C, H, W = pred_frames.shape
-    
-    # Compute frame-wise PSNR
-    frame_wise_psnr = []
-    for t in range(T):
-        psnr = calculate_psnr(pred_frames[t], gt_frames[t], data_range=data_range)
-        frame_wise_psnr.append(psnr)
-    
-    # Compute overall PSNR
-    overall_psnr = calculate_psnr(pred_frames, gt_frames, data_range=data_range)
-    
-    return {
-        'overall_psnr': overall_psnr,
-        'frame_wise_psnr': frame_wise_psnr,
-        'mean_frame_psnr': np.mean(frame_wise_psnr),
-        'std_frame_psnr': np.std(frame_wise_psnr),
-        'num_frames': T
-    }
-
-
-def main():
-    # Directories
-    gt_dir = Path("/mnt/WorldMem/outputs/2025-12-01/23-39-44/videos/test_vis/gt")
-    pred_dir = Path("/mnt/WorldMem/outputs/2025-12-01/23-39-44/videos/test_vis/pred")
-    
-    # Get all video files in gt directory
-    gt_videos = sorted(gt_dir.glob("*.mp4"))
-    
-    print(f"Found {len(gt_videos)} videos to process")
-    print("=" * 80)
-    
-    results = []
-    
-    for gt_path in tqdm(gt_videos, desc="Computing PSNR"):
-        video_name = gt_path.name
-        pred_path = pred_dir / video_name
-        
-        if not pred_path.exists():
-            print(f"Warning: Prediction video not found for {video_name}")
-            continue
-        
-        try:
-            # Load videos
-            gt_frames = load_video(gt_path)
-            pred_frames = load_video(pred_path)
-            
-            # Compute PSNR
-            psnr_results = compute_psnr(pred_frames, gt_frames, data_range=1.0)
-            
-            # Store results
-            result = {
-                'video_name': video_name,
-                'overall_psnr': psnr_results['overall_psnr'],
-                'mean_frame_psnr': psnr_results['mean_frame_psnr'],
-                'std_frame_psnr': psnr_results['std_frame_psnr'],
-                'num_frames': psnr_results['num_frames']
-            }
-            results.append(result)
-            
-        except Exception as e:
-            print(f"Error processing {video_name}: {str(e)}")
-            continue
-    
-    # Create DataFrame for better visualization
-    df = pd.DataFrame(results)
-    
-    # Print results
-    print("\n" + "=" * 80)
-    print("PSNR Results for Each Video:")
-    print("=" * 80)
-    print(df.to_string(index=False))
-    
-    # Print summary statistics
-    print("\n" + "=" * 80)
-    print("Summary Statistics:")
-    print("=" * 80)
-    print(f"Average PSNR across all videos: {df['overall_psnr'].mean():.4f} ± {df['overall_psnr'].std():.4f}")
-    print(f"Min PSNR: {df['overall_psnr'].min():.4f} ({df.loc[df['overall_psnr'].idxmin(), 'video_name']})")
-    print(f"Max PSNR: {df['overall_psnr'].max():.4f} ({df.loc[df['overall_psnr'].idxmax(), 'video_name']})")
-    print(f"Median PSNR: {df['overall_psnr'].median():.4f}")
-    
-    # Group by video ID (video_0, video_1, etc.) and rank
-    df['video_id'] = df['video_name'].str.extract(r'(video_\d+)')[0]
-    df['rank'] = df['video_name'].str.extract(r'rank(\d+)')[0].astype(int)
-    
-    print("\n" + "=" * 80)
-    print("Average PSNR by Video ID:")
-    print("=" * 80)
-    video_group = df.groupby('video_id')['overall_psnr'].agg(['mean', 'std', 'count'])
-    print(video_group.to_string())
-    
-    print("\n" + "=" * 80)
-    print("Average PSNR by Rank:")
-    print("=" * 80)
-    rank_group = df.groupby('rank')['overall_psnr'].agg(['mean', 'std', 'count'])
-    print(rank_group.to_string())
-    
-    # Save results to CSV
-    output_csv = gt_dir.parent / "psnr_results.csv"
-    df.to_csv(output_csv, index=False)
-    print(f"\nResults saved to: {output_csv}")
-    
-    # Save summary to text file
-    output_txt = gt_dir.parent / "psnr_summary.txt"
-    with open(output_txt, 'w') as f:
-        f.write("PSNR Results Summary\n")
-        f.write("=" * 80 + "\n\n")
-        f.write(f"Average PSNR: {df['overall_psnr'].mean():.4f} ± {df['overall_psnr'].std():.4f}\n")
-        f.write(f"Min PSNR: {df['overall_psnr'].min():.4f}\n")
-        f.write(f"Max PSNR: {df['overall_psnr'].max():.4f}\n")
-        f.write(f"Median PSNR: {df['overall_psnr'].median():.4f}\n")
-        f.write("\n" + "=" * 80 + "\n")
-        f.write("Full Results:\n")
-        f.write("=" * 80 + "\n")
-        f.write(df.to_string(index=False))
-    print(f"Summary saved to: {output_txt}")
-
-
-if __name__ == "__main__":
-    main()
-