IQA-Metric-Benchmark/main.py

#!/usr/bin/env python3
"""
Main script for Image Quality Assessment using DeQA scoring.
"""

import argparse
import sys
from pathlib import Path

# Add src directory to path for imports
sys.path.append(str(Path(__file__).parent / "src"))

from src.iqa_analyzer import IQAAnalyzer
from src.logger_config import setup_logging


def main():
    """Main function to run the IQA analysis."""
    parser = argparse.ArgumentParser(
        description="IQA Benchmark Runner - Analyze image quality using DeQA scoring",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run analysis on default image directory
  python main.py
  
  # Run analysis on specific directory
  python main.py --image-dir /path/to/images
  
  # Run with verbose logging and save results as CSV
  python main.py --verbose --output-format csv
  
  # Run analysis and save to custom output directory
  python main.py --output-dir custom_results
        """
    )
    
    parser.add_argument(
        '--image-dir',
        type=str,
        default='data/task/cni/images',
        help='Directory containing images to analyze (default: data/task/cni/images)'
    )
    
    parser.add_argument(
        '--output-dir',
        type=str,
        default='results',
        help='Directory to save results and reports (default: results)'
    )
    
    parser.add_argument(
        '--log-dir',
        type=str,
        default='logs',
        help='Directory to store log files (default: logs)'
    )
    
    parser.add_argument(
        '--output-format',
        choices=['json', 'csv', 'txt'],
        default='json',
        help='Output format for results (default: json)'
    )
    
    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='Enable verbose logging'
    )
    
    parser.add_argument(
        '--enable-deqa',
        action='store_true',
        default=True,
        help='Enable DeQA metric (default: True)'
    )
    
    parser.add_argument(
        '--enable-traditional',
        action='store_true',
        default=False,
        help='Enable traditional metrics (default: False)'
    )

    parser.add_argument(
        '--enable-pyiqa',
        action='store_true',
        default=True,
        help='Enable PyIQA metrics (default: True)'
    )

    parser.add_argument(
        '--pyiqa-top20',
        action='store_true',
        help='Use curated top 20 PyIQA metrics only'
    )
    
    parser.add_argument(
        '--deqa-only',
        action='store_true',
        help='Use only DeQA metric (disable traditional metrics)'
    )
    
    args = parser.parse_args()
    
    # Setup logging
    log_level = "DEBUG" if args.verbose else "INFO"
    logger = setup_logging(log_dir=args.log_dir, log_level=log_level)
    
    # Validate image directory
    image_dir = Path(args.image_dir)
    if not image_dir.exists():
        logger.error(f"Image directory does not exist: {image_dir}")
        sys.exit(1)
    
    if not image_dir.is_dir():
        logger.error(f"Path is not a directory: {image_dir}")
        sys.exit(1)
    
    # Initialize and run IQA analysis
    try:
        logger.info("Starting IQA analysis...")
        logger.info(f"Image directory: {image_dir}")
        logger.info(f"Output directory: {args.output_dir}")
        
        # Determine which metrics to enable
        enable_deqa = args.enable_deqa and not args.deqa_only
        enable_traditional = args.enable_traditional and not args.deqa_only
        
        # User-requested 20 metrics (NR + FR) from PyIQA
        selected_top20 = [
            # No-Reference (for OCR practical use)
            'brisque', 'niqe', 'piqe', 'nrqm', 'nima', 'paq2piq', 'dbcnn', 'hyperiqa',
            'musiq', 'topiq_nr', 'clipiqa+_vitL14_512', 'maniqa', 'ahiq', 'unique', 'uranker',
            # Full-Reference (benchmark)
            'ssim', 'ms_ssim', 'fsim', 'vif', 'dists'
        ]

        analyzer = IQAAnalyzer(
            str(image_dir), 
            args.output_dir,
            enable_deqa=enable_deqa,
            enable_traditional=enable_traditional,
            enable_pyiqa=args.enable_pyiqa,
            pyiqa_selected_metrics=(selected_top20 if args.pyiqa_top20 else None)
        )
        results, report = analyzer.run_analysis()
        
        # Save results only in requested format; default to TXT if none
        if args.output_format == 'csv':
            analyzer.save_results('csv')
        else:
            # TXT as the primary output for this workflow
            analyzer.save_results('txt')
        
        # Optionally display report (keep for visibility)
        print("\n" + report)
        
        logger.info("IQA analysis completed successfully!")
        
    except Exception as e:
        logger.error(f"Error during IQA analysis: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()