""" Main script for ID Card Processing with YOLO Detection """ import argparse import sys from pathlib import Path from typing import Dict, Any import logging # Add src to path for imports sys.path.append(str(Path(__file__).parent / "src")) from src.model.yolo_detector import YOLODetector from src.model.id_card_processor import IDCardProcessor from src.utils import setup_logging def parse_arguments(): """Parse command line arguments""" parser = argparse.ArgumentParser(description="ID Card Processing with YOLO Detection") parser.add_argument( "--input-dir", type=str, required=True, help="Input directory containing ID card images" ) parser.add_argument( "--output-dir", type=str, default="data/processed_id_cards", help="Output directory for processed images" ) parser.add_argument( "--model-path", type=str, help="Path to custom YOLO model (.pt file)" ) parser.add_argument( "--confidence", type=float, default=0.5, help="Confidence threshold for YOLO detection" ) parser.add_argument( "--detect-only", action="store_true", help="Only detect and crop ID cards, skip preprocessing" ) parser.add_argument( "--preprocess-only", action="store_true", help="Skip detection, directly preprocess images" ) parser.add_argument( "--bg-removal", type=str, default="grabcut", choices=["grabcut", "threshold", "contour", "none"], help="Background removal method" ) parser.add_argument( "--target-size", type=str, default="800x600", help="Target size for normalization (width x height)" ) parser.add_argument( "--save-annotated", action="store_true", help="Save annotated images with bounding boxes" ) parser.add_argument( "--log-level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"], help="Logging level" ) return parser.parse_args() def parse_size(size_str: str) -> tuple: """Parse size string like '800x600' to tuple (800, 600)""" try: width, height = map(int, size_str.split('x')) return (width, height) except ValueError: print(f"Invalid size format: {size_str}. Expected format: widthxheight") sys.exit(1) def main(): """Main function""" args = parse_arguments() # Setup logging logging_config = {"level": args.log_level} logger = setup_logging(logging_config.get("level", "INFO")) logger.info("Starting ID Card Processing") # Parse paths input_dir = Path(args.input_dir) output_dir = Path(args.output_dir) # Check if input directory exists if not input_dir.exists(): logger.error(f"Input directory does not exist: {input_dir}") sys.exit(1) # Create output directory output_dir.mkdir(parents=True, exist_ok=True) # Parse target size target_size = parse_size(args.target_size) # Initialize YOLO detector logger.info("Initializing YOLO detector...") yolo_detector = YOLODetector( model_path=args.model_path, confidence=args.confidence ) # Initialize ID card processor logger.info("Initializing ID card processor...") id_processor = IDCardProcessor(yolo_detector) if args.detect_only: # Only detect and crop ID cards logger.info("Running YOLO detection only...") results = yolo_detector.batch_process( input_dir, output_dir / "cropped", save_annotated=args.save_annotated ) print("\n" + "="*50) print("YOLO DETECTION RESULTS") print("="*50) print(f"Total images: {results['total_images']}") print(f"Processed images: {results['processed_images']}") print(f"Total detections: {results['total_detections']}") print(f"Total cropped: {results['total_cropped']}") print(f"Output directory: {output_dir / 'cropped'}") print("="*50) elif args.preprocess_only: # Skip detection, directly preprocess logger.info("Running preprocessing only...") results = id_processor.batch_process_id_cards( input_dir, output_dir / "processed", detect_first=False, remove_bg=args.bg_removal != "none", enhance=True, normalize=True, target_size=target_size ) print("\n" + "="*50) print("PREPROCESSING RESULTS") print("="*50) print(f"Total images: {results['total_images']}") print(f"Processed images: {results['processed_images']}") print(f"Output directory: {output_dir / 'processed'}") print("="*50) else: # Full pipeline: detect + preprocess logger.info("Running full pipeline: detection + preprocessing...") # Step 1: Detect and crop ID cards logger.info("Step 1: Detecting and cropping ID cards...") detection_results = yolo_detector.batch_process( input_dir, output_dir / "cropped", save_annotated=args.save_annotated ) # Step 2: Preprocess cropped images cropped_dir = output_dir / "cropped" if cropped_dir.exists(): logger.info("Step 2: Preprocessing cropped ID cards...") preprocessing_results = id_processor.batch_process_id_cards( cropped_dir, output_dir / "processed", detect_first=False, remove_bg=args.bg_removal != "none", enhance=True, normalize=True, target_size=target_size ) else: logger.warning("No cropped images found, preprocessing original images") preprocessing_results = id_processor.batch_process_id_cards( input_dir, output_dir / "processed", detect_first=False, remove_bg=args.bg_removal != "none", enhance=True, normalize=True, target_size=target_size ) # Print summary print("\n" + "="*50) print("FULL PIPELINE RESULTS") print("="*50) print("DETECTION PHASE:") print(f" - Total images: {detection_results['total_images']}") print(f" - Processed images: {detection_results['processed_images']}") print(f" - Total detections: {detection_results['total_detections']}") print(f" - Total cropped: {detection_results['total_cropped']}") print("\nPREPROCESSING PHASE:") print(f" - Total images: {preprocessing_results['total_images']}") print(f" - Processed images: {preprocessing_results['processed_images']}") print(f"\nOutput directories:") print(f" - Cropped images: {output_dir / 'cropped'}") print(f" - Processed images: {output_dir / 'processed'}") if args.save_annotated: print(f" - Annotated images: {output_dir / 'cropped'}") print("="*50) logger.info("ID Card Processing completed successfully") if __name__ == "__main__": main()