234 lines
7.3 KiB
Python
234 lines
7.3 KiB
Python
|
"""
|
||
|
|
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()
|