133 lines
4.0 KiB
Python
133 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Simple ID Card Cropper using Roboflow API
|
|
Input: folder containing images
|
|
Output: folder with cropped ID cards
|
|
"""
|
|
import sys
|
|
import yaml
|
|
from pathlib import Path
|
|
import logging
|
|
import argparse
|
|
|
|
# Add src to path
|
|
sys.path.append(str(Path(__file__).parent / "src"))
|
|
|
|
from model.roboflow_id_detector import RoboflowIDDetector
|
|
|
|
def setup_logging():
|
|
"""Setup basic logging"""
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(levelname)s - %(message)s'
|
|
)
|
|
|
|
def crop_id_cards(input_folder: str, output_folder: str, api_key: str = "Pkz4puRA0Cy3xMOuNoNr"):
|
|
"""
|
|
Crop ID cards from all images in input folder
|
|
|
|
Args:
|
|
input_folder: Path to input folder containing images
|
|
output_folder: Path to output folder for cropped ID cards
|
|
api_key: Roboflow API key
|
|
"""
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Convert to Path objects
|
|
input_path = Path(input_folder)
|
|
output_path = Path(output_folder)
|
|
|
|
# Check if input folder exists
|
|
if not input_path.exists():
|
|
logger.error(f"Input folder not found: {input_folder}")
|
|
return False
|
|
|
|
# Create output folder
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Initialize detector
|
|
detector = RoboflowIDDetector(
|
|
api_key=api_key,
|
|
model_id="french-card-id-detect",
|
|
version=3,
|
|
confidence=0.5
|
|
)
|
|
|
|
# Get all image files
|
|
image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
|
|
image_files = []
|
|
|
|
for file_path in input_path.rglob('*'):
|
|
if file_path.is_file() and file_path.suffix.lower() in image_extensions:
|
|
image_files.append(file_path)
|
|
|
|
if not image_files:
|
|
logger.error(f"No images found in {input_folder}")
|
|
return False
|
|
|
|
logger.info(f"Found {len(image_files)} images to process")
|
|
|
|
# Process each image
|
|
total_cropped = 0
|
|
|
|
for i, image_path in enumerate(image_files, 1):
|
|
logger.info(f"Processing {i}/{len(image_files)}: {image_path.name}")
|
|
|
|
# Detect ID cards
|
|
detections = detector.detect_id_cards(image_path)
|
|
|
|
if not detections:
|
|
logger.warning(f"No ID cards detected in {image_path.name}")
|
|
continue
|
|
|
|
# Crop each detected ID card
|
|
for j, detection in enumerate(detections):
|
|
bbox = detection['bbox']
|
|
|
|
# Create output filename
|
|
stem = image_path.stem
|
|
suffix = f"_card_{j+1}.jpg"
|
|
output_file = output_path / f"{stem}{suffix}"
|
|
|
|
# Crop ID card
|
|
cropped = detector.crop_id_card(image_path, bbox, output_file)
|
|
|
|
if cropped is not None:
|
|
total_cropped += 1
|
|
logger.info(f" ✓ Cropped card {j+1} to {output_file.name}")
|
|
|
|
# Add delay between requests
|
|
if i < len(image_files):
|
|
import time
|
|
time.sleep(1.0)
|
|
|
|
logger.info(f"Processing completed! Total ID cards cropped: {total_cropped}")
|
|
return True
|
|
|
|
def main():
|
|
"""Main function"""
|
|
parser = argparse.ArgumentParser(description='Crop ID cards from images using Roboflow API')
|
|
parser.add_argument('input_folder', help='Input folder containing images')
|
|
parser.add_argument('output_folder', help='Output folder for cropped ID cards')
|
|
parser.add_argument('--api-key', default="Pkz4puRA0Cy3xMOuNoNr",
|
|
help='Roboflow API key (default: demo key)')
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Setup logging
|
|
setup_logging()
|
|
|
|
# Process images
|
|
success = crop_id_cards(args.input_folder, args.output_folder, args.api_key)
|
|
|
|
if success:
|
|
print(f"\n✓ Successfully processed images from '{args.input_folder}'")
|
|
print(f"✓ Cropped ID cards saved to '{args.output_folder}'")
|
|
else:
|
|
print(f"\n✗ Failed to process images")
|
|
return 1
|
|
|
|
return 0
|
|
|
|
if __name__ == "__main__":
|
|
exit(main()) |