main.py

"""Main entry point for CBCT Report Generation pipeline.

Commands:
    preprocess: Preprocess ToothFairy3 dataset and extract regional tokens
    train: Train the report generation model
    evaluate: Evaluate the model on test set
"""

import argparse
import sys
from pathlib import Path


def preprocess_command(args):
    """Run preprocessing on the dataset."""
    from src.preprocess_dataset import preprocess_dataset
    
    preprocess_dataset(
        dataset_root=args.dataset_root,
        output_dir=args.output_dir,
        skip_existing=args.skip_existing,
        max_samples=args.max_samples,
        generate_tsne=args.tsne,
        generate_pca=args.pca,
    )


def train_command(args):
    """Train the report generation model."""
    print("Training not yet implemented")
    # TODO: Implement training
    pass


def evaluate_command(args):
    """Evaluate the model."""
    print("Evaluation not yet implemented")
    # TODO: Implement evaluation
    pass


def main():
    parser = argparse.ArgumentParser(
        description='CBCT Report Generation Pipeline',
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    
    subparsers = parser.add_subparsers(dest='command', help='Command to run')
    
    # Preprocess command
    preprocess_parser = subparsers.add_parser(
        'preprocess',
        help='Preprocess dataset and extract regional tokens'
    )
    preprocess_parser.add_argument(
        '--dataset-root',
        type=str,
        default='data/ToothFairy3',
        help='Root directory of ToothFairy3 dataset'
    )
    preprocess_parser.add_argument(
        '--output-dir',
        type=str,
        default='data/preprocessed_tokens',
        help='Directory to save preprocessed tokens'
    )
    preprocess_parser.add_argument(
        '--skip-existing',
        action='store_true',
        default=True,
        help='Skip processing if output file exists'
    )
    preprocess_parser.add_argument(
        '--no-skip-existing',
        dest='skip_existing',
        action='store_false',
        help='Reprocess all files even if they exist'
    )
    preprocess_parser.add_argument(
        '--max-samples',
        type=int,
        default=None,
        help='Maximum number of samples to process (for testing)'
    )
    preprocess_parser.add_argument(
        '--tsne',
        action='store_true',
        help='Generate t-SNE visualizations for each patient'
    )
    preprocess_parser.add_argument(
        '--pca',
        action='store_true',
        help='Generate PCA RGB 3D volumes for slice visualization (saved to data/preprocessed_tsne/)'
    )
    
    # Train command
    train_parser = subparsers.add_parser(
        'train',
        help='Train the report generation model'
    )
    train_parser.add_argument(
        '--config',
        type=str,
        help='Path to training configuration file'
    )
    train_parser.add_argument(
        '--data-dir',
        type=str,
        default='data/preprocessed_tokens',
        help='Directory with preprocessed tokens'
    )
    train_parser.add_argument(
        '--output-dir',
        type=str,
        default='outputs/training',
        help='Directory to save model checkpoints'
    )
    
    # Evaluate command
    eval_parser = subparsers.add_parser(
        'evaluate',
        help='Evaluate the trained model'
    )
    eval_parser.add_argument(
        '--checkpoint',
        type=str,
        required=True,
        help='Path to model checkpoint'
    )
    eval_parser.add_argument(
        '--data-dir',
        type=str,
        default='data/preprocessed_tokens',
        help='Directory with preprocessed tokens'
    )
    eval_parser.add_argument(
        '--output-dir',
        type=str,
        default='outputs/evaluation',
        help='Directory to save evaluation results'
    )
    
    args = parser.parse_args()
    
    if args.command is None:
        parser.print_help()
        sys.exit(1)
    
    # Execute command
    if args.command == 'preprocess':
        preprocess_command(args)
    elif args.command == 'train':
        train_command(args)
    elif args.command == 'evaluate':
        evaluate_command(args)
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == '__main__':
    main()