import os import sys import torch import argparse import logging from pathlib import Path from typing import List, Dict, Optional, Tuple from collections import OrderedDict import json from datetime import datetime from tqdm import tqdm import numpy as np import cv2 # OpenCV for video processing import matplotlib.pyplot as plt # Matplotlib for plotting import io # Image processing imports from PIL import Image, ImageDraw, ImageFont import torchvision.transforms as transforms # Import your existing model components # Ensure these files (models/, utils/) are in the same directory or accessible in PYTHONPATH from models.MIQA_base import get_torch_model, get_timm_model from models.RA_MIQA import RegionVisionTransformer from models.hf_model_registry import HF_REPO_ID, HF_REVISION, MODEL_FILENAMES from utils.hf_download_utils import ensure_checkpoint_from_hf SUPPORTED_VIDEO_EXTENSIONS = {'.mp4', '.avi', '.mov', '.mkv'} class MIQAInference: """ MODIFIED Inference wrapper for MIQA models. Now includes a method to predict on PIL Image objects directly. """ def __init__(self, task: str, model_name: str = 'ra_miqa', metric_type: str = 'composite', device: Optional[str] = None): self.task = task.lower() self.model_name = model_name self.metric_type = metric_type self.logger = self._setup_logger() if device is None: self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') else: self.device = torch.device(device) self.logger.info(f"🚀 Initializing MIQA Inference System") self.logger.info(f" Task: {self.task.upper()}") self.logger.info(f" Model: {self.model_name}") self.logger.info(f" Metric Type: {self.metric_type}") self.logger.info(f" Device: {self.device}") self._validate_config() self.model = self._load_model() self.transforms1, self.transforms2 = self._get_transforms() self.logger.info("✅ System ready for inference\n") def _setup_logger(self) -> logging.Logger: """Configure logging with both file and console output.""" logger = logging.getLogger('MIQA_Inference') logger.setLevel(logging.INFO) if logger.hasHandlers(): return logger logger.propagate = False # Console handler with clean formatting console_handler = logging.StreamHandler(sys.stdout) console_handler.setLevel(logging.INFO) console_formatter = logging.Formatter('%(message)s') console_handler.setFormatter(console_formatter) logger.addHandler(console_handler) return logger def _validate_config(self) -> None: """Validate that the requested configuration is supported.""" if self.metric_type not in ['composite', 'consistency', 'accuracy']: raise ValueError( f"Invalid metric_type '{self.metric_type}'. " f"Supported: ['composite', 'consistency', 'accuracy']" ) if self.task not in MODEL_FILENAMES[self.metric_type]: raise ValueError( f"Invalid task '{self.task}'. " f"Supported tasks: {list(MODEL_FILENAMES[self.metric_type].keys())}" ) if self.model_name not in MODEL_FILENAMES[self.metric_type][self.task]: available = list(MODEL_FILENAMES[self.metric_type][self.task].keys()) raise ValueError( f"Model '{self.model_name}' not available for task '{self.task}'. " f"Available models: {available}" ) def _get_checkpoint_path(self) -> str: """Generate the path where model checkpoint should be stored.""" base_dir = Path('models') / 'checkpoints' / f'{self.metric_type}_metric' base_dir.mkdir(parents=True, exist_ok=True) filename = MODEL_FILENAMES[self.metric_type][self.task][self.model_name] return str(base_dir / filename) def _download_weights(self, checkpoint_path: str) -> bool: """ Download model weights if not present locally. Returns: True if weights are available (already existed or successfully downloaded) """ if os.path.exists(checkpoint_path): self.logger.info(f"✓ Found cached model weights") return True self.logger.info( f"⏬ Downloading from Hugging Face: repo={HF_REPO_ID}, " f"file={Path(checkpoint_path).name}, rev={HF_REVISION}" ) try: ensure_checkpoint_from_hf( repo_id=HF_REPO_ID, filename=Path(checkpoint_path).name, local_dir=str(Path(checkpoint_path).parent), revision=HF_REVISION, ) self.logger.info("✓ Successfully downloaded model weights") return True except Exception as e: self.logger.error(f"❌ Failed to download model weights from Hugging Face: {e}") return False def _create_model(self) -> torch.nn.Module: """Create the model architecture.""" if self.model_name == 'ra_miqa': self.logger.info("Building Region-Aware Vision Transformer...") model = RegionVisionTransformer( base_model_name='vit_small_patch16_224', pretrained=False, # We'll load our trained weights mmseg_config_path='models/model_configs/fcn_sere-small_finetuned_fp16_8x32_224x224_3600_imagenets919.py', checkpoint_path='models/checkpoints/sere_finetuned_vit_small_ep100.pth' ) else: try: self.logger.info(f"Building {self.model_name} from PyTorch...") model = get_torch_model(model_name=self.model_name, pretrained=False, num_classes=1) except Exception: self.logger.info(f"Building {self.model_name} from timm library...") model = get_timm_model(model_name=self.model_name, pretrained=False, num_classes=1) return model def _load_model(self) -> torch.nn.Module: """Load model with pre-trained weights.""" checkpoint_path = self._get_checkpoint_path() # Ensure weights are available if not self._download_weights(checkpoint_path): raise RuntimeError("Cannot proceed without model weights") # Create model architecture self.logger.info("🔧 Loading model...") model = self._create_model() # Load weights checkpoint = torch.load(checkpoint_path, map_location='cpu') state_dict = checkpoint.get('state_dict', checkpoint) # Remove 'module.' prefix if present (from DataParallel training) new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k.replace('module.', '') if k.startswith('module.') else k new_state_dict[name] = v model.load_state_dict(new_state_dict, strict=True) model = model.to(self.device) model.eval() # Set to evaluation mode self.logger.info("✓ Model loaded successfully") return model def _get_transforms(self) -> Tuple[transforms.Compose, transforms.Compose | None]: """ Return preprocessing transforms based on model type. """ IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) SIMPLE_MEAN = (0.5, 0.5, 0.5) SIMPLE_STD = (0.5, 0.5, 0.5) # Default (for single-input backbones) transform_imagenet = transforms.Compose([ transforms.Resize(288), transforms.CenterCrop(size=224), transforms.ToTensor(), transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD) ]) transform_simple = transforms.Compose([ transforms.Resize(288), transforms.CenterCrop(size=224), transforms.ToTensor(), transforms.Normalize(mean=SIMPLE_MEAN, std=SIMPLE_STD) ]) # 1️⃣ CNNs（ResNet / EfficientNet） if any(k in self.model_name for k in ['resnet', 'efficientnet']): return transform_imagenet, None # 2️⃣ ViT elif 'vit' in self.model_name: return transform_simple, None # 3️⃣ ra_miqa elif 'ra_miqa' in self.model_name: transform_1 = transforms.Compose([ transforms.Resize(288), transforms.CenterCrop(size=224), transforms.ToTensor(), transforms.Normalize(mean=SIMPLE_MEAN, std=SIMPLE_STD) ]) transform_2 = transforms.Compose([ transforms.Resize(288), transforms.CenterCrop((288, 288)), transforms.ToTensor(), transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD) ]) return transform_1, transform_2 # fallback else: print(f"[Warning] Unknown model type '{self.model_name}', using ImageNet normalization.") return transform_imagenet, None @torch.no_grad() def predict_image_object(self, image: Image.Image) -> float: """ NEW METHOD: Run inference on a PIL Image object. """ # Preprocess the image img1 = self.transforms1(image).unsqueeze(0).to(self.device) img2 = self.transforms2(image).unsqueeze(0).to(self.device) if self.transforms2 else None # Run inference based on model input requirements if img2 is None: output = self.model(img1) else: output = self.model(img1, img2) score = output.item() if torch.is_tensor(output) else float(output) return score class VideoMIQAProcessor: """ A wrapper to process videos using the MIQAInference engine and create a visualized output video with scores and plots. """ # --- Visualization Constants --- PANEL_WIDTH = 480 FONT = cv2.FONT_HERSHEY_SIMPLEX FONT_SCALE_L = 1.0 FONT_SCALE_M = 0.8 FONT_COLOR = (255, 255, 255) # White LINE_THICKNESS = 2 # Plotting style plt.style.use('dark_background') def __init__(self, miqa_engine: MIQAInference): self.miqa_engine = miqa_engine self.logger = miqa_engine.logger def _create_score_plot(self, scores: List[float], width: int, height: int) -> np.ndarray: """ Creates a line chart of scores using Matplotlib and returns it as an OpenCV image. """ fig, ax = plt.subplots(figsize=(width / 100, height / 100), dpi=100) ax.plot(scores, color='#4287f5', linewidth=2) ax.set_xlim(0, max(1, len(scores))) ax.set_ylim(0, 1) ax.set_title("Quality Score Fluctuation", fontsize=10) ax.set_xlabel("Frame", fontsize=8) ax.set_ylabel("Score", fontsize=8) ax.grid(True, alpha=0.3) fig.tight_layout(pad=1.5) # Render plot to an in-memory buffer buf = io.BytesIO() fig.savefig(buf, format='png') buf.seek(0) plt.close(fig) # Convert buffer to a PIL Image and then to an OpenCV image plot_img_pil = Image.open(buf) plot_img_np = np.array(plot_img_pil) plot_img_bgr = cv2.cvtColor(plot_img_np, cv2.COLOR_RGBA2BGR) return plot_img_bgr def process_video(self, input_path: str, output_path: str): """ Reads a video, analyzes each frame for quality, and writes an annotated output video. """ self.logger.info(f"📹 Starting processing for: {Path(input_path).name}") cap = cv2.VideoCapture(input_path) if not cap.isOpened(): self.logger.error(f"❌ Failed to open video: {input_path}") return # Video properties fps = int(cap.get(cv2.CAP_PROP_FPS)) frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) orig_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) orig_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # New dimensions for output video (with side panel) output_width = orig_width + self.PANEL_WIDTH output_height = orig_height # Setup video writer fourcc = cv2.VideoWriter_fourcc(*'avc1') out = cv2.VideoWriter(output_path, fourcc, fps, (output_width, output_height)) scores = [] progress_bar = tqdm(range(frame_count), desc="Analyzing frames", ncols=100) for frame_idx in progress_bar: ret, frame = cap.read() if not ret: break # --- MIQA Inference --- frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) pil_image = Image.fromarray(frame_rgb) score = self.miqa_engine.predict_image_object(pil_image) scores.append(score) # --- Visualization Panel --- panel = np.zeros((orig_height, self.PANEL_WIDTH, 3), dtype=np.uint8) # 1. Task Info task_text = f"Task: {self.miqa_engine.task.upper()}" cv2.putText(panel, task_text, (20, 50), self.FONT, self.FONT_SCALE_M, self.FONT_COLOR, self.LINE_THICKNESS) # 2. Current Score score_text = f"Quality Score: {score:.3f}" # Color coding for score text norm_score = max(0, score) if norm_score < 0.5: color = (0, int(255 * (norm_score * 2)), 255) # Red -> Yellow else: color = (0, 255, int(255 * (2 - norm_score * 2))) # Yellow -> Green cv2.putText(panel, score_text, (20, 110), self.FONT, self.FONT_SCALE_L, color, self.LINE_THICKNESS + 1) # 3. Frame Info frame_text = f"Frame: {frame_idx + 1}/{frame_count}" cv2.putText(panel, frame_text, (20, orig_height - 30), self.FONT, self.FONT_SCALE_M, self.FONT_COLOR, 1) # 4. Score Plot if len(scores) > 1: plot_height = 300 plot_width = self.PANEL_WIDTH - 40 # with margins plot_img = self._create_score_plot(scores, plot_width, plot_height) # Position the plot on the panel y_offset = 160 panel[y_offset:y_offset + plot_img.shape[0], 20:20 + plot_img.shape[1]] = plot_img # --- Combine and Write Frame --- combined_frame = np.concatenate((frame, panel), axis=1) out.write(combined_frame) # Release resources cap.release() out.release() self.logger.info(f"✅ Finished processing. Annotated video saved to: {output_path}\n") def main(): """Command-line interface for Video MIQA inference.""" parser = argparse.ArgumentParser( description='MIQA for Video: Machine-centric Image Quality Assessment on Video Frames', formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Analyze a single video and save the annotated output python video_annotator_inference.py --input my_video.mp4 --task cls --model ra_miqa # Analyze all videos in a directory python video_annotator_inference.py --input ./video_folder/ --task det --model resnet50 """ ) parser.add_argument('--input', type=str, required=True, help='Path to input video file or a directory containing videos.') parser.add_argument('--task', type=str, required=True, choices=['cls', 'det', 'ins'], help='Task type: cls (classification), det (detection), ins (instance).') parser.add_argument('--model', type=str, default='ra_miqa', choices=['ra_miqa'], help='Model architecture (default: ra_miqa; Hub weights are RA-MIQA only).') parser.add_argument('--metric-type', type=str, default='composite', choices=['composite', 'consistency', 'accuracy'], help='Training metric type (default: composite).') parser.add_argument('--device', type=str, default=None, choices=['cuda', 'cpu'], help='Device to run on (auto-detect if not specified).') parser.add_argument('--output-dir', type=str, default='inference_results', help='Directory to save the output annotated videos.') args = parser.parse_args() try: # Initialize the core inference engine miqa_engine = MIQAInference( task=args.task, model_name=args.model, metric_type=args.metric_type, device=args.device ) # Initialize the video processor video_processor = VideoMIQAProcessor(miqa_engine) # Find videos to process input_path = Path(args.input) videos_to_process = [] if input_path.is_dir(): for ext in SUPPORTED_VIDEO_EXTENSIONS: videos_to_process.extend(input_path.glob(f"*{ext}")) elif input_path.is_file() and input_path.suffix.lower() in SUPPORTED_VIDEO_EXTENSIONS: videos_to_process.append(input_path) if not videos_to_process: raise FileNotFoundError(f"No supported video files found in '{args.input}'") # Create output directory output_dir = Path(args.output_dir) / 'video' /args.task / args.metric_type output_dir.mkdir(parents=True, exist_ok=True) # Process each video for video_path in videos_to_process: output_filename = f"{video_path.stem}_miqa_{args.model}_{args.task}.mp4" output_filepath = str(output_dir / output_filename) video_processor.process_video(str(video_path), output_filepath) except Exception as e: # Use the logger if it exists, otherwise print try: miqa_engine.logger.error(f"\n❌ An error occurred: {str(e)}") except: print(f"\n❌ An error occurred: {str(e)}", file=sys.stderr) sys.exit(1) if __name__ == '__main__': main()