init
This commit is contained in:
Nguyễn Phước Thành
234
id_card_processor_main.py
Normal file
234
id_card_processor_main.py
Normal file
@@ -0,0 +1,234 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user