main.py

import argparse
import os
import sys

def main():
    parser = argparse.ArgumentParser(description='Watermark detection and removal system')
    parser.add_argument('--mode', type=str, choices=['train', 'inference'], 
                        default='inference', help='Mode to run: train or inference')
    parser.add_argument('--image', type=str, default=None, 
                        help='Path to input image for inference mode')
    
    args = parser.parse_args()
    
    if args.mode == 'train':
        # 运行训练脚本
        print("Starting training mode...")
        os.system("python train.py")
    elif args.mode == 'inference':
        # 运行推理脚本
        print("Starting inference mode...")
        if args.image:
            os.system(f"python inference.py --image {args.image}")
        else:
            os.system("python inference.py")

if __name__ == "__main__":
    main()