cli.py

"""
Command Line Interface for Damaged Poster Detection
=================================================

Provides command-line interface for single image processing,
batch processing, and model training.
"""

import click
import os
import json
import logging
from pathlib import Path
from typing import Optional
import sys

from damaged_poster_detector import DamagedPosterDetector, DamageConfig
from train_model import DamageTrainer, TrainingConfig, create_sample_data_structure

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


@click.group()
@click.version_option(version="1.0.0")
def cli():
    """
    Damaged Poster Detection System CLI
    
    A production-ready system for detecting damage in poster images
    using computer vision and deep learning techniques.
    """
    pass


@cli.command()
@click.argument('image_path', type=click.Path(exists=True))
@click.option('--config', '-c', type=click.Path(exists=True), 
              help='Path to configuration file')
@click.option('--output', '-o', type=click.Path(), 
              help='Output JSON file for results')
@click.option('--verbose', '-v', is_flag=True, 
              help='Enable verbose output')
def detect(image_path: str, config: Optional[str], 
           output: Optional[str], verbose: bool):
    """
    Detect damage in a single poster image.
    
    IMAGE_PATH: Path to the poster image to analyze
    """
    if verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    try:
        # Initialize detector
        detector = DamagedPosterDetector(config_path=config)
        
        click.echo(f"Analyzing image: {image_path}")
        
        # Detect damage
        result = detector.detect_damage(image_path)
        
        if 'error' in result:
            click.echo(f"Error: {result['error']}", err=True)
            sys.exit(1)
        
        # Display results
        click.echo("\n" + "="*60)
        click.echo("DAMAGE DETECTION RESULTS")
        click.echo("="*60)
        
        overall = result['overall_status']
        click.echo(f"Overall Status: {'DAMAGED' if overall['damaged'] else 'UNDAMAGED'}")
        click.echo(f"Severity: {overall['severity'].upper()}")
        click.echo(f"Confidence: {overall['overall_confidence']:.2f}")
        click.echo(f"Detected Damage Types: {', '.join(overall['detected_damage_types'])}")
        
        click.echo("\nDetailed Analysis:")
        for damage_type, details in result['damage_details'].items():
            status = "✓" if details['detected'] else "✗"
            click.echo(f"  {status} {damage_type.replace('_', ' ').title()}: "
                      f"{details['confidence']:.3f}")
        
        click.echo("\nRecommendations:")
        for i, rec in enumerate(result['recommendations'], 1):
            click.echo(f"  {i}. {rec}")
        
        # Save to file if requested
        if output:
            with open(output, 'w') as f:
                json.dump(result, f, indent=2, default=str)
            click.echo(f"\nResults saved to: {output}")
        
    except Exception as e:
        click.echo(f"Error: {str(e)}", err=True)
        sys.exit(1)


@cli.command()
@click.argument('input_dir', type=click.Path(exists=True, file_okay=False))
@click.option('--output', '-o', type=click.Path(), 
              help='Output JSON file for batch results')
@click.option('--config', '-c', type=click.Path(exists=True), 
              help='Path to configuration file')
@click.option('--pattern', '-p', default='*', 
              help='File pattern to match (e.g., "*.jpg")')
@click.option('--verbose', '-v', is_flag=True, 
              help='Enable verbose output')
def batch(input_dir: str, output: Optional[str], config: Optional[str],
          pattern: str, verbose: bool):
    """
    Process multiple images in batch.
    
    INPUT_DIR: Directory containing images to process
    """
    if verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    try:
        # Initialize detector
        detector = DamagedPosterDetector(config_path=config)
        
        click.echo(f"Processing images in: {input_dir}")
        click.echo(f"File pattern: {pattern}")
        
        # Set default output file if not provided
        if not output:
            output = f"batch_results_{Path(input_dir).name}.json"
        
        # Process batch
        with click.progressbar(label='Processing images') as bar:
            results = detector.batch_process(input_dir, output)
            bar.update(len(results))
        
        # Summary statistics
        total_images = len(results)
        damaged_images = sum(1 for r in results if r.get('overall_status', {}).get('damaged', False))
        error_count = sum(1 for r in results if 'error' in r)
        
        click.echo(f"\n" + "="*50)
        click.echo("BATCH PROCESSING SUMMARY")
        click.echo("="*50)
        click.echo(f"Total images processed: {total_images}")
        click.echo(f"Damaged images found: {damaged_images}")
        click.echo(f"Undamaged images: {total_images - damaged_images - error_count}")
        click.echo(f"Processing errors: {error_count}")
        click.echo(f"Results saved to: {output}")
        
        # Show damage type distribution
        damage_counts = {}
        for result in results:
            if 'overall_status' in result:
                for damage_type in result['overall_status'].get('detected_damage_types', []):
                    damage_counts[damage_type] = damage_counts.get(damage_type, 0) + 1
        
        if damage_counts:
            click.echo("\nDamage Type Distribution:")
            for damage_type, count in sorted(damage_counts.items()):
                click.echo(f"  {damage_type.replace('_', ' ').title()}: {count}")
        
    except Exception as e:
        click.echo(f"Error: {str(e)}", err=True)
        sys.exit(1)


@cli.command()
@click.argument('data_dir', type=click.Path())
@click.option('--config', '-c', type=click.Path(), 
              help='Path to training configuration file')
@click.option('--epochs', '-e', type=int, default=30, 
              help='Number of training epochs')
@click.option('--batch-size', '-b', type=int, default=16, 
              help='Batch size for training')
@click.option('--learning-rate', '-lr', type=float, default=1e-4, 
              help='Learning rate')
@click.option('--create-structure', is_flag=True, 
              help='Create sample data directory structure')
@click.option('--wandb', is_flag=True, 
              help='Use Weights & Biases for experiment tracking')
@click.option('--verbose', '-v', is_flag=True, 
              help='Enable verbose output')
def train(data_dir: str, config: Optional[str], epochs: int, batch_size: int,
          learning_rate: float, create_structure: bool, wandb: bool, verbose: bool):
    """
    Train the damage detection model.
    
    DATA_DIR: Directory containing training data organized by damage type
    """
    if verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    try:
        # Create data structure if requested
        if create_structure:
            create_sample_data_structure(data_dir)
            click.echo(f"Created sample data structure in: {data_dir}")
            click.echo("Please organize your training images and run training again.")
            return
        
        # Check if data directory exists and has proper structure
        if not os.path.exists(data_dir):
            click.echo(f"Data directory not found: {data_dir}")
            click.echo("Use --create-structure to create the directory structure.")
            sys.exit(1)
        
        # Initialize training configuration
        train_config = TrainingConfig(
            data_dir=data_dir,
            num_epochs=epochs,
            batch_size=batch_size,
            learning_rate=learning_rate
        )
        
        # Load custom config if provided
        if config and os.path.exists(config):
            import yaml
            with open(config, 'r') as f:
                config_dict = yaml.safe_load(f)
            
            for key, value in config_dict.items():
                if hasattr(train_config, key):
                    setattr(train_config, key, value)
        
        click.echo("Training Configuration:")
        click.echo(f"  Data directory: {train_config.data_dir}")
        click.echo(f"  Epochs: {train_config.num_epochs}")
        click.echo(f"  Batch size: {train_config.batch_size}")
        click.echo(f"  Learning rate: {train_config.learning_rate}")
        click.echo(f"  Use WandB: {wandb}")
        
        # Initialize trainer
        trainer = DamageTrainer(train_config)
        
        # Start training
        click.echo("\nStarting training...")
        trainer.train(data_dir, use_wandb=wandb)
        
        # Plot training history
        trainer.plot_training_history()
        
        click.echo("Training completed successfully!")
        click.echo(f"Model saved in: {train_config.model_save_path}")
        click.echo(f"Logs saved in: {train_config.log_dir}")
        
    except Exception as e:
        click.echo(f"Error: {str(e)}", err=True)
        sys.exit(1)


@cli.command()
@click.option('--host', default='0.0.0.0', help='Host to bind to')
@click.option('--port', default=8000, help='Port to bind to')
@click.option('--workers', default=1, help='Number of worker processes')
@click.option('--reload', is_flag=True, help='Enable auto-reload for development')
def serve(host: str, port: int, workers: int, reload: bool):
    """
    Start the FastAPI web service.
    """
    try:
        import uvicorn
        
        click.echo(f"Starting web service on {host}:{port}")
        click.echo(f"Workers: {workers}")
        click.echo(f"Reload: {reload}")
        click.echo(f"API docs will be available at: http://{host}:{port}/docs")
        
        uvicorn.run(
            "api_service:app",
            host=host,
            port=port,
            workers=workers if not reload else 1,
            reload=reload
        )
        
    except ImportError:
        click.echo("FastAPI and uvicorn are required to run the web service.", err=True)
        click.echo("Install with: pip install fastapi uvicorn", err=True)
        sys.exit(1)
    except Exception as e:
        click.echo(f"Error starting service: {str(e)}", err=True)
        sys.exit(1)


@cli.command()
@click.option('--format', 'output_format', type=click.Choice(['json', 'table']), 
              default='table', help='Output format')
def info(output_format: str):
    """
    Display system information and configuration.
    """
    import torch
    import cv2
    import numpy as np
    
    info_data = {
        "system": {
            "python_version": sys.version,
            "opencv_version": cv2.__version__,
            "numpy_version": np.__version__,
            "pytorch_version": torch.__version__,
            "cuda_available": torch.cuda.is_available(),
            "cuda_version": torch.version.cuda if torch.cuda.is_available() else "N/A",
            "gpu_count": torch.cuda.device_count() if torch.cuda.is_available() else 0
        },
        "damage_types": [
            "tear", "crease", "stain", "fade", "burn", "water_damage"
        ],
        "supported_formats": [
            "jpg", "jpeg", "png", "bmp", "tiff", "tif"
        ]
    }
    
    if output_format == 'json':
        click.echo(json.dumps(info_data, indent=2))
    else:
        click.echo("System Information:")
        click.echo("=" * 50)
        for key, value in info_data["system"].items():
            click.echo(f"  {key.replace('_', ' ').title()}: {value}")
        
        click.echo(f"\nSupported Damage Types:")
        for damage_type in info_data["damage_types"]:
            click.echo(f"  - {damage_type.replace('_', ' ').title()}")
        
        click.echo(f"\nSupported Image Formats:")
        for fmt in info_data["supported_formats"]:
            click.echo(f"  - {fmt.upper()}")


if __name__ == '__main__':
    cli()