demo.py

#!/usr/bin/env python3
"""
Demo script for the Image Classifier project
Shows the complete workflow from data preparation to prediction
"""

import os
import sys
import time
import logging
from pathlib import Path

# Add src to path
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))

from image_classifier import ImageClassifier
from utils import create_sample_dataset, validate_data_directory, check_gpu_availability
import config

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


def print_banner():
    """Print project banner"""
    banner = """
    ╔══════════════════════════════════════════════════════════════╗
    ║                    Image Classifier Demo                     ║
    ║                                                              ║
    ║  A comprehensive machine learning project for image          ║
    ║  classification using TensorFlow and CNN                    ║
    ╚══════════════════════════════════════════════════════════════╝
    """
    print(banner)


def check_environment():
    """Check and display environment information"""
    logger.info("Checking environment...")
    
    # Check Python version
    python_version = sys.version_info
    logger.info(f"Python version: {python_version.major}.{python_version.minor}.{python_version.micro}")
    
    # Check GPU availability
    gpu_available = check_gpu_availability()
    
    # Check TensorFlow
    try:
        import tensorflow as tf
        tf_version = tf.__version__
        logger.info(f"TensorFlow version: {tf_version}")
    except ImportError:
        logger.error("TensorFlow not installed!")
        return False
    
    # Check other dependencies
    dependencies = ['numpy', 'cv2', 'matplotlib', 'sklearn']
    for dep in dependencies:
        try:
            __import__(dep)
            logger.info(f"✓ {dep} available")
        except ImportError:
            logger.error(f"✗ {dep} not available")
            return False
    
    logger.info("Environment check completed successfully!")
    return True


def prepare_demo_data():
    """Prepare demo data"""
    logger.info("Preparing demo data...")
    
    data_dir = config.PATHS['data']
    
    # Create sample dataset
    create_sample_dataset(data_dir, num_samples=30)
    logger.info(f"Created sample dataset in {data_dir}")
    
    # Validate the data
    is_valid, message = validate_data_directory(data_dir)
    if is_valid:
        logger.info("✓ Data validation passed")
    else:
        logger.error(f"✗ Data validation failed: {message}")
        return False
    
    return True


def run_training_demo():
    """Run the training demo"""
    logger.info("Starting training demo...")
    
    # Initialize classifier
    classifier = ImageClassifier(input_shape=config.MODEL_CONFIG['input_shape'])
    
    # Load data
    train_dataset, val_dataset, test_dataset = classifier.load_data(
        config.PATHS['data'],
        batch_size=config.MODEL_CONFIG['batch_size'],
        validation_split=config.MODEL_CONFIG['validation_split']
    )
    
    # Build model
    classifier.build_model()
    
    # Display model summary
    logger.info("Model architecture:")
    classifier.model.summary()
    
    # Train model
    start_time = time.time()
    history = classifier.train(
        train_dataset,
        val_dataset,
        epochs=10,  # Reduced for demo
        patience=3
    )
    training_time = time.time() - start_time
    
    logger.info(f"Training completed in {training_time:.2f} seconds")
    
    # Evaluate model
    results = classifier.evaluate(test_dataset)
    logger.info(f"Test Results: {results}")
    
    # Plot training history
    classifier.plot_training_history("demo_training_history.png")
    
    # Save model
    classifier.save_model(config.PATHS['trained_model'])
    
    return classifier, results


def run_prediction_demo(classifier):
    """Run the prediction demo"""
    logger.info("Starting prediction demo...")
    
    # Find a sample image
    data_dir = Path(config.PATHS['data'])
    sample_image = data_dir / "class1" / "sample_0.jpg"
    
    if sample_image.exists():
        # Make prediction
        predicted_class, confidence, processed_image = classifier.predict_image(str(sample_image))
        
        logger.info(f"Prediction Demo Results:")
        logger.info(f"  Image: {sample_image.name}")
        logger.info(f"  Predicted Class: {predicted_class}")
        logger.info(f"  Confidence: {confidence:.3f}")
        
        return True
    else:
        logger.warning("No sample image found for prediction demo")
        return False


def display_results(results):
    """Display demo results"""
    logger.info("Demo Results Summary:")
    logger.info("=" * 50)
    logger.info(f"Accuracy:  {results['accuracy']:.3f}")
    logger.info(f"Precision: {results['precision']:.3f}")
    logger.info(f"Recall:    {results['recall']:.3f}")
    logger.info("=" * 50)


def cleanup_demo():
    """Clean up demo files"""
    logger.info("Cleaning up demo files...")
    
    files_to_remove = [
        "demo_training_history.png",
        "training.log"
    ]
    
    for file in files_to_remove:
        if os.path.exists(file):
            os.remove(file)
            logger.info(f"Removed {file}")


def main():
    """Main demo function"""
    print_banner()
    
    # Check environment
    if not check_environment():
        logger.error("Environment check failed. Please install required dependencies.")
        return False
    
    try:
        # Prepare data
        if not prepare_demo_data():
            logger.error("Failed to prepare demo data")
            return False
        
        # Run training demo
        classifier, results = run_training_demo()
        
        # Run prediction demo
        prediction_success = run_prediction_demo(classifier)
        
        # Display results
        display_results(results)
        
        # Cleanup
        cleanup_demo()
        
        logger.info("Demo completed successfully!")
        logger.info("You can now:")
        logger.info("1. Use the trained model for predictions")
        logger.info("2. Modify the code for your own dataset")
        logger.info("3. Explore the project structure and documentation")
        
        return True
        
    except Exception as e:
        logger.error(f"Demo failed: {e}")
        return False


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)