img_inference.py

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

# Image processing imports
from PIL import Image, ImageDraw, ImageFont
import torchvision.transforms as transforms

# Import your existing model components
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 image file extensions
SUPPORTED_EXTENSIONS = {'.jpg', '.jpeg', 'JPEG', '.png', '.bmp', '.tiff', '.tif'}


class MIQAInference:
    """
    Inference wrapper for MIQA models.

    This class handles model initialization, automatic weight downloading,
    image preprocessing, batch prediction, and result visualization.
    """

    def __init__(self, task: str, model_name: str = 'ra_miqa',
                 metric_type: str = 'composite', device: Optional[str] = None):
        """
        Initialize the MIQA inference system.

        Args:
            task: Task type - 'cls' (classification), 'det' (detection), or 'ins' (instance)
            model_name: Model architecture to use (default: RA_MIQA for best performance)
            metric_type: Training objective - 'composite', 'consistency', or 'accuracy'
            device: Device to run inference on ('cuda' or 'cpu'). Auto-detects if None.
        """
        self.task = task.lower()
        self.model_name = model_name
        self.metric_type = metric_type

        # Setup logging with clean formatting
        self.logger = self._setup_logger()

        # Determine computation device
        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}")

        # Validate configuration
        self._validate_config()

        # Initialize model
        self.model = self._load_model()

        # Setup image preprocessing pipeline
        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

    def _prepare_image(self, image_path: str):
        """
        Load and preprocess a single image.
        Return (img1, img2, original_img)
        """
        img = Image.open(image_path).convert('RGB')
        img1 = self.transforms1(img).unsqueeze(0)
        img2 = self.transforms2(img).unsqueeze(0) if self.transforms2 else None
        return img1, img2, img

    @torch.no_grad()
    def predict_single(self, image_path: str) -> Dict:
        """
        Prediction interface for different backbones.
        """
        img1, img2, original_img = self._prepare_image(image_path)

        img1 = img1.to(self.device)
        if img2 is not None:
            img2 = img2.to(self.device)

        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 {
            'image_path': image_path,
            'image_name': Path(image_path).name,
            'quality_score': score,
            'original_image': original_img
        }

    @torch.no_grad()
    def predict_batch(self, image_paths: List[str],
                      show_progress: bool = True) -> List[Dict]:
        """
        Run inference on multiple images with progress tracking.

        Args:
            image_paths: List of paths to image files
            show_progress: Whether to display progress bar

        Returns:
            List of prediction results for each image
        """
        results = []

        # Create progress bar if requested
        iterator = tqdm(image_paths, desc="Processing images",
                        disable=not show_progress, ncols=80)

        for img_path in iterator:
            try:
                result = self.predict_single(img_path)
                results.append(result)
            except Exception as e:
                self.logger.warning(f"⚠️  Failed to process {img_path}: {str(e)}")
                results.append({
                    'image_path': img_path,
                    'image_name': Path(img_path).name,
                    'quality_score': None,
                    'error': str(e)
                })

        return results

    def predict(self, input_path: str, show_progress: bool = True) -> List[Dict]:
        """
        Main prediction interface - handles both single images and directories.

        Args:
            input_path: Path to an image file or directory containing images
            show_progress: Whether to show progress bar for batch processing

        Returns:
            List of prediction results
        """
        input_path = Path(input_path)

        # Handle single file
        if input_path.is_file():
            if input_path.suffix.lower() not in SUPPORTED_EXTENSIONS:
                raise ValueError(
                    f"Unsupported file extension: {input_path.suffix}. "
                    f"Supported: {SUPPORTED_EXTENSIONS}"
                )
            return [self.predict_single(str(input_path))]

        # Handle directory
        elif input_path.is_dir():
            # Find all supported images in directory
            image_paths = []
            for ext in SUPPORTED_EXTENSIONS:
                image_paths.extend(input_path.glob(f"*{ext}"))
                # image_paths.extend(input_path.glob(f"*{ext.upper()}"))

            image_paths = sorted([str(p) for p in image_paths])

            if not image_paths:
                raise ValueError(f"No supported images found in {input_path}")

            self.logger.info(f"📁 Found {len(image_paths)} images in directory")
            return self.predict_batch(image_paths, show_progress)

        else:
            raise ValueError(f"Input path does not exist: {input_path}")

    def visualize_results(self, results: List[Dict], output_dir: str = 'inference_results',
                          score_range: Tuple[float, float] = (0, 100)) -> None:
        """
        Create annotated visualizations of predictions.

        This generates images with quality scores overlaid in a color-coded box.
        Low scores appear in red, high scores in green.

        Args:
            results: List of prediction results
            output_dir: Directory to save visualizations
            score_range: Expected range of quality scores for color normalization
        """
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)

        self.logger.info(f"\n🎨 Creating visualizations...")

        for result in tqdm(results, desc="Generating visualizations", ncols=80):
            if result.get('quality_score') is None:
                continue  # Skip failed predictions

            img = result['original_image'].copy()
            draw = ImageDraw.Draw(img)

            # Prepare score display
            score = result['quality_score']
            score_text = f"Quality: {score:.4f}"

            # Normalize score to [0, 1] for color interpolation
            norm_score = (score - score_range[0]) / (score_range[1] - score_range[0])
            norm_score = max(0, min(1, norm_score))  # Clamp to [0, 1]

            # Color interpolation: red (low quality) -> yellow -> green (high quality)
            if norm_score < 0.5:
                # Red to yellow
                r = 255
                g = int(255 * (norm_score * 2))
                b = 0
            else:
                # Yellow to green
                r = int(255 * (2 - norm_score * 2))
                g = 255
                b = 0

            color = (r, g, b)

            # Draw colored box with score
            box_width = 250
            box_height = 50
            margin = 10

            # Position in top-left corner
            box_coords = [margin, margin, margin + box_width, margin + box_height]
            draw.rectangle(box_coords, fill=color)

            # Add text (try to load a nice font, fall back to default)
            try:
                font = ImageFont.truetype("arial.ttf", 30)
            except:
                font = ImageFont.load_default()

            # Calculate text position to center it in the box
            bbox = draw.textbbox((0, 0), score_text, font=font)
            text_width = bbox[2] - bbox[0]
            text_height = bbox[3] - bbox[1]
            text_x = margin + (box_width - text_width) // 2
            text_y = margin + (box_height - text_height) // 2

            draw.text((text_x, text_y), score_text, fill='black', font=font)

            # Save annotated image
            output_file = output_path / f"miqa_{self.model_name}_{result['image_name']}"
            img.save(output_file)

        self.logger.info(f"✓ Visualizations saved to: {output_dir}/")

    def save_results(self, results: List[Dict], output_path: str = 'predictions.json',
                     format: str = 'json') -> None:
        """
        Save prediction results to file.

        Args:
            results: List of prediction results
            output_path: Path to save results
            format: Output format - 'json' or 'csv'
        """
        # Remove PIL image objects before saving
        clean_results = []
        for r in results:
            clean_r = {k: v for k, v in r.items() if k != 'original_image'}
            clean_results.append(clean_r)

        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        if format == 'json':
            with open(output_path, 'w') as f:
                json.dump({
                    'metadata': {
                        'task': self.task,
                        'model': self.model_name,
                        'metric_type': self.metric_type,
                        'timestamp': datetime.now().isoformat(),
                        'num_images': len(clean_results)
                    },
                    'predictions': clean_results
                }, f, indent=2)

        elif format == 'csv':
            import csv
            with open(output_path, 'w', newline='') as f:
                if clean_results:
                    writer = csv.DictWriter(f, fieldnames=clean_results[0].keys())
                    writer.writeheader()
                    writer.writerows(clean_results)

        self.logger.info(f"💾 Results saved to: {output_path}")

    def print_summary(self, results: List[Dict]) -> None:
        """Print a formatted summary of prediction results."""
        valid_results = [r for r in results if r.get('quality_score') is not None]
        failed_results = [r for r in results if r.get('quality_score') is None]

        self.logger.info("\n" + "=" * 80)
        self.logger.info("PREDICTION SUMMARY")
        self.logger.info("=" * 80)

        if valid_results:
            scores = [r['quality_score'] for r in valid_results]
            self.logger.info(f"✓ Successfully processed: {len(valid_results)} images")
            self.logger.info(f"  Average quality score: {np.mean(scores):.4f}")
            self.logger.info(f"  Score range: [{np.min(scores):.4f}, {np.max(scores):.4f}]")
            self.logger.info(f"  Standard deviation: {np.std(scores):.4f}")

            # Show top and bottom quality images
            sorted_results = sorted(valid_results, key=lambda x: x['quality_score'], reverse=True)

            self.logger.info("\n🏆 Top 3 quality images:")
            for i, r in enumerate(sorted_results[:3], 1):
                self.logger.info(f"  {i}. {r['image_name']}: {r['quality_score']:.4f}")

            if len(sorted_results) > 3:
                self.logger.info("\n⚠️  Bottom 3 quality images:")
                for i, r in enumerate(sorted_results[-3:], 1):
                    self.logger.info(f"  {i}. {r['image_name']}: {r['quality_score']:.4f}")

        if failed_results:
            self.logger.info(f"\n❌ Failed to process: {len(failed_results)} images")

        self.logger.info("=" * 80 + "\n")


def main():
    """Command-line interface for MIQA inference."""
    parser = argparse.ArgumentParser(
        description='MIQA: Machine-centric Image Quality Assessment',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
    Examples:
      # Predict quality of a single image
      python img_inference.py --input image.jpg --task cls --model ra_miqa
    
      # Process all images in a directory
      python img_inference.py --input ./assets/demo_images/imagenet_demo --task det --model ra_miqa
    
      # Save results and create visualizations
      python img_inference.py --input /assets/demo_images/imagenet_demo --task ins --save-results --visualize
            """
        )

    parser.add_argument('--input', type=str, required=True,
                        help='Path to input image or directory containing images')
    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('--save-results', action='store_true',
                        help='Save prediction results to file')
    parser.add_argument('--output-file', type=str, default='predictions.json',
                        help='Output file path for results (default: predictions.json)')
    parser.add_argument('--output-format', type=str, default='json',
                        choices=['json', 'csv'],
                        help='Output file format (default: json)')
    parser.add_argument('--visualize', action='store_true',
                        help='Create annotated visualizations of predictions')
    parser.add_argument('--save-dir', type=str, default='inference_results',
                        help='Directory for save (default: inference_results)')
    parser.add_argument('--no-progress', action='store_true',
                        help='Disable progress bar')

    args = parser.parse_args()

    try:
        # Initialize inference system
        miqa = MIQAInference(
            task=args.task,
            model_name=args.model,
            metric_type=args.metric_type,
            device=args.device
        )

        # Run predictions
        results = miqa.predict(args.input, show_progress=not args.no_progress)

        # Print summary
        miqa.print_summary(results)

        args.save_dir = Path(args.save_dir)/ 'image' / args.task / args.metric_type
        args.output_file = f'miqa_{args.model}_'+args.output_file

        # Save results if requested
        if args.save_results:
            miqa.save_results(results, Path(args.save_dir)/args.output_file, args.output_format)

        # Create visualizations if requested
        if args.visualize:
            miqa.visualize_results(results, args.save_dir)

    except Exception as e:
        print(f"\n❌ Error: {str(e)}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()