main.py

import os
import argparse
from pathlib import Path
import tensorflow as tf
from typing import Dict

from model import FogRemovalGAN
from training import FogDataset, GANTrainer, GANInference

def parse_args() -> argparse.Namespace:
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(description='Fog Removal GAN')
    
    # Model architecture parameters
    parser.add_argument('--growth_rate', type=int, default=12,
                      help='Growth rate for dense blocks (default: 12)')
    parser.add_argument('--layers', type=int, default=4,
                      help='Number of layers per dense block (default: 4)')
    parser.add_argument('--D_filters', type=int, default=64,
                      help='Number of filters in first discriminator layer (default: 64)')
    
    # Training parameters
    parser.add_argument('--batch_size', type=int, default=1,
                      help='Batch size for training (default: 1)')
    parser.add_argument('--epochs', type=int, default=200,
                      help='Number of epochs to train (default: 200)')
    parser.add_argument('--lr', type=float, default=0.0002,
                      help='Learning rate (default: 0.0002)')
    parser.add_argument('--gan_wt', type=float, default=2.0,
                      help='Weight for GAN loss (default: 2.0)')
    parser.add_argument('--l1_wt', type=float, default=100.0,
                      help='Weight for L1 loss (default: 100.0)')
    parser.add_argument('--vgg_wt', type=float, default=10.0,
                      help='Weight for VGG perceptual loss (default: 10.0)')
    parser.add_argument('--val_frequency', type=int, default=5,
                      help='Validation frequency in epochs (default: 5)')
    
    # Data parameters
    parser.add_argument('--data_root', type=str, required=True,
                      help='Root directory containing train/val/test folders')
    parser.add_argument('--image_size', type=int, default=256,
                      help='Size of input images (default: 256)')
    
    # Model saving/loading
    parser.add_argument('--model_name', type=str, default='fog_removal_gan',
                      help='Name for saving model and logs (default: fog_removal_gan)')
    parser.add_argument('--restore', action='store_true',
                      help='Restore from latest checkpoint')
    
    # Execution mode
    parser.add_argument('--mode', type=str, choices=['train', 'test', 'inference'],
                      default='train', help='Execution mode (default: train)')
    parser.add_argument('--gpu', type=str, default='0',
                      help='GPU to use (default: 0)')
    
    return parser.parse_args()

def setup_gpu(gpu: str):
    """Setup GPU device"""
    os.environ['CUDA_VISIBLE_DEVICES'] = gpu
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        try:
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
        except RuntimeError as e:
            print(e)

def train(args: argparse.Namespace):
    """Training execution path"""
    # Create dataset handler
    dataset = FogDataset(
        root_dir=args.data_root,
        batch_size=args.batch_size,
        image_size=(args.image_size, args.image_size)
    )
    
    # Create and compile model
    model = FogRemovalGAN(args)
    model.compile()
    
    # Create trainer and start training
    trainer = GANTrainer(model, dataset, args)
    trainer.train()

def test(args: argparse.Namespace):
    """Testing execution path"""
    # Create dataset handler for test set
    dataset = FogDataset(
        root_dir=args.data_root,
        batch_size=1,  # Always use batch size 1 for testing
        image_size=(args.image_size, args.image_size)
    )
    test_dataset = dataset.get_dataset('test')
    
    # Create model and load weights
    model = FogRemovalGAN(args)
    checkpoint_dir = Path(args.model_name) / 'checkpoints'
    inference = GANInference(model, str(checkpoint_dir))
    
    # Process test set
    output_dir = Path(args.model_name) / 'test_results'
    output_dir.mkdir(parents=True, exist_ok=True)
    inference.process_directory(
        str(Path(args.data_root) / 'test'),
        str(output_dir)
    )

def inference(args: argparse.Namespace):
    """Inference execution path"""
    # Create model and load weights
    model = FogRemovalGAN(args)
    checkpoint_dir = Path(args.model_name) / 'checkpoints'
    inference = GANInference(model, str(checkpoint_dir))
    
    # Setup input/output directories
    input_dir = Path(args.data_root)
    output_dir = Path(args.model_name) / 'inference_results'
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Process images
    inference.process_directory(str(input_dir), str(output_dir))

def main():
    # Parse arguments
    args = parse_args()
    
    # Setup GPU
    setup_gpu(args.gpu)
    
    # Print key settings
    print("\nRunning with settings:")
    print(f"Mode: {args.mode}")
    print(f"Data root: {args.data_root}")
    print(f"Model name: {args.model_name}")
    print(f"Batch size: {args.batch_size}")
    if args.mode == 'train':
        print(f"Epochs: {args.epochs}")
        print(f"Learning rate: {args.lr}")
    print()
    
    # Execute appropriate mode
    if args.mode == 'train':
        train(args)
    elif args.mode == 'test':
        test(args)
    else:
        inference(args)

if __name__ == '__main__':
    main()