240 lines
8.2 KiB
Python
240 lines
8.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
import argparse
|
|
import os
|
|
import json
|
|
from fastvideo.logger import init_logger
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def skip_if_gated_repo_inaccessible(repo_id: str,
|
|
*,
|
|
local_path: str | None = None,
|
|
test_name: str = "test",
|
|
allow_module_level: bool = False) -> None:
|
|
"""Skip a test when ``repo_id`` is gated/private and the token lacks access.
|
|
|
|
Local weights win: when ``local_path`` already exists the check is skipped
|
|
entirely, so cached machines keep running offline. Transient hub failures
|
|
(offline, DNS, 429/5xx) do NOT skip either — the subsequent download will
|
|
serve from cache or fail loudly, preserving the CI failure signal. Only a
|
|
positively-identified authorization problem produces a skip.
|
|
"""
|
|
import pytest
|
|
|
|
if local_path is not None and os.path.exists(local_path):
|
|
return
|
|
|
|
from huggingface_hub import HfApi
|
|
from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError
|
|
|
|
try:
|
|
HfApi().auth_check(repo_id)
|
|
except (GatedRepoError, RepositoryNotFoundError) as exc:
|
|
pytest.skip(
|
|
f"Skipping {test_name}: the configured HuggingFace token cannot "
|
|
f"access the gated repo {repo_id}: {exc}",
|
|
allow_module_level=allow_module_level,
|
|
)
|
|
except Exception as exc: # noqa: BLE001 - hub unreachable, proxies, 5xx
|
|
logger.warning(
|
|
"Gated-repo access probe for %s failed (%s); proceeding — the "
|
|
"download will serve from cache or fail loudly.", repo_id, exc)
|
|
|
|
|
|
def _read_video_frames(path: str) -> torch.Tensor:
|
|
"""Read video frames as a ``(T, C, H, W)`` uint8 tensor.
|
|
|
|
Tries ``torchvision.io.read_video`` first (available in older
|
|
torchvision) and falls back to PyAV when the function has been
|
|
removed (torchvision >= 0.26).
|
|
"""
|
|
try:
|
|
from torchvision.io import read_video
|
|
|
|
frames, _, _ = read_video(
|
|
path,
|
|
pts_unit="sec",
|
|
output_format="TCHW",
|
|
)
|
|
return frames
|
|
except (ImportError, AttributeError):
|
|
pass
|
|
|
|
import av
|
|
|
|
container = av.open(path)
|
|
frames = []
|
|
for frame in container.decode(video=0):
|
|
arr = frame.to_ndarray(format="rgb24")
|
|
frames.append(torch.from_numpy(arr).permute(2, 0, 1))
|
|
container.close()
|
|
if not frames:
|
|
raise RuntimeError(f"No video frames decoded from {path}")
|
|
return torch.stack(frames)
|
|
|
|
|
|
def _read_image_as_single_frame_video(path: str) -> torch.Tensor:
|
|
"""Read one image as a single-frame ``(1, C, H, W)`` uint8 tensor."""
|
|
from torchvision.io import read_image
|
|
|
|
img = read_image(path)
|
|
return img.unsqueeze(0)
|
|
|
|
|
|
def _read_visual_frames(path: str) -> torch.Tensor:
|
|
"""Read a video or a single image as ``(T, C, H, W)`` uint8."""
|
|
ext = os.path.splitext(path)[1].lower()
|
|
if ext in {".png", ".jpg", ".jpeg", ".webp"}:
|
|
return _read_image_as_single_frame_video(path)
|
|
return _read_video_frames(path)
|
|
|
|
|
|
def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
|
|
"""
|
|
Compute SSIM between two videos or single-frame image files.
|
|
|
|
Image paths (``.png``, ``.jpg``, ``.jpeg``, ``.webp``) are treated as
|
|
one-frame clips so T2I SSIM can share the same MS-SSIM path as video.
|
|
|
|
Args:
|
|
video1_path: Path to the first video or image.
|
|
video2_path: Path to the second video or image.
|
|
use_ms_ssim: Whether to use Multi-Scale Structural Similarity(MS-SSIM) instead of SSIM.
|
|
"""
|
|
from pytorch_msssim import ms_ssim, ssim
|
|
|
|
print(f"Computing SSIM between {video1_path} and {video2_path}...")
|
|
if not os.path.exists(video1_path):
|
|
raise FileNotFoundError(f"Video1 not found: {video1_path}")
|
|
if not os.path.exists(video2_path):
|
|
raise FileNotFoundError(f"Video2 not found: {video2_path}")
|
|
|
|
frames1 = _read_visual_frames(video1_path)
|
|
frames2 = _read_visual_frames(video2_path)
|
|
|
|
# Ensure same number of frames
|
|
min_frames = min(frames1.shape[0], frames2.shape[0])
|
|
frames1 = frames1[:min_frames]
|
|
frames2 = frames2[:min_frames]
|
|
|
|
frames1 = frames1.float() / 255.0
|
|
frames2 = frames2.float() / 255.0
|
|
|
|
if torch.cuda.is_available():
|
|
frames1 = frames1.cuda()
|
|
frames2 = frames2.cuda()
|
|
|
|
ssim_values = []
|
|
|
|
# Process each frame individually
|
|
for i in range(min_frames):
|
|
img1 = frames1[i:i + 1]
|
|
img2 = frames2[i:i + 1]
|
|
|
|
with torch.no_grad():
|
|
value = ms_ssim(img1, img2, data_range=1.0) if use_ms_ssim else ssim(img1, img2, data_range=1.0)
|
|
|
|
ssim_values.append(value.item())
|
|
|
|
if ssim_values:
|
|
mean_ssim = np.mean(ssim_values)
|
|
min_ssim = np.min(ssim_values)
|
|
max_ssim = np.max(ssim_values)
|
|
min_frame_idx = np.argmin(ssim_values)
|
|
max_frame_idx = np.argmax(ssim_values)
|
|
|
|
print(f"Mean SSIM: {mean_ssim:.4f}")
|
|
print(f"Min SSIM: {min_ssim:.4f} (at frame {min_frame_idx})")
|
|
print(f"Max SSIM: {max_ssim:.4f} (at frame {max_frame_idx})")
|
|
|
|
return mean_ssim, min_ssim, max_ssim
|
|
else:
|
|
print("No SSIM values calculated")
|
|
return 0, 0, 0
|
|
|
|
|
|
def compare_folders(reference_folder, generated_folder, use_ms_ssim=True):
|
|
"""
|
|
Compare videos with the same filename between reference_folder and generated_folder
|
|
|
|
Example usage:
|
|
results = compare_folders(reference_folder, generated_folder,
|
|
args.use_ms_ssim)
|
|
for video_name, ssim_value in results.items():
|
|
if ssim_value is not None:
|
|
print(
|
|
f"{video_name}: {ssim_value[0]:.4f}, Min SSIM: {ssim_value[1]:.4f}, Max SSIM: {ssim_value[2]:.4f}"
|
|
)
|
|
else:
|
|
print(f"{video_name}: Error during comparison")
|
|
|
|
valid_ssims = [v for v in results.values() if v is not None]
|
|
if valid_ssims:
|
|
avg_ssim = np.mean([v[0] for v in valid_ssims])
|
|
print(f"\nAverage SSIM across all videos: {avg_ssim:.4f}")
|
|
else:
|
|
print("\nNo valid SSIM values to average")
|
|
"""
|
|
|
|
reference_videos = [f for f in os.listdir(reference_folder) if f.endswith(".mp4")]
|
|
|
|
results = {}
|
|
|
|
for video_name in reference_videos:
|
|
ref_path = os.path.join(reference_folder, video_name)
|
|
gen_path = os.path.join(generated_folder, video_name)
|
|
|
|
if os.path.exists(gen_path):
|
|
print(f"\nComparing {video_name}...")
|
|
try:
|
|
ssim_value = compute_video_ssim_torchvision(ref_path, gen_path, use_ms_ssim)
|
|
results[video_name] = ssim_value
|
|
except Exception as e:
|
|
print(f"Error comparing {video_name}: {e}")
|
|
results[video_name] = None
|
|
else:
|
|
print(f"\nSkipping {video_name} - no matching file in generated folder")
|
|
|
|
return results
|
|
|
|
|
|
def write_ssim_results(output_dir, ssim_values, reference_path, generated_path, num_inference_steps, prompt):
|
|
"""
|
|
Write SSIM results to a JSON file in the same directory as the generated videos.
|
|
"""
|
|
try:
|
|
logger.info(f"Attempting to write SSIM results to directory: {output_dir}")
|
|
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
mean_ssim, min_ssim, max_ssim = ssim_values
|
|
|
|
result = {
|
|
"mean_ssim": mean_ssim,
|
|
"min_ssim": min_ssim,
|
|
"max_ssim": max_ssim,
|
|
"reference_video": reference_path,
|
|
"generated_video": generated_path,
|
|
"parameters": {
|
|
"num_inference_steps": num_inference_steps,
|
|
"prompt": prompt
|
|
},
|
|
}
|
|
|
|
test_name = f"steps{num_inference_steps}_{prompt[:100]}"
|
|
result_file = os.path.join(output_dir, f"{test_name}_ssim.json")
|
|
logger.info(f"Writing JSON results to: {result_file}")
|
|
with open(result_file, "w") as f:
|
|
json.dump(result, f, indent=2)
|
|
|
|
logger.info(f"SSIM results written to {result_file}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"ERROR writing SSIM results: {str(e)}")
|
|
return False
|