Coverage for markdiffusion / evaluation / tools / video_quality_analyzer.py: 93.37%

1from typing import List 

2from PIL import Image 

3import torch 

4import torch.nn as nn 

5import torch.nn.functional as F 

6from torchvision import transforms 

7from torchvision.transforms import Compose, Resize, ToTensor, Normalize, CenterCrop 

8import numpy as np 

9from tqdm import tqdm 

10import cv2 

11import os 

12import subprocess 

13from markdiffusion.utils.media_utils import pil_to_torch 

14 

15try: 

16 from torchvision.transforms import InterpolationMode 

17 BICUBIC = InterpolationMode.BICUBIC 

18except ImportError: 

19 BICUBIC = Image.BICUBIC 

20 

21from pathlib import Path 

22 

23# Package root (parent of `markdiffusion/evaluation`). Used to resolve bundled 

24# model checkpoints (.pth) and a hyphen-named module file by absolute path, 

25# regardless of the user's current working directory. 

26_MARKDIFFUSION_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 

27_DEFAULT_AMT_S_CKPT = os.path.join(_MARKDIFFUSION_ROOT, "model", "amt", "amt-s.pth") 

28_DEFAULT_AMT_S_MODULE = os.path.join(_MARKDIFFUSION_ROOT, "model", "amt", "networks", "AMT-S.py") 

29_DEFAULT_RAFT_CKPT = os.path.join(_MARKDIFFUSION_ROOT, "model", "raft", "raft-things.pth") 

30 

31if not hasattr(np, 'sctypes'): 

32 np.sctypes = { 

33 'int': [np.int8, np.int16, np.int32, np.int64], 

34 'uint': [np.uint8, np.uint16, np.uint32, np.uint64], 

35 'float': [np.float16, np.float32, np.float64], 

36 'complex': [np.complex64, np.complex128], 

37 'others': [bool, object, bytes, str, np.void] 

38 } 

39 

40def dino_transform_Image(n_px): 

41 """DINO transform for PIL Images.""" 

42 return Compose([ 

43 Resize(size=n_px, antialias=False), 

44 ToTensor(), 

45 Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) 

46 ]) 

47 

48 

49class VideoQualityAnalyzer: 

50 """Video quality analyzer base class.""" 

51 

52 def __init__(self): 

53 pass 

54 

55 def analyze(self, frames: List[Image.Image]): 

56 """Analyze video quality. 

57  

58 Args: 

59 frames: List of PIL Image frames representing the video 

60  

61 Returns: 

62 Quality score(s) 

63 """ 

64 raise NotImplementedError("Subclasses must implement analyze method") 

65 

66 

67class SubjectConsistencyAnalyzer(VideoQualityAnalyzer): 

68 """Analyzer for evaluating subject consistency across video frames using DINO features. 

69  

70 This analyzer measures how consistently the main subject appears across frames by: 

71 1. Extracting DINO features from each frame 

72 2. Computing cosine similarity between consecutive frames and with the first frame 

73 3. Averaging these similarities to get a consistency score 

74 """ 

75 def __init__( 

76 self, 

77 model_url: str = "https://dl.fbaipublicfiles.com/dino/dino_vitbase16_pretrain/dino_vitbase16_pretrain_full_checkpoint.pth", 

78 model_path: str = "dino_vitb16_full.pth", 

79 device: str = "cuda" 

80 ): 

81 self.device = torch.device(device if torch.cuda.is_available() else "cpu") 

82 self.model_path = model_path 

83 self.model_url = model_url 

84 

85 # ensure weights exist / download automatically 

86 self._download_weights() 

87 

88 # load model via timm 

89 self.model = self._load_dino_model() 

90 self.model.eval() 

91 self.model.to(self.device) 

92 

93 def _download_weights(self): 

94 if not os.path.exists(self.model_path): 

95 import urllib 

96 print("Downloading DINO ViT-B/16 weights...") 

97 urllib.request.urlretrieve(self.model_url, self.model_path) 

98 print("Download complete:", self.model_path) 

99 else: 

100 print("Weights already exist:", self.model_path) 

101 

102 def _load_dino_model(self): 

103 import timm 

104 # timm vit-base-p16 structure 

105 model = timm.create_model( 

106 "vit_base_patch16_224", 

107 pretrained=False, 

108 num_classes=0 

109 ) 

110 

111 # load full checkpoint 

112 ckpt = torch.load(self.model_path, map_location="cpu") 

113 

114 # for full checkpoint the state dict is nested 

115 if "teacher" in ckpt: 

116 state_dict = ckpt["teacher"] 

117 elif "student" in ckpt: 

118 state_dict = ckpt["student"] 

119 else: 

120 state_dict = ckpt 

121 

122 # remove classifier head keys 

123 state_dict = {k: v for k, v in state_dict.items() if "head" not in k} 

124 

125 model.load_state_dict(state_dict, strict=False) 

126 return model 

127 

128 def transform(self, img: Image.Image) -> torch.Tensor: 

129 """Transform PIL Image to tensor for DINO model.""" 

130 transform = dino_transform_Image(224) 

131 return transform(img) 

132 

133 def analyze(self, frames: List[Image.Image]) -> float: 

134 """Analyze subject consistency across video frames. 

135  

136 Args: 

137 frames: List of PIL Image frames representing the video 

138  

139 Returns: 

140 Subject consistency score (higher is better, range [0, 1]) 

141 """ 

142 if len(frames) < 2: 

143 return 1.0 # Single frame is perfectly consistent with itself 

144 

145 video_sim = 0.0 

146 frame_count = 0 

147 

148 # Process frames and extract features 

149 with torch.no_grad(): 

150 for i, frame in enumerate(frames): 

151 # Transform and prepare frame 

152 frame_tensor = self.transform(frame).unsqueeze(0).to(self.device) 

153 

154 # Extract features 

155 features = self.model(frame_tensor) 

156 features = F.normalize(features, dim=-1, p=2) 

157 

158 if i == 0: 

159 # Store first frame features 

160 first_frame_features = features 

161 else: 

162 # Compute similarity with previous frame 

163 sim_prev = max(0.0, F.cosine_similarity(prev_features, features).item()) 

164 

165 # Compute similarity with first frame 

166 sim_first = max(0.0, F.cosine_similarity(first_frame_features, features).item()) 

167 

168 # Average the two similarities 

169 frame_sim = (sim_prev + sim_first) / 2.0 

170 video_sim += frame_sim 

171 frame_count += 1 

172 

173 # Store current features as previous for next iteration 

174 prev_features = features 

175 

176 # Return average similarity across all frame pairs 

177 if frame_count > 0: 

178 return video_sim / frame_count 

179 else: 

180 return 1.0 

181 

182# from contextlib import contextmanager 

183 

184# @contextmanager 

185# def isolated_import_context(code_dir, isolated_prefixes, prefix_tag=None): 

186# """Context manager for isolated module imports to avoid conflicts with main project. 

187 

188# Args: 

189# code_dir: External code directory to add to sys.path 

190# isolated_prefixes: List of module name prefixes to isolate (e.g., ['utils', 'networks']) 

191# prefix_tag: Tag to prefix external modules with after loading (default: code_dir.name + '_ext_') 

192 

193# Example: 

194# with isolated_import_context(CODE_DIR, ['utils', 'networks']): 

195# # imports here will use CODE_DIR's modules 

196# spec = importlib.util.spec_from_file_location("entry", CODE_DIR / "main.py") 

197# ... 

198# # after exiting, main project's 'utils' is restored 

199# """ 

200# import sys 

201 

202# if prefix_tag is None: 

203# prefix_tag = code_dir.name + '_ext_' 

204 

205# original_path = sys.path.copy() 

206# saved_modules = {} 

207 

208# # Remove potentially conflicting modules 

209# for prefix in isolated_prefixes: 

210# for mod_name in list(sys.modules.keys()): 

211# if mod_name == prefix or mod_name.startswith(prefix + '.'): 

212# saved_modules[mod_name] = sys.modules.pop(mod_name) 

213 

214# sys.path.insert(0, str(code_dir)) 

215 

216# try: 

217# yield 

218# finally: 

219# sys.path[:] = original_path 

220 

221# # Rename external modules with prefix tag to avoid future conflicts 

222# for prefix in isolated_prefixes: 

223# for mod_name in list(sys.modules.keys()): 

224# if mod_name == prefix or mod_name.startswith(prefix + '.'): 

225# if mod_name not in saved_modules: 

226# sys.modules[prefix_tag + mod_name] = sys.modules.pop(mod_name) 

227 

228# # Restore main project modules 

229# sys.modules.update(saved_modules) 

230 

231class MotionSmoothnessAnalyzer(VideoQualityAnalyzer): 

232 """Analyzer for evaluating motion smoothness in videos using AMT-S model. 

233  

234 This analyzer measures motion smoothness by: 

235 1. Extracting frames at even indices from the video 

236 2. Using AMT-S model to interpolate between consecutive frames 

237 3. Comparing interpolated frames with actual frames to compute smoothness score 

238  

239 The score represents how well the motion can be predicted/interpolated, 

240 with smoother motion resulting in higher scores. 

241 """ 

242 

243 def __init__(self, model_path: str = _DEFAULT_AMT_S_CKPT, 

244 device: str = "cuda", niters: int = 1): 

245 """Initialize the MotionSmoothnessAnalyzer. 

246  

247 Args: 

248 model_path: Path to the AMT-S model checkpoint 

249 device: Device to run the model on ('cuda' or 'cpu') 

250 niters: Number of interpolation iterations (default: 1) 

251 """ 

252 self.device = torch.device(device if torch.cuda.is_available() else "cpu") 

253 self.niters = niters 

254 

255 # Initialize model parameters 

256 self._initialize_params() 

257 

258 # Load AMT-S model 

259 self.model = self._load_amt_model(model_path) 

260 self.model.eval() 

261 self.model.to(self.device) 

262 

263 def _initialize_params(self): 

264 """Initialize parameters for video processing.""" 

265 if self.device.type == 'cuda': 

266 self.anchor_resolution = 1024 * 512 

267 self.anchor_memory = 1500 * 1024**2 

268 self.anchor_memory_bias = 2500 * 1024**2 

269 self.vram_avail = torch.cuda.get_device_properties(self.device).total_memory 

270 else: 

271 # Do not resize in cpu mode 

272 self.anchor_resolution = 8192 * 8192 

273 self.anchor_memory = 1 

274 self.anchor_memory_bias = 0 

275 self.vram_avail = 1 

276 

277 # Time embedding for interpolation (t=0.5) 

278 self.embt = torch.tensor(1/2).float().view(1, 1, 1, 1).to(self.device) 

279 

280 def _load_amt_model(self, model_path: str): 

281 """Load AMT-S model. 

282  

283 Args: 

284 model_path: Path to the model checkpoint 

285  

286 Returns: 

287 Loaded AMT-S model 

288 """ 

289 # Import AMT-S model (note the hyphen in filename) 

290 import sys 

291 import importlib.util 

292 

293 # Load the module with hyphen in filename 

294 spec = importlib.util.spec_from_file_location("amt_s", _DEFAULT_AMT_S_MODULE) 

295 amt_s_module = importlib.util.module_from_spec(spec) 

296 spec.loader.exec_module(amt_s_module) 

297 Model = amt_s_module.Model 

298 

299 # Create model with default parameters 

300 model = Model( 

301 corr_radius=3, 

302 corr_lvls=4, 

303 num_flows=3 

304 ) 

305 

306 # Load checkpoint 

307 if os.path.exists(model_path): 

308 ckpt = torch.load(model_path, map_location="cpu", weights_only=False) 

309 model.load_state_dict(ckpt['state_dict']) 

310 

311 return model 

312 

313 def _extract_frames(self, frames: List[Image.Image], start_from: int = 0) -> List[np.ndarray]: 

314 """Extract frames at even indices starting from start_from. 

315  

316 Args: 

317 frames: List of PIL Image frames 

318 start_from: Starting index (default: 0) 

319  

320 Returns: 

321 List of extracted frames as numpy arrays 

322 """ 

323 extracted = [] 

324 for i in range(start_from, len(frames), 2): 

325 # Convert PIL Image to numpy array 

326 frame_np = np.array(frames[i]) 

327 extracted.append(frame_np) 

328 return extracted 

329 

330 def _img2tensor(self, img: np.ndarray) -> torch.Tensor: 

331 """Convert numpy image to tensor. 

332  

333 Args: 

334 img: Image as numpy array (H, W, C) 

335  

336 Returns: 

337 Image tensor (1, C, H, W) 

338 """ 

339 from markdiffusion.model.amt.utils.utils import img2tensor 

340 return img2tensor(img) 

341 

342 def _tensor2img(self, tensor: torch.Tensor) -> np.ndarray: 

343 """Convert tensor to numpy image. 

344  

345 Args: 

346 tensor: Image tensor (1, C, H, W) 

347  

348 Returns: 

349 Image as numpy array (H, W, C) 

350 """ 

351 from markdiffusion.model.amt.utils.utils import tensor2img 

352 return tensor2img(tensor) 

353 

354 def _check_dim_and_resize(self, tensor_list: List[torch.Tensor]) -> List[torch.Tensor]: 

355 """Check dimensions and resize tensors if needed. 

356  

357 Args: 

358 tensor_list: List of image tensors 

359  

360 Returns: 

361 List of resized tensors 

362 """ 

363 from markdiffusion.model.amt.utils.utils import check_dim_and_resize 

364 return check_dim_and_resize(tensor_list) 

365 

366 def _calculate_scale(self, h: int, w: int) -> float: 

367 """Calculate scaling factor based on available VRAM. 

368  

369 Args: 

370 h: Height of the image 

371 w: Width of the image 

372  

373 Returns: 

374 Scaling factor 

375 """ 

376 scale = self.anchor_resolution / (h * w) * np.sqrt((self.vram_avail - self.anchor_memory_bias) / self.anchor_memory) 

377 scale = 1 if scale > 1 else scale 

378 scale = 1 / np.floor(1 / np.sqrt(scale) * 16) * 16 

379 return scale 

380 

381 def _interpolate_frames(self, inputs: List[torch.Tensor], scale: float) -> List[torch.Tensor]: 

382 """Interpolate frames using AMT-S model. 

383 

384 Args: 

385 inputs: List of input frame tensors 

386 scale: Scaling factor for processing 

387 

388 Returns: 

389 List of interpolated frame tensors 

390 """ 

391 

392 from markdiffusion.model.amt.utils.utils import InputPadder 

393 # Pad inputs 

394 padding = int(16 / scale) 

395 padder = InputPadder(inputs[0].shape, padding) 

396 inputs = padder.pad(*inputs) 

397 

398 # Perform interpolation for specified iterations 

399 for _ in range(self.niters): 

400 outputs = [inputs[0]] 

401 for in_0, in_1 in zip(inputs[:-1], inputs[1:]): 

402 in_0 = in_0.to(self.device) 

403 in_1 = in_1.to(self.device) 

404 with torch.no_grad(): 

405 imgt_pred = self.model(in_0, in_1, self.embt, scale_factor=scale, eval=True)['imgt_pred'] 

406 outputs += [imgt_pred.cpu(), in_1.cpu()] 

407 inputs = outputs 

408 

409 # Unpad outputs 

410 outputs = padder.unpad(*outputs) 

411 return outputs 

412 

413 def _compute_frame_difference(self, img1: np.ndarray, img2: np.ndarray) -> float: 

414 """Compute average absolute difference between two images. 

415  

416 Args: 

417 img1: First image 

418 img2: Second image 

419  

420 Returns: 

421 Average pixel difference 

422 """ 

423 diff = cv2.absdiff(img1, img2) 

424 return np.mean(diff) 

425 

426 def _compute_vfi_score(self, original_frames: List[np.ndarray], interpolated_frames: List[np.ndarray]) -> float: 

427 """Compute video frame interpolation score. 

428  

429 Args: 

430 original_frames: Original video frames 

431 interpolated_frames: Interpolated frames 

432  

433 Returns: 

434 VFI score (lower difference means better interpolation) 

435 """ 

436 # Extract frames at odd indices for comparison 

437 ori_compare = self._extract_frames([Image.fromarray(f) for f in original_frames], start_from=1) 

438 interp_compare = self._extract_frames([Image.fromarray(f) for f in interpolated_frames], start_from=1) 

439 

440 scores = [] 

441 for ori, interp in zip(ori_compare, interp_compare): 

442 score = self._compute_frame_difference(ori, interp) 

443 scores.append(score) 

444 

445 return np.mean(scores) if scores else 0.0 

446 

447 def analyze(self, frames: List[Image.Image]) -> float: 

448 """Analyze motion smoothness in video frames. 

449  

450 Args: 

451 frames: List of PIL Image frames representing the video 

452  

453 Returns: 

454 Motion smoothness score (higher is better, range [0, 1]) 

455 """ 

456 if len(frames) < 2: 

457 return 1.0 # Single frame has perfect smoothness 

458 

459 # Convert PIL Images to numpy arrays 

460 np_frames = [np.array(frame) for frame in frames] 

461 

462 # Extract frames at even indices 

463 frame_list = self._extract_frames(frames, start_from=0) 

464 

465 # Convert to tensors 

466 inputs = [self._img2tensor(frame).to(self.device) for frame in frame_list] 

467 

468 if len(inputs) <= 1: 

469 return 1.0 # Not enough frames for interpolation 

470 

471 # Check dimensions and resize if needed 

472 inputs = self._check_dim_and_resize(inputs) 

473 h, w = inputs[0].shape[-2:] 

474 

475 # Calculate scale based on available memory 

476 scale = self._calculate_scale(h, w) 

477 

478 # Perform frame interpolation 

479 outputs = self._interpolate_frames(inputs, scale) 

480 

481 # Convert outputs back to images 

482 output_images = [self._tensor2img(out) for out in outputs] 

483 

484 # Compute VFI score 

485 vfi_score = self._compute_vfi_score(np_frames, output_images) 

486 

487 # Normalize score to [0, 1] range (higher is better) 

488 # Original score is average pixel difference [0, 255], we normalize and invert 

489 normalized_score = (255.0 - vfi_score) / 255.0 

490 

491 return normalized_score 

492 

493 

494class DynamicDegreeAnalyzer(VideoQualityAnalyzer): 

495 """Analyzer for evaluating dynamic degree (motion intensity) in videos using RAFT optical flow. 

496  

497 This analyzer measures the amount and intensity of motion in videos by: 

498 1. Computing optical flow between consecutive frames using RAFT 

499 2. Calculating flow magnitude for each pixel 

500 3. Extracting top 5% highest flow magnitudes 

501 4. Determining if video has sufficient dynamic motion based on thresholds 

502  

503 The score represents whether the video contains dynamic motion (1.0) or is mostly static (0.0). 

504 """ 

505 

506 def __init__(self, model_path: str = _DEFAULT_RAFT_CKPT, 

507 device: str = "cuda", sample_fps: int = 8): 

508 """Initialize the DynamicDegreeAnalyzer. 

509  

510 Args: 

511 model_path: Path to the RAFT model checkpoint 

512 device: Device to run the model on ('cuda' or 'cpu') 

513 sample_fps: Target FPS for frame sampling (default: 8) 

514 """ 

515 self.device = torch.device(device if torch.cuda.is_available() else "cpu") 

516 self.sample_fps = sample_fps 

517 

518 # Load RAFT model 

519 self.model = self._load_raft_model(model_path) 

520 self.model.eval() 

521 self.model.to(self.device) 

522 

523 def _load_raft_model(self, model_path: str): 

524 """Load RAFT optical flow model. 

525  

526 Args: 

527 model_path: Path to the model checkpoint 

528  

529 Returns: 

530 Loaded RAFT model 

531 """ 

532 from markdiffusion.model.raft.core.raft import RAFT 

533 from easydict import EasyDict as edict 

534 

535 # Configure RAFT arguments 

536 args = edict({ 

537 "model": model_path, 

538 "small": False, 

539 "mixed_precision": False, 

540 "alternate_corr": False 

541 }) 

542 

543 # Create and load model 

544 model = RAFT(args) 

545 

546 if os.path.exists(model_path): 

547 ckpt = torch.load(model_path, map_location="cpu") 

548 # Remove 'module.' prefix if present (from DataParallel) 

549 new_ckpt = {k.replace('module.', ''): v for k, v in ckpt.items()} 

550 model.load_state_dict(new_ckpt) 

551 

552 return model 

553 

554 def _extract_frames_for_flow(self, frames: List[Image.Image], target_fps: int = 8) -> List[torch.Tensor]: 

555 """Extract and prepare frames for optical flow computation. 

556  

557 Args: 

558 frames: List of PIL Image frames 

559 target_fps: Target sampling rate (default: 8 fps) 

560  

561 Returns: 

562 List of prepared frame tensors 

563 """ 

564 # Estimate original FPS and calculate sampling interval 

565 # Assuming 30fps original video, adjust sampling to get ~8fps 

566 total_frames = len(frames) 

567 assumed_fps = 30 # Common video fps 

568 interval = max(1, round(assumed_fps / target_fps)) 

569 

570 # Sample frames at interval 

571 sampled_frames = [] 

572 for i in range(0, total_frames, interval): 

573 frame = frames[i] 

574 # Convert PIL to numpy array 

575 frame_np = np.array(frame) 

576 # Convert to tensor and normalize 

577 frame_tensor = torch.from_numpy(frame_np.astype(np.uint8)).permute(2, 0, 1).float() 

578 frame_tensor = frame_tensor[None].to(self.device) 

579 sampled_frames.append(frame_tensor) 

580 

581 return sampled_frames 

582 

583 def _compute_flow_magnitude(self, flow: torch.Tensor) -> float: 

584 """Compute flow magnitude score from optical flow. 

585  

586 Args: 

587 flow: Optical flow tensor (B, 2, H, W) 

588  

589 Returns: 

590 Flow magnitude score 

591 """ 

592 # Extract flow components 

593 flow_np = flow[0].permute(1, 2, 0).cpu().numpy() 

594 u = flow_np[:, :, 0] 

595 v = flow_np[:, :, 1] 

596 

597 # Compute flow magnitude 

598 magnitude = np.sqrt(np.square(u) + np.square(v)) 

599 

600 # Get top 5% highest magnitudes 

601 h, w = magnitude.shape 

602 magnitude_flat = magnitude.flatten() 

603 cut_index = int(h * w * 0.05) 

604 

605 # Sort in descending order and take mean of top 5% 

606 top_magnitudes = np.sort(-magnitude_flat)[:cut_index] 

607 mean_magnitude = np.mean(np.abs(top_magnitudes)) 

608 

609 return mean_magnitude.item() 

610 

611 def _determine_dynamic_threshold(self, frame_shape: tuple, num_frames: int) -> dict: 

612 """Determine thresholds for dynamic motion detection. 

613  

614 Args: 

615 frame_shape: Shape of the frame tensor 

616 num_frames: Number of frames in the video 

617  

618 Returns: 

619 Dictionary with threshold parameters 

620 """ 

621 # Scale threshold based on image resolution 

622 scale = min(frame_shape[-2:]) # min of height and width 

623 magnitude_threshold = 6.0 * (scale / 256.0) 

624 

625 # Scale count threshold based on number of frames 

626 count_threshold = round(4 * (num_frames / 16.0)) 

627 

628 return { 

629 "magnitude_threshold": magnitude_threshold, 

630 "count_threshold": count_threshold 

631 } 

632 

633 def _check_dynamic_motion(self, flow_scores: List[float], thresholds: dict) -> bool: 

634 """Check if video has dynamic motion based on flow scores. 

635  

636 Args: 

637 flow_scores: List of optical flow magnitude scores 

638 thresholds: Threshold parameters 

639  

640 Returns: 

641 True if video has dynamic motion, False otherwise 

642 """ 

643 magnitude_threshold = thresholds["magnitude_threshold"] 

644 count_threshold = thresholds["count_threshold"] 

645 

646 # Count frames with significant motion 

647 motion_count = 0 

648 for score in flow_scores: 

649 if score > magnitude_threshold: 

650 motion_count += 1 

651 if motion_count >= count_threshold: 

652 return True 

653 

654 return False 

655 

656 def analyze(self, frames: List[Image.Image]) -> float: 

657 """Analyze dynamic degree (motion intensity) in video frames. 

658  

659 Args: 

660 frames: List of PIL Image frames representing the video 

661  

662 Returns: 

663 Dynamic degree score: 1.0 if video has dynamic motion, 0.0 if mostly static 

664 """ 

665 if len(frames) < 2: 

666 return 0.0 # Cannot compute optical flow with less than 2 frames 

667 

668 # Extract and prepare frames for optical flow 

669 prepared_frames = self._extract_frames_for_flow(frames, self.sample_fps) 

670 

671 if len(prepared_frames) < 2: 

672 return 0.0 

673 

674 # Determine thresholds based on video characteristics 

675 thresholds = self._determine_dynamic_threshold( 

676 prepared_frames[0].shape, 

677 len(prepared_frames) 

678 ) 

679 

680 # Compute optical flow between consecutive frames 

681 flow_scores = [] 

682 

683 with torch.no_grad(): 

684 for frame1, frame2 in zip(prepared_frames[:-1], prepared_frames[1:]): 

685 # Pad frames if necessary 

686 from markdiffusion.model.raft.core.utils_core.utils import InputPadder 

687 padder = InputPadder(frame1.shape) 

688 frame1_padded, frame2_padded = padder.pad(frame1, frame2) 

689 

690 # Compute optical flow 

691 _, flow_up = self.model(frame1_padded, frame2_padded, iters=20, test_mode=True) 

692 

693 # Calculate flow magnitude score 

694 magnitude_score = self._compute_flow_magnitude(flow_up) 

695 flow_scores.append(magnitude_score) 

696 

697 # Check if video has dynamic motion 

698 has_dynamic_motion = self._check_dynamic_motion(flow_scores, thresholds) 

699 

700 # Return binary score: 1.0 for dynamic, 0.0 for static 

701 return 1.0 if has_dynamic_motion else 0.0 

702 

703 

704class BackgroundConsistencyAnalyzer(VideoQualityAnalyzer): 

705 """Analyzer for evaluating background consistency across video frames using CLIP features. 

706  

707 This analyzer measures how consistently the background appears across frames by: 

708 1. Extracting CLIP visual features from each frame 

709 2. Computing cosine similarity between consecutive frames and with the first frame 

710 3. Averaging these similarities to get a consistency score 

711  

712 Similar to SubjectConsistencyAnalyzer but focuses on overall visual consistency 

713 including background elements, making it suitable for detecting background stability. 

714 """ 

715 

716 def __init__(self, model_name: str = "ViT-B/32", device: str = "cuda"): 

717 """Initialize the BackgroundConsistencyAnalyzer. 

718  

719 Args: 

720 model_name: CLIP model name (default: "ViT-B/32") 

721 device: Device to run the model on ('cuda' or 'cpu') 

722 """ 

723 self.device = torch.device(device if torch.cuda.is_available() else "cpu") 

724 

725 # Load CLIP model 

726 self.model, self.preprocess = self._load_clip_model(model_name) 

727 self.model.eval() 

728 self.model.to(self.device) 

729 

730 # Image transform for CLIP (when processing tensor inputs) 

731 self.tensor_transform = self._get_clip_tensor_transform(224) 

732 

733 def _load_clip_model(self, model_name: str): 

734 """Load CLIP model. 

735  

736 Args: 

737 model_name: Name of the CLIP model to load 

738  

739 Returns: 

740 Tuple of (model, preprocess_function) 

741 """ 

742 import clip 

743 

744 model, preprocess = clip.load(model_name, device=self.device) 

745 return model, preprocess 

746 

747 def _get_clip_tensor_transform(self, n_px: int): 

748 """Get CLIP transform for tensor inputs. 

749  

750 Args: 

751 n_px: Target image size 

752  

753 Returns: 

754 Transform composition for tensor inputs 

755 """ 

756 try: 

757 from torchvision.transforms import InterpolationMode 

758 BICUBIC = InterpolationMode.BICUBIC 

759 except ImportError: 

760 BICUBIC = Image.BICUBIC 

761 

762 return Compose([ 

763 Resize(n_px, interpolation=BICUBIC, antialias=False), 

764 CenterCrop(n_px), 

765 transforms.Lambda(lambda x: x.float().div(255.0)), 

766 Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)), 

767 ]) 

768 

769 def _prepare_images_for_clip(self, frames: List[Image.Image]) -> torch.Tensor: 

770 """Prepare PIL images for CLIP processing. 

771  

772 Args: 

773 frames: List of PIL Image frames 

774  

775 Returns: 

776 Batch tensor of preprocessed images 

777 """ 

778 # Use CLIP's built-in preprocess for PIL images 

779 images = [] 

780 for frame in frames: 

781 processed = self.preprocess(frame) 

782 images.append(processed) 

783 

784 # Stack into batch tensor 

785 return torch.stack(images).to(self.device) 

786 

787 def analyze(self, frames: List[Image.Image]) -> float: 

788 """Analyze background consistency across video frames. 

789  

790 Args: 

791 frames: List of PIL Image frames representing the video 

792  

793 Returns: 

794 Background consistency score (higher is better, range [0, 1]) 

795 """ 

796 if len(frames) < 2: 

797 return 1.0 # Single frame is perfectly consistent with itself 

798 

799 # Prepare images for CLIP 

800 images = self._prepare_images_for_clip(frames) 

801 

802 # Extract CLIP features 

803 with torch.no_grad(): 

804 image_features = self.model.encode_image(images) 

805 image_features = F.normalize(image_features, dim=-1, p=2) 

806 

807 video_sim = 0.0 

808 frame_count = 0 

809 

810 # Compute similarity between frames 

811 for i in range(len(image_features)): 

812 image_feature = image_features[i].unsqueeze(0) 

813 

814 if i == 0: 

815 # Store first frame features 

816 first_image_feature = image_feature 

817 else: 

818 # Compute similarity with previous frame 

819 sim_prev = max(0.0, F.cosine_similarity(former_image_feature, image_feature).item()) 

820 

821 # Compute similarity with first frame 

822 sim_first = max(0.0, F.cosine_similarity(first_image_feature, image_feature).item()) 

823 

824 # Average the two similarities 

825 frame_sim = (sim_prev + sim_first) / 2.0 

826 video_sim += frame_sim 

827 frame_count += 1 

828 

829 # Store current features as previous for next iteration 

830 former_image_feature = image_feature 

831 

832 # Return average similarity across all frame pairs 

833 if frame_count > 0: 

834 return video_sim / frame_count 

835 else: 

836 return 1.0 

837 

838class ImagingQualityAnalyzer(VideoQualityAnalyzer): 

839 """Analyzer for evaluating imaging quality of videos. 

840  

841 This analyzer measures the quality of videos by: 

842 1. Inputting frames into MUSIQ image quality predictor 

843 2. Determining if the video is blurry or has artifacts 

844  

845 The score represents the quality of the video (higher is better). 

846 """ 

847 def __init__(self, model_path: str = "musiq_spaq_ckpt-358bb6af.pth", device: str = "cuda"): 

848 self.device = torch.device(device if torch.cuda.is_available() else "cpu") 

849 self.model = self._load_musiq(model_path) 

850 self.model.to(self.device) 

851 self.model.eval() 

852 

853 def _load_musiq(self, model_path: str): 

854 """Load MUSIQ model. 

855  

856 Args: 

857 model_path: Path to the MUSIQ model checkpoint 

858  

859 Returns: 

860 MUSIQ model 

861 """ 

862 from pathlib import Path 

863 CACHE_DIR = Path("/tmp/musiq_cache") 

864 model_path = CACHE_DIR / model_path 

865 

866 # if the model_path not exists 

867 # then makedir and wget 

868 if not os.path.exists(model_path): 

869 os.makedirs(os.path.dirname(model_path), exist_ok=True) 

870 wget_command = ['wget', 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/musiq_spaq_ckpt-358bb6af.pth', '-P', os.path.dirname(model_path)] 

871 subprocess.run(wget_command, check=True) 

872 try: 

873 from pyiqa.archs.musiq_arch import MUSIQ 

874 except ImportError: 

875 raise ImportError("Please install pyiqa to use ImagingQualityAnalyzer: pip install pyiqa") 

876 model = MUSIQ(pretrained_model_path=str(model_path)) 

877 

878 return model 

879