From 24060e4ce7233a97d636857f1cd14128c91d4aab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nguy=E1=BB=85n=20Ph=C6=B0=E1=BB=9Bc=20Th=C3=A0nh?= <93478665+Zeres-Engel@users.noreply.github.com> Date: Tue, 5 Aug 2025 19:09:55 +0700 Subject: [PATCH] init --- .gitignore | 17 + README_ID_Card_Processing.md | 183 ++++++++++ config/config.yaml | 48 +++ data_augmentation.log | 222 ++++++++++++ id_card_processor_main.py | 234 ++++++++++++ main.py | 276 ++++++++++++++ src/__init__.py | 23 ++ src/__pycache__/__init__.cpython-313.pyc | Bin 0 -> 589 bytes src/__pycache__/__init__.cpython-39.pyc | Bin 0 -> 559 bytes src/__pycache__/config_manager.cpython-39.pyc | Bin 0 -> 6334 bytes .../data_augmentation.cpython-39.pyc | Bin 0 -> 5323 bytes .../image_processor.cpython-39.pyc | Bin 0 -> 5218 bytes src/__pycache__/utils.cpython-313.pyc | Bin 0 -> 5478 bytes src/__pycache__/utils.cpython-39.pyc | Bin 0 -> 3426 bytes src/config.py | 40 ++ src/config_manager.py | 175 +++++++++ src/data_augmentation.py | 161 ++++++++ src/image_processor.py | 174 +++++++++ src/model/__init__.py | 8 + src/model/__pycache__/__init__.cpython-39.pyc | Bin 0 -> 352 bytes .../id_card_processor.cpython-39.pyc | Bin 0 -> 8960 bytes .../__pycache__/yolo_detector.cpython-39.pyc | Bin 0 -> 7104 bytes src/model/id_card_processor.py | 343 ++++++++++++++++++ src/model/yolo_detector.py | 266 ++++++++++++++ src/utils.py | 98 +++++ 25 files changed, 2268 insertions(+) create mode 100644 .gitignore create mode 100644 README_ID_Card_Processing.md create mode 100644 config/config.yaml create mode 100644 data_augmentation.log create mode 100644 id_card_processor_main.py create mode 100644 main.py create mode 100644 src/__init__.py create mode 100644 src/__pycache__/__init__.cpython-313.pyc create mode 100644 src/__pycache__/__init__.cpython-39.pyc create mode 100644 src/__pycache__/config_manager.cpython-39.pyc create mode 100644 src/__pycache__/data_augmentation.cpython-39.pyc create mode 100644 src/__pycache__/image_processor.cpython-39.pyc create mode 100644 src/__pycache__/utils.cpython-313.pyc create mode 100644 src/__pycache__/utils.cpython-39.pyc create mode 100644 src/config.py create mode 100644 src/config_manager.py create mode 100644 src/data_augmentation.py create mode 100644 src/image_processor.py create mode 100644 src/model/__init__.py create mode 100644 src/model/__pycache__/__init__.cpython-39.pyc create mode 100644 src/model/__pycache__/id_card_processor.cpython-39.pyc create mode 100644 src/model/__pycache__/yolo_detector.cpython-39.pyc create mode 100644 src/model/id_card_processor.py create mode 100644 src/model/yolo_detector.py create mode 100644 src/utils.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..70d391f --- /dev/null +++ b/.gitignore @@ -0,0 +1,17 @@ +*.png +*.json +*.jpg +*.zip +*.rar +*.pdf +*.docx +*.doc +*.xls +*.xlsx +*.ppt +*.pptx +*.txt +*.csv +*.json +*.pt +*.ipynb diff --git a/README_ID_Card_Processing.md b/README_ID_Card_Processing.md new file mode 100644 index 0000000..3e922a4 --- /dev/null +++ b/README_ID_Card_Processing.md @@ -0,0 +1,183 @@ +# ID Card Processing with YOLO Detection + +Hệ thống xử lý ID cards sử dụng YOLO để detect và crop, kết hợp với các phương pháp tiền xử lý để clean background và enhance chất lượng ảnh. + +## Tính năng chính + +- **YOLO Detection**: Detect và crop ID cards từ ảnh gốc +- **Background Removal**: 3 phương pháp loại bỏ background (GrabCut, Threshold, Contour) +- **Image Enhancement**: Cải thiện chất lượng ảnh cho OCR +- **Batch Processing**: Xử lý hàng loạt ảnh +- **Flexible Pipeline**: Có thể chạy từng bước riêng biệt + +## Cài đặt + +1. Cài đặt dependencies: +```bash +pip install -r requirements.txt +``` + +2. Cấu trúc thư mục: +``` +OCR/ +├── src/ +│ ├── model/ +│ │ ├── __init__.py +│ │ ├── yolo_detector.py +│ │ └── id_card_processor.py +│ └── ... +├── data/ +│ ├── IDcards/ # Thư mục chứa ảnh ID cards gốc +│ └── processed_id_cards/ # Thư mục output +├── id_card_processor_main.py +└── requirements.txt +``` + +## Sử dụng + +### 1. Full Pipeline (Detect + Preprocess) + +```bash +python id_card_processor_main.py \ + --input-dir "data/IDcards" \ + --output-dir "data/processed_id_cards" \ + --confidence 0.5 \ + --bg-removal grabcut \ + --target-size 800x600 \ + --save-annotated +``` + +### 2. Chỉ Detect và Crop + +```bash +python id_card_processor_main.py \ + --input-dir "data/IDcards" \ + --output-dir "data/processed_id_cards" \ + --detect-only \ + --save-annotated +``` + +### 3. Chỉ Preprocess (bỏ qua detection) + +```bash +python id_card_processor_main.py \ + --input-dir "data/IDcards" \ + --output-dir "data/processed_id_cards" \ + --preprocess-only \ + --bg-removal threshold \ + --target-size 800x600 +``` + +## Các tham số + +### Detection Parameters +- `--model-path`: Đường dẫn đến custom YOLO model (.pt file) +- `--confidence`: Ngưỡng confidence cho detection (default: 0.5) + +### Preprocessing Parameters +- `--bg-removal`: Phương pháp loại bỏ background + - `grabcut`: Sử dụng GrabCut algorithm (recommended) + - `threshold`: Sử dụng thresholding + - `contour`: Sử dụng contour detection + - `none`: Không loại bỏ background +- `--target-size`: Kích thước chuẩn hóa (width x height) + +### Output Options +- `--save-annotated`: Lưu ảnh với bounding boxes +- `--detect-only`: Chỉ chạy detection +- `--preprocess-only`: Chỉ chạy preprocessing + +## Output Structure + +``` +data/processed_id_cards/ +├── cropped/ # Ảnh đã được crop từ YOLO +│ ├── image1_card_1.jpg +│ ├── image1_card_2.jpg +│ └── ... +├── processed/ # Ảnh đã được preprocess +│ ├── image1_card_1_processed.jpg +│ ├── image1_card_2_processed.jpg +│ └── ... +└── annotated/ # Ảnh với bounding boxes (nếu có) + ├── image1_annotated.jpg + └── ... +``` + +## Ví dụ sử dụng + +### Ví dụ 1: Xử lý toàn bộ dataset +```bash +# Xử lý tất cả ảnh trong thư mục IDcards +python id_card_processor_main.py \ + --input-dir "data/IDcards" \ + --output-dir "data/processed_id_cards" \ + --confidence 0.6 \ + --bg-removal grabcut \ + --target-size 1024x768 \ + --save-annotated +``` + +### Ví dụ 2: Test với một vài ảnh +```bash +# Tạo thư mục test với một vài ảnh +mkdir -p data/test_images +# Copy một vài ảnh vào test_images + +# Chạy detection +python id_card_processor_main.py \ + --input-dir "data/test_images" \ + --output-dir "data/test_output" \ + --detect-only \ + --save-annotated +``` + +### Ví dụ 3: Sử dụng custom model +```bash +# Nếu bạn có custom YOLO model đã train +python id_card_processor_main.py \ + --input-dir "data/IDcards" \ + --output-dir "data/processed_id_cards" \ + --model-path "models/custom_id_card_model.pt" \ + --confidence 0.7 +``` + +## Lưu ý + +1. **YOLO Model**: Mặc định sử dụng YOLOv8n pre-trained. Nếu có custom model tốt hơn, hãy sử dụng `--model-path` + +2. **Background Removal**: + - `grabcut`: Tốt nhất cho ID cards có background phức tạp + - `threshold`: Nhanh, phù hợp với background đơn giản + - `contour`: Phù hợp với ID cards có viền rõ ràng + +3. **Performance**: + - Sử dụng GPU nếu có thể để tăng tốc độ detection + - Có thể điều chỉnh `--confidence` để cân bằng giữa precision và recall + +4. **Memory**: Với dataset lớn, có thể cần tăng memory hoặc xử lý theo batch nhỏ hơn + +## Troubleshooting + +### Lỗi thường gặp + +1. **No detections found**: + - Giảm `--confidence` xuống 0.3-0.4 + - Kiểm tra chất lượng ảnh input + +2. **Memory error**: + - Giảm batch size hoặc xử lý từng ảnh một + - Sử dụng CPU thay vì GPU + +3. **Poor background removal**: + - Thử các phương pháp khác nhau: `grabcut`, `threshold`, `contour` + - Điều chỉnh parameters trong code + +### Debug mode + +```bash +python id_card_processor_main.py \ + --input-dir "data/IDcards" \ + --output-dir "data/processed_id_cards" \ + --log-level DEBUG +``` \ No newline at end of file diff --git a/config/config.yaml b/config/config.yaml new file mode 100644 index 0000000..96b942f --- /dev/null +++ b/config/config.yaml @@ -0,0 +1,48 @@ +# Data Augmentation Configuration +# Main configuration file for image data augmentation + +# Paths configuration +paths: + input_dir: "data/Archive" + output_dir: "out" + log_file: "logs/data_augmentation.log" + +# Data augmentation parameters - ONLY ROTATION +augmentation: + # Geometric transformations - ONLY ROTATION + rotation: + enabled: true + angles: [30, 60, 120, 150, 180, 210, 240, 300, 330] # Specific rotation angles + probability: 1.0 # Always apply rotation + +# Processing configuration +processing: + target_size: [224, 224] # [width, height] + batch_size: 32 + num_augmentations: 3 # number of augmented versions per image + save_format: "jpg" + quality: 95 + +# Supported image formats +supported_formats: + - ".jpg" + - ".jpeg" + - ".png" + - ".bmp" + - ".tiff" + +# Logging configuration +logging: + level: "INFO" # DEBUG, INFO, WARNING, ERROR + format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + handlers: + - type: "file" + filename: "logs/data_augmentation.log" + - type: "console" + +# Performance settings +performance: + num_workers: 4 + prefetch_factor: 2 + pin_memory: true + use_gpu: false \ No newline at end of file diff --git a/data_augmentation.log b/data_augmentation.log new file mode 100644 index 0000000..2405a6d --- /dev/null +++ b/data_augmentation.log @@ -0,0 +1,222 @@ +2025-08-05 18:53:06,981 - src.model.yolo_detector - INFO - Using pre-trained YOLOv8n model +2025-08-05 18:53:07,004 - src.model.yolo_detector - INFO - Using device: cuda +2025-08-05 18:53:07,038 - src.model.yolo_detector - INFO - Using pre-trained YOLOv8n model +2025-08-05 18:53:07,038 - src.model.yolo_detector - INFO - Using device: cuda +2025-08-05 18:53:07,361 - src.model.yolo_detector - INFO - Using pre-trained YOLOv8n model +2025-08-05 18:53:07,362 - src.model.yolo_detector - INFO - Using device: cuda +2025-08-05 18:53:07,363 - src.model.id_card_processor - INFO - Detecting and cropping ID cards... +2025-08-05 18:53:07,363 - src.model.yolo_detector - ERROR - No images found in data\IDcards +2025-08-05 18:53:07,364 - src.model.id_card_processor - INFO - Processing cropped ID cards... +2025-08-05 18:53:07,364 - src.model.id_card_processor - ERROR - No images found in data\test_output\cropped +2025-08-05 19:04:14,903 - src.model.yolo_detector - INFO - Using pre-trained YOLOv8n model +2025-08-05 19:04:14,995 - src.model.yolo_detector - INFO - Using device: cuda +2025-08-05 19:04:14,996 - src.model.id_card_processor - INFO - Detecting and cropping ID cards... +2025-08-05 19:04:14,997 - src.model.yolo_detector - INFO - Processing 29 images from data\IDcards and subdirectories +2025-08-05 19:04:14,998 - src.model.yolo_detector - INFO - Processing 1/29: im10.png +2025-08-05 19:04:19,785 - src.model.yolo_detector - INFO - Found 1 detections in im10.png +2025-08-05 19:04:19,813 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im10_card_1.jpg +2025-08-05 19:04:19,813 - src.model.yolo_detector - INFO - Processed im10.png: 1 cards cropped +2025-08-05 19:04:19,814 - src.model.yolo_detector - INFO - Processing 2/29: im11.png +2025-08-05 19:04:19,926 - src.model.yolo_detector - INFO - Found 2 detections in im11.png +2025-08-05 19:04:19,937 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im11_card_1.jpg +2025-08-05 19:04:19,946 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im11_card_2.jpg +2025-08-05 19:04:19,946 - src.model.yolo_detector - INFO - Processed im11.png: 2 cards cropped +2025-08-05 19:04:19,946 - src.model.yolo_detector - INFO - Processing 3/29: im12.png +2025-08-05 19:04:20,056 - src.model.yolo_detector - INFO - Found 2 detections in im12.png +2025-08-05 19:04:20,069 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im12_card_1.jpg +2025-08-05 19:04:20,082 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im12_card_2.jpg +2025-08-05 19:04:20,083 - src.model.yolo_detector - INFO - Processed im12.png: 2 cards cropped +2025-08-05 19:04:20,083 - src.model.yolo_detector - INFO - Processing 4/29: im13.png +2025-08-05 19:04:20,116 - src.model.yolo_detector - INFO - Found 0 detections in im13.png +2025-08-05 19:04:20,117 - src.model.yolo_detector - WARNING - No ID cards detected in im13.png +2025-08-05 19:04:20,117 - src.model.yolo_detector - INFO - Processing 5/29: im14.png +2025-08-05 19:04:20,156 - src.model.yolo_detector - INFO - Found 1 detections in im14.png +2025-08-05 19:04:20,172 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im14_card_1.jpg +2025-08-05 19:04:20,173 - src.model.yolo_detector - INFO - Processed im14.png: 1 cards cropped +2025-08-05 19:04:20,174 - src.model.yolo_detector - INFO - Processing 6/29: im15.png +2025-08-05 19:04:20,208 - src.model.yolo_detector - INFO - Found 1 detections in im15.png +2025-08-05 19:04:20,222 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im15_card_1.jpg +2025-08-05 19:04:20,222 - src.model.yolo_detector - INFO - Processed im15.png: 1 cards cropped +2025-08-05 19:04:20,223 - src.model.yolo_detector - INFO - Processing 7/29: im1_.png +2025-08-05 19:04:20,466 - src.model.yolo_detector - INFO - Found 0 detections in im1_.png +2025-08-05 19:04:20,466 - src.model.yolo_detector - WARNING - No ID cards detected in im1_.png +2025-08-05 19:04:20,466 - src.model.yolo_detector - INFO - Processing 8/29: im2.png +2025-08-05 19:04:20,534 - src.model.yolo_detector - INFO - Found 2 detections in im2.png +2025-08-05 19:04:20,564 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im2_card_1.jpg +2025-08-05 19:04:20,594 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im2_card_2.jpg +2025-08-05 19:04:20,594 - src.model.yolo_detector - INFO - Processed im2.png: 2 cards cropped +2025-08-05 19:04:20,595 - src.model.yolo_detector - INFO - Processing 9/29: im3.png +2025-08-05 19:04:20,648 - src.model.yolo_detector - INFO - Found 1 detections in im3.png +2025-08-05 19:04:20,671 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im3_card_1.jpg +2025-08-05 19:04:20,671 - src.model.yolo_detector - INFO - Processed im3.png: 1 cards cropped +2025-08-05 19:04:20,672 - src.model.yolo_detector - INFO - Processing 10/29: im4.png +2025-08-05 19:04:20,724 - src.model.yolo_detector - INFO - Found 1 detections in im4.png +2025-08-05 19:04:20,753 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im4_card_1.jpg +2025-08-05 19:04:20,754 - src.model.yolo_detector - INFO - Processed im4.png: 1 cards cropped +2025-08-05 19:04:20,754 - src.model.yolo_detector - INFO - Processing 11/29: im5.png +2025-08-05 19:04:20,798 - src.model.yolo_detector - INFO - Found 2 detections in im5.png +2025-08-05 19:04:20,816 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im5_card_1.jpg +2025-08-05 19:04:20,835 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im5_card_2.jpg +2025-08-05 19:04:20,836 - src.model.yolo_detector - INFO - Processed im5.png: 2 cards cropped +2025-08-05 19:04:20,837 - src.model.yolo_detector - INFO - Processing 12/29: im6.png +2025-08-05 19:04:20,994 - src.model.yolo_detector - INFO - Found 2 detections in im6.png +2025-08-05 19:04:21,052 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im6_card_1.jpg +2025-08-05 19:04:21,118 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im6_card_2.jpg +2025-08-05 19:04:21,119 - src.model.yolo_detector - INFO - Processed im6.png: 2 cards cropped +2025-08-05 19:04:21,120 - src.model.yolo_detector - INFO - Processing 13/29: im7.png +2025-08-05 19:04:21,159 - src.model.yolo_detector - INFO - Found 3 detections in im7.png +2025-08-05 19:04:21,168 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im7_card_1.jpg +2025-08-05 19:04:21,176 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im7_card_2.jpg +2025-08-05 19:04:21,184 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im7_card_3.jpg +2025-08-05 19:04:21,184 - src.model.yolo_detector - INFO - Processed im7.png: 3 cards cropped +2025-08-05 19:04:21,185 - src.model.yolo_detector - INFO - Processing 14/29: im8.png +2025-08-05 19:04:21,353 - src.model.yolo_detector - INFO - Found 2 detections in im8.png +2025-08-05 19:04:21,387 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im8_card_1.jpg +2025-08-05 19:04:21,423 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im8_card_2.jpg +2025-08-05 19:04:21,424 - src.model.yolo_detector - INFO - Processed im8.png: 2 cards cropped +2025-08-05 19:04:21,425 - src.model.yolo_detector - INFO - Processing 15/29: im9.png +2025-08-05 19:04:21,522 - src.model.yolo_detector - INFO - Found 1 detections in im9.png +2025-08-05 19:04:21,532 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\Archive\im9_card_1.jpg +2025-08-05 19:04:21,532 - src.model.yolo_detector - INFO - Processed im9.png: 1 cards cropped +2025-08-05 19:04:21,532 - src.model.yolo_detector - INFO - Processing 16/29: im10.png +2025-08-05 19:04:21,585 - src.model.yolo_detector - INFO - Found 3 detections in im10.png +2025-08-05 19:04:21,601 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im10_card_1.jpg +2025-08-05 19:04:21,618 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im10_card_2.jpg +2025-08-05 19:04:21,636 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im10_card_3.jpg +2025-08-05 19:04:21,636 - src.model.yolo_detector - INFO - Processed im10.png: 3 cards cropped +2025-08-05 19:04:21,638 - src.model.yolo_detector - INFO - Processing 17/29: im11.png +2025-08-05 19:04:21,679 - src.model.yolo_detector - INFO - Found 2 detections in im11.png +2025-08-05 19:04:21,696 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im11_card_1.jpg +2025-08-05 19:04:21,712 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im11_card_2.jpg +2025-08-05 19:04:21,713 - src.model.yolo_detector - INFO - Processed im11.png: 2 cards cropped +2025-08-05 19:04:21,713 - src.model.yolo_detector - INFO - Processing 18/29: im12.png +2025-08-05 19:04:21,755 - src.model.yolo_detector - INFO - Found 0 detections in im12.png +2025-08-05 19:04:21,756 - src.model.yolo_detector - WARNING - No ID cards detected in im12.png +2025-08-05 19:04:21,756 - src.model.yolo_detector - INFO - Processing 19/29: im13.png +2025-08-05 19:04:21,793 - src.model.yolo_detector - INFO - Found 1 detections in im13.png +2025-08-05 19:04:21,806 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im13_card_1.jpg +2025-08-05 19:04:21,806 - src.model.yolo_detector - INFO - Processed im13.png: 1 cards cropped +2025-08-05 19:04:21,806 - src.model.yolo_detector - INFO - Processing 20/29: im14.png +2025-08-05 19:04:21,846 - src.model.yolo_detector - INFO - Found 2 detections in im14.png +2025-08-05 19:04:21,862 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im14_card_1.jpg +2025-08-05 19:04:21,877 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im14_card_2.jpg +2025-08-05 19:04:21,877 - src.model.yolo_detector - INFO - Processed im14.png: 2 cards cropped +2025-08-05 19:04:21,878 - src.model.yolo_detector - INFO - Processing 21/29: im15.png +2025-08-05 19:04:21,914 - src.model.yolo_detector - INFO - Found 0 detections in im15.png +2025-08-05 19:04:21,914 - src.model.yolo_detector - WARNING - No ID cards detected in im15.png +2025-08-05 19:04:21,914 - src.model.yolo_detector - INFO - Processing 22/29: im1_.png +2025-08-05 19:04:21,959 - src.model.yolo_detector - INFO - Found 3 detections in im1_.png +2025-08-05 19:04:21,971 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im1__card_1.jpg +2025-08-05 19:04:21,983 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im1__card_2.jpg +2025-08-05 19:04:21,996 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im1__card_3.jpg +2025-08-05 19:04:21,997 - src.model.yolo_detector - INFO - Processed im1_.png: 3 cards cropped +2025-08-05 19:04:21,997 - src.model.yolo_detector - INFO - Processing 23/29: im2.png +2025-08-05 19:04:22,101 - src.model.yolo_detector - INFO - Found 1 detections in im2.png +2025-08-05 19:04:22,174 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im2_card_1.jpg +2025-08-05 19:04:22,174 - src.model.yolo_detector - INFO - Processed im2.png: 1 cards cropped +2025-08-05 19:04:22,176 - src.model.yolo_detector - INFO - Processing 24/29: im3.png +2025-08-05 19:04:22,220 - src.model.yolo_detector - INFO - Found 2 detections in im3.png +2025-08-05 19:04:22,235 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im3_card_1.jpg +2025-08-05 19:04:22,251 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im3_card_2.jpg +2025-08-05 19:04:22,252 - src.model.yolo_detector - INFO - Processed im3.png: 2 cards cropped +2025-08-05 19:04:22,252 - src.model.yolo_detector - INFO - Processing 25/29: im5.png +2025-08-05 19:04:22,307 - src.model.yolo_detector - INFO - Found 1 detections in im5.png +2025-08-05 19:04:22,316 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im5_card_1.jpg +2025-08-05 19:04:22,316 - src.model.yolo_detector - INFO - Processed im5.png: 1 cards cropped +2025-08-05 19:04:22,317 - src.model.yolo_detector - INFO - Processing 26/29: im6.png +2025-08-05 19:04:22,375 - src.model.yolo_detector - INFO - Found 2 detections in im6.png +2025-08-05 19:04:22,387 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im6_card_1.jpg +2025-08-05 19:04:22,397 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im6_card_2.jpg +2025-08-05 19:04:22,398 - src.model.yolo_detector - INFO - Processed im6.png: 2 cards cropped +2025-08-05 19:04:22,399 - src.model.yolo_detector - INFO - Processing 27/29: im7.png +2025-08-05 19:04:22,441 - src.model.yolo_detector - INFO - Found 1 detections in im7.png +2025-08-05 19:04:22,458 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im7_card_1.jpg +2025-08-05 19:04:22,459 - src.model.yolo_detector - INFO - Processed im7.png: 1 cards cropped +2025-08-05 19:04:22,460 - src.model.yolo_detector - INFO - Processing 28/29: im8.png +2025-08-05 19:04:22,492 - src.model.yolo_detector - INFO - Found 2 detections in im8.png +2025-08-05 19:04:22,502 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im8_card_1.jpg +2025-08-05 19:04:22,509 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im8_card_2.jpg +2025-08-05 19:04:22,510 - src.model.yolo_detector - INFO - Processed im8.png: 2 cards cropped +2025-08-05 19:04:22,510 - src.model.yolo_detector - INFO - Processing 29/29: im9.png +2025-08-05 19:04:22,540 - src.model.yolo_detector - INFO - Found 1 detections in im9.png +2025-08-05 19:04:22,546 - src.model.yolo_detector - INFO - Saved cropped image to data\processed_id_cards\cropped\titre-sejour-fr\im9_card_1.jpg +2025-08-05 19:04:22,546 - src.model.yolo_detector - INFO - Processed im9.png: 1 cards cropped +2025-08-05 19:04:22,546 - src.model.yolo_detector - INFO - Batch processing completed: +2025-08-05 19:04:22,548 - src.model.yolo_detector - INFO - - Total images: 29 +2025-08-05 19:04:22,548 - src.model.yolo_detector - INFO - - Processed: 25 +2025-08-05 19:04:22,548 - src.model.yolo_detector - INFO - - Total detections: 42 +2025-08-05 19:04:22,549 - src.model.yolo_detector - INFO - - Total cropped: 42 +2025-08-05 19:04:22,549 - src.model.id_card_processor - INFO - Processing cropped ID cards... +2025-08-05 19:04:22,552 - src.model.id_card_processor - INFO - Processing 42 images from data\processed_id_cards\cropped and subdirectories +2025-08-05 19:04:22,552 - src.model.id_card_processor - INFO - Processing 1/42: im10_card_1.jpg +2025-08-05 19:04:22,564 - src.model.id_card_processor - INFO - Removing background from im10_card_1.jpg +2025-08-05 19:04:22,877 - src.model.id_card_processor - INFO - Enhancing im10_card_1.jpg +2025-08-05 19:04:23,016 - src.model.id_card_processor - INFO - Normalizing im10_card_1.jpg +2025-08-05 19:04:23,023 - src.model.id_card_processor - INFO - Processed im10_card_1.jpg +2025-08-05 19:04:23,023 - src.model.id_card_processor - INFO - Processing 2/42: im11_card_1.jpg +2025-08-05 19:04:23,034 - src.model.id_card_processor - INFO - Removing background from im11_card_1.jpg +2025-08-05 19:04:23,264 - src.model.id_card_processor - INFO - Enhancing im11_card_1.jpg +2025-08-05 19:04:23,265 - src.model.id_card_processor - INFO - Normalizing im11_card_1.jpg +2025-08-05 19:04:23,270 - src.model.id_card_processor - INFO - Processed im11_card_1.jpg +2025-08-05 19:04:23,271 - src.model.id_card_processor - INFO - Processing 3/42: im11_card_2.jpg +2025-08-05 19:04:23,282 - src.model.id_card_processor - INFO - Removing background from im11_card_2.jpg +2025-08-05 19:04:23,312 - src.model.id_card_processor - INFO - Enhancing im11_card_2.jpg +2025-08-05 19:04:23,313 - src.model.id_card_processor - INFO - Normalizing im11_card_2.jpg +2025-08-05 19:04:23,316 - src.model.id_card_processor - INFO - Processed im11_card_2.jpg +2025-08-05 19:04:23,316 - src.model.id_card_processor - INFO - Processing 4/42: im12_card_1.jpg +2025-08-05 19:04:23,328 - src.model.id_card_processor - INFO - Removing background from im12_card_1.jpg +2025-08-05 19:04:23,670 - src.model.id_card_processor - INFO - Enhancing im12_card_1.jpg +2025-08-05 19:04:23,671 - src.model.id_card_processor - INFO - Normalizing im12_card_1.jpg +2025-08-05 19:04:23,675 - src.model.id_card_processor - INFO - Processed im12_card_1.jpg +2025-08-05 19:04:23,676 - src.model.id_card_processor - INFO - Processing 5/42: im12_card_2.jpg +2025-08-05 19:04:23,686 - src.model.id_card_processor - INFO - Removing background from im12_card_2.jpg +2025-08-05 19:04:29,279 - src.model.id_card_processor - INFO - Enhancing im12_card_2.jpg +2025-08-05 19:04:29,284 - src.model.id_card_processor - INFO - Normalizing im12_card_2.jpg +2025-08-05 19:04:29,289 - src.model.id_card_processor - INFO - Processed im12_card_2.jpg +2025-08-05 19:04:29,290 - src.model.id_card_processor - INFO - Processing 6/42: im14_card_1.jpg +2025-08-05 19:04:29,301 - src.model.id_card_processor - INFO - Removing background from im14_card_1.jpg +2025-08-05 19:04:29,774 - src.model.id_card_processor - INFO - Enhancing im14_card_1.jpg +2025-08-05 19:04:29,775 - src.model.id_card_processor - INFO - Normalizing im14_card_1.jpg +2025-08-05 19:04:29,779 - src.model.id_card_processor - INFO - Processed im14_card_1.jpg +2025-08-05 19:04:29,780 - src.model.id_card_processor - INFO - Processing 7/42: im15_card_1.jpg +2025-08-05 19:04:29,791 - src.model.id_card_processor - INFO - Removing background from im15_card_1.jpg +2025-08-05 19:04:30,009 - src.model.id_card_processor - INFO - Enhancing im15_card_1.jpg +2025-08-05 19:04:30,010 - src.model.id_card_processor - INFO - Normalizing im15_card_1.jpg +2025-08-05 19:04:30,015 - src.model.id_card_processor - INFO - Processed im15_card_1.jpg +2025-08-05 19:04:30,015 - src.model.id_card_processor - INFO - Processing 8/42: im2_card_1.jpg +2025-08-05 19:04:30,017 - src.model.id_card_processor - INFO - Removing background from im2_card_1.jpg +2025-08-05 19:04:31,861 - src.model.id_card_processor - INFO - Enhancing im2_card_1.jpg +2025-08-05 19:04:31,863 - src.model.id_card_processor - INFO - Normalizing im2_card_1.jpg +2025-08-05 19:04:31,869 - src.model.id_card_processor - INFO - Processed im2_card_1.jpg +2025-08-05 19:04:31,869 - src.model.id_card_processor - INFO - Processing 9/42: im2_card_2.jpg +2025-08-05 19:04:31,884 - src.model.id_card_processor - INFO - Removing background from im2_card_2.jpg +2025-08-05 19:04:38,985 - src.model.id_card_processor - INFO - Enhancing im2_card_2.jpg +2025-08-05 19:04:38,996 - src.model.id_card_processor - INFO - Normalizing im2_card_2.jpg +2025-08-05 19:04:39,007 - src.model.id_card_processor - INFO - Processed im2_card_2.jpg +2025-08-05 19:04:39,008 - src.model.id_card_processor - INFO - Processing 10/42: im3_card_1.jpg +2025-08-05 19:04:39,009 - src.model.id_card_processor - INFO - Removing background from im3_card_1.jpg +2025-08-05 19:04:39,177 - src.model.id_card_processor - INFO - Enhancing im3_card_1.jpg +2025-08-05 19:04:39,178 - src.model.id_card_processor - INFO - Normalizing im3_card_1.jpg +2025-08-05 19:04:39,182 - src.model.id_card_processor - INFO - Processed im3_card_1.jpg +2025-08-05 19:04:39,182 - src.model.id_card_processor - INFO - Processing 11/42: im4_card_1.jpg +2025-08-05 19:04:39,184 - src.model.id_card_processor - INFO - Removing background from im4_card_1.jpg +2025-08-05 19:04:39,374 - src.model.id_card_processor - INFO - Enhancing im4_card_1.jpg +2025-08-05 19:04:39,375 - src.model.id_card_processor - INFO - Normalizing im4_card_1.jpg +2025-08-05 19:04:39,379 - src.model.id_card_processor - INFO - Processed im4_card_1.jpg +2025-08-05 19:04:39,379 - src.model.id_card_processor - INFO - Processing 12/42: im5_card_1.jpg +2025-08-05 19:04:39,389 - src.model.id_card_processor - INFO - Removing background from im5_card_1.jpg +2025-08-05 19:04:39,842 - src.model.id_card_processor - INFO - Enhancing im5_card_1.jpg +2025-08-05 19:04:39,843 - src.model.id_card_processor - INFO - Normalizing im5_card_1.jpg +2025-08-05 19:04:39,846 - src.model.id_card_processor - INFO - Processed im5_card_1.jpg +2025-08-05 19:04:39,846 - src.model.id_card_processor - INFO - Processing 13/42: im5_card_2.jpg +2025-08-05 19:04:39,859 - src.model.id_card_processor - INFO - Removing background from im5_card_2.jpg +2025-08-05 19:04:42,430 - src.model.id_card_processor - INFO - Enhancing im5_card_2.jpg +2025-08-05 19:04:42,434 - src.model.id_card_processor - INFO - Normalizing im5_card_2.jpg +2025-08-05 19:04:42,438 - src.model.id_card_processor - INFO - Processed im5_card_2.jpg +2025-08-05 19:04:42,439 - src.model.id_card_processor - INFO - Processing 14/42: im6_card_1.jpg +2025-08-05 19:04:42,449 - src.model.id_card_processor - INFO - Removing background from im6_card_1.jpg +2025-08-05 19:04:47,647 - src.model.id_card_processor - INFO - Enhancing im6_card_1.jpg +2025-08-05 19:04:47,652 - src.model.id_card_processor - INFO - Normalizing im6_card_1.jpg +2025-08-05 19:04:47,657 - src.model.id_card_processor - INFO - Processed im6_card_1.jpg +2025-08-05 19:04:47,657 - src.model.id_card_processor - INFO - Processing 15/42: im6_card_2.jpg +2025-08-05 19:04:47,680 - src.model.id_card_processor - INFO - Removing background from im6_card_2.jpg diff --git a/id_card_processor_main.py b/id_card_processor_main.py new file mode 100644 index 0000000..5dabff9 --- /dev/null +++ b/id_card_processor_main.py @@ -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() \ No newline at end of file diff --git a/main.py b/main.py new file mode 100644 index 0000000..c544951 --- /dev/null +++ b/main.py @@ -0,0 +1,276 @@ +""" +Main script for data augmentation +""" +import argparse +import sys +from pathlib import Path +from typing import Dict, Any + +# Add src to path for imports +sys.path.append(str(Path(__file__).parent / "src")) + +from src.config_manager import ConfigManager +from src.data_augmentation import DataAugmentation +from src.image_processor import ImageProcessor +from src.utils import setup_logging, get_image_files, print_progress + +def parse_arguments(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description="Image Data Augmentation Tool") + + parser.add_argument( + "--config", + type=str, + default="config/config.yaml", + help="Path to configuration file" + ) + + parser.add_argument( + "--preset", + type=str, + help="Apply augmentation preset (light, medium, heavy, ocr_optimized, document)" + ) + + parser.add_argument( + "--input-dir", + type=str, + help="Input directory containing images (overrides config)" + ) + + parser.add_argument( + "--output-dir", + type=str, + help="Output directory for augmented images (overrides config)" + ) + + parser.add_argument( + "--num-augmentations", + type=int, + help="Number of augmented versions per image (overrides config)" + ) + + parser.add_argument( + "--target-size", + type=str, + help="Target size for images (width x height) (overrides config)" + ) + + parser.add_argument( + "--preview", + action="store_true", + help="Preview augmentation on first image only" + ) + + parser.add_argument( + "--info", + action="store_true", + help="Show information about images in input directory" + ) + + + + parser.add_argument( + "--list-presets", + action="store_true", + help="List available presets and exit" + ) + + parser.add_argument( + "--log-level", + type=str, + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Logging level" + ) + + return parser.parse_args() + +def parse_range(range_str: str) -> tuple: + """Parse range string like '0.8-1.2' to tuple (0.8, 1.2)""" + try: + min_val, max_val = map(float, range_str.split('-')) + return (min_val, max_val) + except ValueError: + print(f"Invalid range format: {range_str}. Expected format: min-max") + sys.exit(1) + +def parse_size(size_str: str) -> tuple: + """Parse size string like '224x224' to tuple (224, 224)""" + 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 show_image_info(input_dir: Path): + """Show information about images in input directory""" + image_files = get_image_files(input_dir) + + if not image_files: + print(f"No images found in {input_dir}") + return + + print(f"\nFound {len(image_files)} images in {input_dir}") + print("\nImage Information:") + print("-" * 80) + + processor = ImageProcessor() + total_size = 0 + + for i, image_path in enumerate(image_files[:10]): # Show first 10 images + info = processor.get_image_info(image_path) + if info: + print(f"{i+1:2d}. {image_path.name}") + print(f" Size: {info['width']}x{info['height']} pixels") + print(f" Channels: {info['channels']}") + print(f" File size: {info['file_size_mb']} MB") + print(f" Format: {info['format']}") + total_size += info['file_size_mb'] + + if len(image_files) > 10: + print(f"\n... and {len(image_files) - 10} more images") + + print(f"\nTotal file size: {total_size:.2f} MB") + print(f"Average file size: {total_size/len(image_files):.2f} MB") + +def preview_augmentation(input_dir: Path, output_dir: Path, config: Dict[str, Any]): + """Preview augmentation on first image""" + image_files = get_image_files(input_dir) + + if not image_files: + print(f"No images found in {input_dir}") + return + + print(f"\nPreviewing augmentation on: {image_files[0].name}") + + # Create augmentation instance + augmenter = DataAugmentation(config) + + # Augment first image + augmented_paths = augmenter.augment_image_file( + image_files[0], + output_dir, + num_augmentations=3 + ) + + if augmented_paths: + print(f"Created {len(augmented_paths)} augmented versions:") + for i, path in enumerate(augmented_paths, 1): + print(f" {i}. {path.name}") + else: + print("Failed to create augmented images") + +def main(): + """Main function""" + args = parse_arguments() + + # Initialize config manager + config_manager = ConfigManager(args.config) + + # List presets if requested + if args.list_presets: + presets = config_manager.list_presets() + print("\nAvailable presets:") + for preset in presets: + print(f" - {preset}") + return + + # Apply preset if specified + if args.preset: + if not config_manager.apply_preset(args.preset): + sys.exit(1) + + + + # Override config with command line arguments + if args.input_dir: + config_manager.update_config({"paths": {"input_dir": args.input_dir}}) + + if args.output_dir: + config_manager.update_config({"paths": {"output_dir": args.output_dir}}) + + if args.num_augmentations: + config_manager.update_config({"processing": {"num_augmentations": args.num_augmentations}}) + + if args.target_size: + target_size = parse_size(args.target_size) + config_manager.update_config({"processing": {"target_size": list(target_size)}}) + + # Get configuration + config = config_manager.get_config() + paths_config = config_manager.get_paths_config() + processing_config = config_manager.get_processing_config() + augmentation_config = config_manager.get_augmentation_config() + logging_config = config_manager.get_logging_config() + + # Setup logging + logger = setup_logging(logging_config.get("level", "INFO")) + logger.info("Starting data augmentation process") + + # Parse paths + input_dir = Path(paths_config.get("input_dir", "data/dataset/training_data/images")) + output_dir = Path(paths_config.get("output_dir", "data/augmented_data")) + + # 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) + + # Show image information if requested + if args.info: + show_image_info(input_dir) + return + + # Preview augmentation if requested + if args.preview: + preview_augmentation(input_dir, output_dir, augmentation_config) + return + + # Get image files + image_files = get_image_files(input_dir) + + if not image_files: + logger.error(f"No images found in {input_dir}") + sys.exit(1) + + logger.info(f"Found {len(image_files)} images to process") + logger.info(f"Output directory: {output_dir}") + logger.info(f"Number of augmentations per image: {processing_config.get('num_augmentations', 3)}") + logger.info(f"Target size: {processing_config.get('target_size', [224, 224])}") + + # Create augmentation instance with new config + augmenter = DataAugmentation(augmentation_config) + + # Update target size + target_size = tuple(processing_config.get("target_size", [224, 224])) + augmenter.image_processor.target_size = target_size + + # Perform batch augmentation + logger.info("Starting batch augmentation...") + results = augmenter.batch_augment( + input_dir, + output_dir, + num_augmentations=processing_config.get("num_augmentations", 3) + ) + + # Get and display summary + summary = augmenter.get_augmentation_summary(results) + + print("\n" + "="*50) + print("AUGMENTATION SUMMARY") + print("="*50) + print(f"Original images: {summary['total_original_images']}") + print(f"Augmented images: {summary['total_augmented_images']}") + print(f"Augmentation ratio: {summary['augmentation_ratio']:.2f}") + print(f"Successful augmentations: {summary['successful_augmentations']}") + print(f"Output directory: {output_dir}") + print("="*50) + + logger.info("Data augmentation completed successfully") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..47078f4 --- /dev/null +++ b/src/__init__.py @@ -0,0 +1,23 @@ +""" +Data Augmentation Package +""" + +__version__ = "1.0.0" +__author__ = "OCR Data Augmentation Tool" + +from .utils import * +from .image_processor import ImageProcessor +from .data_augmentation import DataAugmentation +from .config_manager import ConfigManager + +__all__ = [ + "ImageProcessor", + "DataAugmentation", + "ConfigManager", + "setup_logging", + "get_image_files", + "load_image", + "save_image", + "validate_image", + "print_progress", +] \ No newline at end of file diff --git a/src/__pycache__/__init__.cpython-313.pyc b/src/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..49581cad59e510881d38143517717adc2f997416 GIT binary patch literal 589 zcmZWlO>5gg5Z#p}%U&b9ff7o4(N!rp>8TJ(p(GS)(>ATzn_yYw&1Q+L6}u}#d`kX7 zFZ~g{_kZ+S7JN!Cy`}XZ=txRGpmP}B3^Q-un{5;k1o7qbhqG&p&~;P%4(}Fhp9A=a zQlwFeHBLS4rM=Xj_a1vO0u;>sS)hX!_vigtsKXYA^TBMWhcVhm>EI34L<|$Men5_y zW%09GEO}{}m8y(SSboj~C+pzR^k90hzW@5@ZT$b(JEe;6pb{n6;XO#>7fV<>(JJT0 zC=H!E&DOu%(2I^#c`C(AR>H7OhS~=T9W<#A$wh|Sx}t>=LYBf!1h-T+#Pn1a+&EGw zw&*&J7`EcKc)VhTTrm5muCy#It+Wyv^t?(&E~Ip!a!TEZ(iPVR%tApZWwkw18gNis zSr|9z4#I9{rPb(mfStkY-&eJ=-8gR@(Pd|;vjMkQQBe8`ZNU|Cd--AZ-oVkU6!itq zj=4Fvs>0yk2cpmm|257Q-Ir0hcAv;+VE literal 0 HcmV?d00001 diff --git a/src/__pycache__/__init__.cpython-39.pyc b/src/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f1a7a0866e72a66b9e1b0a522636a7980d2075b GIT binary patch literal 559 zcmZWly>8nu5GG|Qwv}3G)^_pg3`V*ZMbRQ&I;0JnbRh&B2$5lkP@+Im2I8m4Gj!|T z_sGsQQ(vJ|k5+&LC<%VN?{I(L@x`JfcpiS=-#k)69us+aG9;hy`3-^~j0no8$XF&O zA~%yOf`4K1wqONK_JKKQr);{nr)Jj9*_;wZ#OxD-u}bMq#UmvH=u{Z4wUM5c#&R)wSn2szKH1fl z8!dQv;<4(T?gI4ADu?rYT^+_LfUpffJOtQE=TR(@5`Yik#yYf%A!y^{(by1kM=ARm z<1tp}82h}c{n)W;_KKh#C5P@Abk588_`d-bgN5T$W-!q%G zXV)EJ)_h@dJG4zZsy7`kl(nVDc)hY(W{Li;9ll@I6aA7~4HNBx*G>v+EgHdYB)Quj z3RP;`f93FS9(V9MQxr!5c!hvhH2%D!J2}U2@-=l)SyZ^TgHPbO2f0dt8$AC&sT7?O z>PM)a7f@T|r3ZRtbVxhGN1+`X(vID7Beu%P5vBaU|t&R z#bhkC*Y(s+@u{>BnAHI@Hhi1AUfr}kZr-ySE=&oNkmW|whE@v(tg3zXp)%bU$wVh1 zI!ug)klV~NLU`3S>nIj;6-c$C>}q=oe2u(Fi-VGxUDrG}bn(L2F$aKl8*~d7L_JvP zK2i5RWGf~)m>K$TNY~41%$nP9x;^5lwEZjT&vM&tHdf2Jpl=Wbh{Q;1l2MDM!%CmZ zYN7{Dqb8d?Qp6})etD+4vVJ>oM6mAFqqeht$qDX;ertX0;*Iq{RM*q4R<<~nTkv+v za`CB_Mg0QA6}h0k{{TF$iK%ct_nyLaUD1_ojVa#5SCrt~8JL0o?%)m$o}=EsP(D+)8Co2L z#^~3;aQQtlAx84f>PlE>|C%XVDpTSY(Kg>v_xtv@yo0xHsUN9cDg8a8oPh?WFT_jl zU~c=Zw9h1Dtb~a#b5U43{l+7f6dbSWlUJ@r$MMa7-uaG>hZDf0M=86BlIpZ(OlEClFk-{(_ft(fIv<_6A(ZV(2E?zbE-(aBFL z#lWsP(k(<0gCx0@aJ?{@w(3r3ai?ZSjWBbVq8*mhIny#a);(_>{Yg`9;fXnL& z&;OF~0x#}pF5@L$s;e9%$U%Ah$W9LR6Q~EN@o5ms44(y^&Vi(KzOZA!e2ekoRVOsF z4<58b{2cNBD*xU7qhHvD=Z`&s>)$1>DqEh)NQ zE5aPRA}RS%n3XF9)LPPI3c^o!n_K1LV{(+oXGwfW;!_g8A@K(iUqHC0FiGy)5OqA^ z`r<9TCGU9lM#JHWVSDw46I4c8!r!nr+=d&rAO8u>9-Aa+C04Y}R{gPsQaP85gtj0b z2*6;K(uN&Y-Ud~*BsZEZ zs6w|^OZ022SJvX!PMoxZYUnnda$ue|Pn`7ZZn5Fqa~gfmO(zI2;pG7Am6Mf;k-qn_ zPM<-XccrXXMhN76U)%+iRgScTQ*-D8YIZgBMWxhoJ*(+7ebKJuqrkE1t*C6k_d&Hm zGTLvfq=fmZ;;pcFWip6bEnkE%aQf09$-`WA)FmS=M@XetO;^MtCfCUD>X-+@mv9HC zAQW9?`1_Y`O#DaH4Jy;$KlHqTb))`U&t;l_>VJ!$8(K97SLlHyL(4{TEAcY1A6%WK zu3n9VKsYn#>SZQbKzswuhkz#~dS6N@mr)(iR4Zz{Fo>$x=l~cpkOVW10)g2@)aELC zv~#NQ@vB&W{WCv^OneQU9;qowX?V+_80YEr{extHuc7lk;1dS{AJXF~(JvADo~I09 ze;uzE-z0I^#Lg1vUf&Gz^*+ci9E5zJ{Zj&-mw?INhp@i(M644MD|O3YezygO=QDC9dy9Vv>l;tRIw$9|*Ez#` zrw{Lohs^4M4xqz))WRm@1mYMv4+ONt?cL4uT|C5dh-qjEGDK5|vMh#Zirqvs6{>CQ zATkPC3N^0wqM-K>S8tGO^ea?zWnHA+zqMZ(tRZvW9w_1&McMJ-hyAL;H5HS>3-x85yx-~N$f}259C2p@p zjlN-SMKpuGa7s1?#b=+iVyV6U<1%q+^4YALxOtePjdiNH#$FB58TD_WnhbOW7vrg7Xj^wyW$ZTUy9M zHPk4JJgmgoj`ky^CPS4L%8MDlj7czodNc(z?2(ZRzyxdy04m`GDHdfLPmB<&&P4fKKW=`nGIw zR{3C}qmxdS2sJ}Kl;~9-`AW_Woo0}fT%xzNUg^Y^M zFwm;-Fvk=yPH`ZA`~E=wR^X)$+t0wd>{M+33hbmIM_URDe59lGT7aF3TiQn&M>wCB zBs&#h{n7ZhU0S$_NYy+v_8~qQ$g8`34sui9r1cVX3WBpmFgVcv&r3AK|6(Pg_?AsE zTYwCbGC9k?o3Q(00(KFT5Ld*jkY%khwV$}v8vcDN2-`?N4cFu7TN$$(q>MZlIjFOm zavd_z)KB(v))WY4B*H0N9~VO3&xGYt|0zR?F3C0TBB>EG1Zx%o>nm!uQi3FI6vUGI zPJM$E^o~PWGEFS1@*v>{Xq2Xu;K(U9fpqdUb$p*I7YF^E@MqLGu$|{I2*tv*RD!=k zq~|$sTAbj}K>>a>`UM$*gmC15^d5wyQkJ>bJraO6FSR(9>uJ#kXF+gyIABLbpHhD< zv?EeG+OI_)9euRALPWQllf-ut#FyJLLZf4Zd63;F9Rs5wm))bXzFX+z_JHw&PcY_@ zjxK)BCw7Y+eUEu6pWI&(sFPp9r?3R3<74kqVn0aI*#DYM?5pvsr8l4Yif1pbtzNly z_4bVmx2~zkcDyo%pLyJ<^-;2(W-Dm>9Qeb+1JQ^uCDv<&_!d#TK3x%fba zg5wy%zliJ8S7KD2}}q@Env4PXazRM=DJeH{xio$)g?84esE13 zHW=-U_+Ox7Jc40aNy)OBK93qy9M}KYSfF(q$Ndr%$BJYdP94dR-KHfJu++}%iVC@; zhg`>6a4(e$^b!=g^w7G1lw0n-w0}X)Jx-53=@y{3qG&COW~JlFpE>HflsqWldjlaB_K+fb5Uf}o0{P(sy!I7I%cjt0M4pofNI5bVic<+K#|o=XRraYt&8f+GbIu(68bLj5I1Sln2J_*9vk)#ii{X;96fQf< zM7gKX5-mShX!%HWE>QK6(ya7vTbo|$*%q_Mo$bf#kMD19S*NfQWvfYq{?JSJn}*OI_(>{^$GtdUqP!jBl3pP6O}~{2ZM}O4 zW9#>E;fFkGu_TFjQx{ecdDNBP6IS9qWyAN&EzZ1@xkCjRb=!Wxx?adcJ?4Hlbz>fN zI4sjbMJAszRBl5_{tSXCU}VKnDFGv2U^+9H*;X$rRC@ttrp9x_DY5cFg;^-6`9hP% z(w!<@k?l1$hrckR0(};(!i+k-$Qtku`x?+wd!f@gT8Ex_dTH2$Uy{A@YYiANPZwas zg4Chtnb2tQg-Vy`@^jr;luo)V=Fvy%6E5n%xiRA-dd7=0jzSlit*G1fJ1rFqGJ5bG z%LbGLFso!FHhm>iTqPw(_zN}zsHpu_@RW}6e44fgU3_GlVYK*rpD`_E=M(Zj zUw}dXc&oL#`;!C~+3j|EhirF~B_E|xyt}=zvzzeNF4}MeG^x||_)wHx*N5%8?mR4! z=pdAmMf#U#H-3FIuBizl$FbpcJHP>xv$V&%EhPWwL8VP7(&`W3B_jti82kz_xvLJe zfu8A9y`p3q%CrhXInXkanVFtxZ53)d)j!b(rK|+LGWuT>cG`)6 z)ojS*y+ga>KV@B8YF3OXUe5YT?qt=z-;H}|zDO|`T`AXDwdF@JbjZ5xcWugmlx*xS z99+!Ykp`JTMXo5F`N-BOg0lk!o-f~N635z$IA#3(h%$&aMzZh4OlYmAH^fqhr6Wz= z_fqaZyRj*(k3Al*x7&V~iTQW8cQ)_rxI1?qKDfWJ{`gL_%+alU8N~$fiVyj8Fv~S%?LYmL%Fb1k`0VXWaVNJ=hyP~JbGP$mIOhnR2Rv-R@dr8|F^02 zzjg*Zld}NWsE=6a_QEMRCLH-Wg`)~I7*HgRVtC2a*ra5jc9emZDF@0i95B9eq9UH) zUF#@R|BOK0NYzX|BAd#GKOC4DIo3`TKo5>FsClCDU!?l6kr^jsP|lQNGb?3fm{rQO zW5C!6A2o8EY|l=7WWQ3X})eB{t}gA1B2pUkyt<&8C-=N{x|yioABpqg@MAav*6|POn18=^Bb> zD1M9Lk0}0v;!6^jY0tBe&n}c zmpIJy;+S=5voe9?H?U)Vd}QVtPk=a&!2BAHBfe89hDs95VtRSl?dI;wvHd2@mFI<0 zEfXD(J+IbCS*sC?)YO{V|K=HtpSKF0{8BqNhU`IDQ7?r+<5Hh<+(Zs|8FZnBPZdf5 z_y8&?giZiDIZ%%z$SLV5{7Of81V`7ghT##x4s7`^EG8)r2zHJI#T(K%shs-A=pg-yXJiWzADhOo*&6Gz_u|%fk;Sg%wufGO?p8( z#>;s2`EZw(ShX>mLcsbjiM(S=t=iktYBpGnwbE$%00sQd02VR-RWb)XE4hc}pcA~D z(QE4zYzsLl$Dqa=AVf9wIM4uq*Ju1Dwp19X&LCxA9HAg|Oa(b+-5e3wtmZWM9yWX* zg@loe?>C_8RHYKq#tDB5TM+y?q@Bjls6~od9X8+=W?8W~!gG#aoU;w<<(!zS6J$kA z(}5RtVrj^V{W#z&#f!#IK|Z3 zJSBbeZq$Qd)L)!fZHQg|1HesxVHn587VKB8VH*&(zdV|l^jer+w|hZw*ena9m-<1% zZ^0VEke4x`!QG7i2pV_|1s;AvOH$sna$dfTtz%HjsnrQ(MQIo=B{wasddmbJR7le)T0`d80zaB6x}<(sIZq}${@38`0t zvVrSwfxv?YC3=v*hC?Ae&yOIx&7;ts+HUd5SHtwtP@(ZXt(e?G=#3sU zng%7of@aK=59#}Dtzx6v9 z`gbTwf9;G55%NQn2!}BGsrm%` zGojZ#7&1|nX}`4=K%Ur&!uZxPcK;Ye|Ly-r&te@o2ri^>>j>OYG{{+eC;08Z%#=HNmBqIpH_3 zT7C<%%49`TiAom9@<02{|JpnfD*zYye+jSTcQi`*r~oJ-cq#n-Q1dE*;jmwn9@{sE4=3_=)?z{H3AXZT)1%m@CS{E{%! zLyZ2p0vK;NX1e(fr1IavIXc`x_hot{|MbD%!7(yrl35UDMS1W=jm}e7l-F{j-9p+X TxC31)L%vcYuY<(jOX`0B6u*8o literal 0 HcmV?d00001 diff --git a/src/__pycache__/image_processor.cpython-39.pyc b/src/__pycache__/image_processor.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93f0a66971a81c47e97ff7ce74af5194986183ff GIT binary patch literal 5218 zcmbtY%WoUU8J~SGDT}wjXUtBzl4I47NHu|Bdw*$uii4?YeZ((YFQeNnUURf zT29w(x!p>uBKxeU+MR37b!)9!cfK{>t+(nLStHb@&T~SYL%p>?^$$tY&2LxNyTK0g zdOT^fG!5e&uN?(x>U9$CZ3k)C_Ch%!=`kK;VG^g+6WAEpXllZI6lA+iLs*ady@(0- zVGkz+kuV>GY1TAFHA({NO9G;r22a_bUVj=yAq}!&dtrxV<&fVABbK(2!}w!C;|{#( zKcHX)z#uK1Y5>MHW-^Oe9sM$)#tVRhTFk%3c+E?H4+AsQcV(&0)`IKDKD( zg+Z%y4n}Ko6#HthVxHEaZ+@b0fi6N{9cV8|s%Pj~XfHzBVP}{Fyv_l&20ahGXX*7B z+%CfSC3+clcy5C76?zrcG^9|LL|w|{Q7Mvy=Wngf6v>#(M$*|h*2G*Ea6~i>bJo@Y zvEhStd{^O3>rfP0&!$>Id?HyJ`ERP;;8~BuECey-%$wYfnxQIp`A)i`TDX2}11sKR zxf>66Bh{)@D9)e*^=BC2Wz;L$SfQ) z(+hxm&p6Z%^g^e4pYU7QU+70>X6@N8!MLvy90T(-^S>!H9)KlzK}^9?+2AmxWvUsf za#|dt{jYdvvtE`A%o@#2Ekwzw#-=>ylw2Pv=8O&~1gRJIyS;rc;5^v(z-??KG4sNX z*9m|)RRr_KdMxdJjP;7Q@o&J3kS~)vO^3Ts6IPFhamH<^OF;`eNVEMO6MEbeP6u2{ zcKcS-;Hba+tXv__C)J@W7P0r7bUpwX2!9CK1sF(ilvG`0Ln{rO8Mu zQFBz~6(U{n*5h73n<)vk7Nz=1dwiAgy257J}2mqapL92>%DKfQ%^ zgXb01^^M2(H~qVt_wNp6I06+7dQ5_~a!nZ-egI>otvSD<#wmpRvdT{q3_9&F2?4-}c`5L`_$fuOlfn z`SmEPZA=5r8*lXQCOK>*c9STq3t=;d|~5ClVtAx?=^DZ zXe9qmR-v%|Gsu(*2V#|pu?jiZ!bJG?k z?m$>F7ZG-uiwL*98^ke-Qc=S|DHAHcyDjYUwE6Y>oI}8aTq~R8CG(tIU-6n%z61<} z9dOn!OI)p)2vmp@H4O`RY4?^Nsu3Og{T7G@X zF-+2d+`kfH1oCOfMHS>esGkmx6=0^qKZ^Fj zHc%_1dJ7|Bv8XUeCGnogL|tZqn78jhbWiyem?_LSfZ^}qA_v=c(Ce|7HY=rV{168) ztcf{k_T^==wJ_K#qh*KCLYX&8GZJ$Hfbq=JX4BXXvi9!q@=&_>vpJ~2d+eEI~coTskvVVc5iue?Z^L1RV`csxH#e08udfLMxm zHBzlZC^@C~v(_{<`1gq4hc(Ufco|2l_I*+H{cb}05NCZ+^Zh6NAR3G~zE6|3?{hqc zFz|YK25KCo#Zg=`0~St98?{Dlap}T$y?HsKxyT>kY>B8P@4|!api7CuZMjWh!yOj< zqdY>nXWRy)mRT63aU%9od$2} zP=8FcS{E8MimO$crFD~|O&VJ(EkE|7sV1#I?Km3jmTJ+KKauvsS(0|@{_H)+zQ!O_ zbw~KT=i@zJ?|JU?zK6$ly9Gh|*>4svU9=a;cfj4p5`7qH z3Rr#5t<9~M=xhebQC+85YL^-|*OZ*FQe(gkb2asuq^9bK7D<#`&;x5p9(ooYdYh}g ztx^l!-xkm#WP}n2zsLWA>;{^ub@&7eCDrvNomFi18Q40;(%tgH}|irTZ`qXEAH>je^# z^qd@cMOI!{lc~U&L0>?@;XoSTr|eD5X?BI!Bv{=NZh_Inq{>;y=FtP*y3DuS<8RvLyL*j&rD_Pe(Bf_S)OdtsGSciy$_EZI8C0@q~ux6vpV zo;VS-Tesh}ZvXIH$@)xzf95evjWqL26w6^XiRaeglbnI#(?q#wAF$1)Lgp%ZX|1m|acWVQ_nW~$+OtxID4oCEdNP^n`_ME~zB-}ZfM zzQU>V2w4#zTjyd}Q(qsj%Vj!7WK|5sVq%nh<>D-W11JM0<6_N@ULnKpy_||<1jywG zR8k3&PsdZxq(*0FHRd_5pc&+=Kp_f|s+^e7_(&`{jT^xbnn6ioRi5$kcnhgEO8$w7 zGrq~8SAxT5e5d*+CCy6L57PBD4rWkz&=euzW-_D!pc$;w!JW|B3O_}KY##lib)dj| z?l-yTPp%l9%f^fJv(&ld$wy1=FWj>o z*`)K2-rsva-o3P=u=nLs`|-cp`X33%z3Yhq*_sx2m#l3Cz73PrHA`hrDIn2xcr!pJ z#oMo;a09i~bIYcT=tug!uqeU!naz--W0Lws@S!>D(MueM1m40Bdov(#rQJ@}KjLuRlh*3ohiC0K#1tF&=hv60p zPiuTKB_}k)6&Xio=QQ)s)v!#xhL^=%q@Rn%BMLdIeuoY^56oVogbA#aXhsg#UrNoH zy$FWf6IGBN0mbLxr_4h}Kmg|-ed$XCJzg%+e zUJzES&bN=PxLSZA8r?T-*KJw~Uu7`or!Uaq(x3HhBvmVs^DXQ*+Ynft6+fHFP|8k#kkR#Ry;NN7nj{vZwCuR0gzpi56Rh#&fn z5ck)^jMt}GlQnR1z+Icb+gI~IYhxgqhxfn{c3tGPDTJrsZBa>nMhBjstAlk7wz?j8 zYjz)duzF@-LBfAdy|)T(e#)h{0Js?nZXP`8dT6Ifj%Gk|4Iw-~z{BeNY=PbB+&lzb$BE8QAS1B137!Xl-|_3?InQso#|8Kn#8>OHeVp6WAHVYZP%>tSl=JaT1)P(T=bL9<({!(&)DZ9 zg#(lKO#YPymkxdp1%GGmWbW8{@VhDaf!>M3eXDK{t(zCPCtLvaoPEBC(_%dWS5Cm-?H>Lz9F@6>R0j^kL8vqhC zP^|Y7)ggopi6+dm$!7sz;`Ue(qQO?@AP^WhR}$rr=$gLGszJV7b+8%})VTdUGvrC8 zsB#=1hjs7?$bh$s8k~#Gi>54--Ie3A+Y0tw^TR8oyS->_&rRlr^M~@?OUIUbM~b~8g`N|Ij?toZ zw7`#6Vjwd)4>kL`5Qx}{+WAE&P&)(FY)pDcv)MCIVyc#7GcFRQl4Zn-4@rop2BHW# z5qxc6CU^-Mm5L*W_sVHlr1D;=#Ewql1U0FK;Q5=v@~P01Kzy}og6D3S(l7@~ z3OPC@JDu(vG$07-d&Xq17vQJ73mJf-(KSB?F8f0vd%4)w^NIiV!H)yQw*FGEg ze`>YCz39)LeJ}V?I`9AGkBXk&rB_Q0hv!eMSnP}TqGd;p%Vmns_7}y0qGe!yaFsVN z3@lt(xU?9~c4j@<>0E2hlJCq%^5>UAw}quQZifn%Du+;DazBD@?c$6K_pjkR~M3ad^D5iIm1ojE>Vdzv62QmH!LShdh0IY5a6D8KW^g z{yt2g851)hj6-vnL`-;yWQ1|>G|9*ZM-6-n(Zf_eOw)_;WF_tfo<4e@~chp@(-CrQ6Hhqaak9Z67Y*s2G5c$aO61NghNwM5&X4vU}6XVYIl@W-%yfk>*wRGZEl2$)&E%+vveSsqU zjb408G#+(37|-JL&CBmyE+Z)BVJ;B2!&H3U_mtEgi99pLSkPb5txzUy&NM*8czvf|n}* literal 0 HcmV?d00001 diff --git a/src/__pycache__/utils.cpython-39.pyc b/src/__pycache__/utils.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b62b19a15b19d1660e5af26da604618214e16804 GIT binary patch literal 3426 zcmZWsTXPi074Gi2>_yr|hy@Fa*aMi$5Sw5^B`$&;gOG?4At{nD#YEK9XuEeAF}t(w z=~dXRmCIG4@ezuz^_ZPF9XD`0rJ&R)2W2B; zNhPS5`PEXAVGCE-rwOk@Ho+&K+re?jrg;th zIzNW~34R`ul4yuVi=8H-EGD*3V#U-eiy!CH&~l1@re&Dbz{V4(pW+_&aCRM!_98HKMV=E13NLs^l_k{AIB=KVHMP3W0 zf=kM@hvH|ajTWnE54Jv;u%~t)3p*^p4;st-IB*} z`rOS@sRZ|J*+4fqzW!i!b!Ba1{%*LivUdOW#(F+s_7Bbefdvi5_(d_TRU89S;gLx* z5N350gq308GX74{GgQt(0_#?bsXYYGPtD6P=bd@i&>4OO_l6Po>EC8{>4?5z`(Uhx zFuF}_vIqC?PJKzTE7u^^e~*e$1hq|+HCQ^7a%-Jj!zvRbXTOV*7+x*BxMd9M#hDj* zJ6c+pdwF94Ceh)T+i6d919wNrxHZTt^Us>Xgt*V-8H@*3+-ZMklDbia4GO!Q$Gl9t z9bpU}cagd%<)`TTmQj*Vs7PuvX7NxIR?(I?{It_)IyX-D7c_}vaSgMt_S+C_Zh~hyo0tgLJPIa*$Ydrcu0(mN`oIM*@Yd4=TwRNTnrqj%& zb=sivGDhBpD4_Wws(d2tt6pD)+I{oV_kFmG8qDfl7fWe7QW>0ZvH9z^MY9sKLnwdA!HT*|mSL zxeJ_3%i!d|<4eQo^=M}dj;B)N^Bai8QZJH7Y*}753=yVJhbIc4vS{Z$PVy{nu)N?` zfSu+SO1Kl-@ zj_JQsCyqT`r!1`}(a`{H%z#qv+{)(+HY)(|Sf^8!^J^h{5V$A=i_3pA`*#1-iK{2bgF45N8dXD{=b7$;FNW|^FZfpzUV*N=!yiH^LA zz2prPa5wz~kiR?>$Y1H!2X;8!G)lo{JJa+<;K|ce{u(Wv+4O0Q4Sq_! z`u(9v$6Q*|cahDu;iBeBYLkfo6v3nEVaed3jZ)ZNzv6la!i$sl79IRH`C8ntDVR=nQ?ud-&Qzm~@)*f;XuLJ#K!|7E`M_zAy8*eL;Z#AH8{O~>c z_N_Me0sZ=besMr+L+hsZry+an4Xq2_tdKnIEL0j&~w2$oVw7_4g zh@6uA8hYe6y7(=Mpo*RF$n3V&lb{AuCvb{!w~QxOQ%8GXVbX`#)D3}#hbIo8Y@Ebx zI)en>pr3;tEfOUI;wYs=ul5jyIZt)&8&ouLz$nC8(O-v${i%<}h zx*Ru)28Oq#j$GsU@(;SD9n75n$kk0inO98tRww2CM0}%x$#hP!E&X(v#LZ338k3EQ Osmi&=%+x7&%Kje`FF>yV literal 0 HcmV?d00001 diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000..2231176 --- /dev/null +++ b/src/config.py @@ -0,0 +1,40 @@ +""" +Configuration file for data augmentation +""" +import os +from pathlib import Path + +# Paths +BASE_DIR = Path(__file__).parent.parent +DATA_DIR = BASE_DIR / "data" +INPUT_IMAGES_DIR = DATA_DIR / "dataset" / "training_data" / "images" +OUTPUT_DIR = DATA_DIR / "augmented_data" + +# Data augmentation parameters +AUGMENTATION_CONFIG = { + "rotation_range": 15, # degrees + "width_shift_range": 0.1, # fraction of total width + "height_shift_range": 0.1, # fraction of total height + "brightness_range": [0.8, 1.2], # brightness factor + "zoom_range": [0.9, 1.1], # zoom factor + "horizontal_flip": True, + "vertical_flip": False, + "fill_mode": "nearest", + "cval": 0, + "rescale": 1./255, +} + +# Processing parameters +PROCESSING_CONFIG = { + "target_size": (224, 224), # (width, height) + "batch_size": 32, + "num_augmentations": 3, # number of augmented versions per image + "save_format": "jpg", + "quality": 95, +} + +# Supported image formats +SUPPORTED_FORMATS = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff'] + +# Create output directory if it doesn't exist +OUTPUT_DIR.mkdir(parents=True, exist_ok=True) \ No newline at end of file diff --git a/src/config_manager.py b/src/config_manager.py new file mode 100644 index 0000000..62d612f --- /dev/null +++ b/src/config_manager.py @@ -0,0 +1,175 @@ +""" +Configuration manager for data augmentation +""" +import yaml +import os +from pathlib import Path +from typing import Dict, Any, Optional, Union + +class ConfigManager: + """Manages configuration loading and validation""" + + def __init__(self, config_path: Optional[Union[str, Path]] = None): + """ + Initialize ConfigManager + + Args: + config_path: Path to main config file + """ + self.config_path = Path(config_path) if config_path else Path("config/config.yaml") + self.config = {} + + self._load_config() + + def _load_config(self): + """Load main configuration file""" + try: + if self.config_path.exists(): + with open(self.config_path, 'r', encoding='utf-8') as f: + self.config = yaml.safe_load(f) + print(f"✅ Loaded configuration from {self.config_path}") + else: + print(f"⚠️ Config file not found: {self.config_path}") + self.config = self._get_default_config() + except Exception as e: + print(f"❌ Error loading config: {e}") + self.config = self._get_default_config() + + def _get_default_config(self) -> Dict[str, Any]: + """Get default configuration""" + return { + "paths": { + "input_dir": "data/dataset/training_data/images", + "output_dir": "data/augmented_data", + "log_file": "logs/data_augmentation.log" + }, + "augmentation": { + "rotation": {"enabled": True, "angles": [30, 60, 120, 150, 180, 210, 240, 300, 330], "probability": 1.0} + }, + "processing": { + "target_size": [224, 224], + "batch_size": 32, + "num_augmentations": 3, + "save_format": "jpg", + "quality": 95 + }, + "supported_formats": [".jpg", ".jpeg", ".png", ".bmp", ".tiff"], + "logging": { + "level": "INFO", + "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + }, + "performance": { + "num_workers": 4, + "prefetch_factor": 2, + "pin_memory": True, + "use_gpu": False + } + } + + def get_config(self) -> Dict[str, Any]: + """Get current configuration""" + return self.config + + def get_augmentation_config(self) -> Dict[str, Any]: + """Get augmentation configuration""" + return self.config.get("augmentation", {}) + + def get_processing_config(self) -> Dict[str, Any]: + """Get processing configuration""" + return self.config.get("processing", {}) + + def get_paths_config(self) -> Dict[str, Any]: + """Get paths configuration""" + return self.config.get("paths", {}) + + def get_logging_config(self) -> Dict[str, Any]: + """Get logging configuration""" + return self.config.get("logging", {}) + + def get_performance_config(self) -> Dict[str, Any]: + """Get performance configuration""" + return self.config.get("performance", {}) + + + + def update_config(self, updates: Dict[str, Any]) -> bool: + """ + Update configuration with new values + + Args: + updates: Dictionary with updates to apply + + Returns: + True if updated successfully + """ + try: + self.config = self._merge_configs(self.config, updates) + return True + except Exception as e: + print(f"❌ Error updating config: {e}") + return False + + def _merge_configs(self, base_config: Dict[str, Any], updates: Dict[str, Any]) -> Dict[str, Any]: + """Merge updates with base configuration""" + merged = base_config.copy() + + def deep_merge(base: Dict[str, Any], update: Dict[str, Any]) -> Dict[str, Any]: + result = base.copy() + for key, value in update.items(): + if key in result and isinstance(result[key], dict) and isinstance(value, dict): + result[key] = deep_merge(result[key], value) + else: + result[key] = value + return result + + return deep_merge(merged, updates) + + def save_config(self, output_path: Optional[Union[str, Path]] = None) -> bool: + """ + Save current configuration to file + + Args: + output_path: Path to save config file + + Returns: + True if saved successfully + """ + try: + output_path = Path(output_path) if output_path else self.config_path + output_path.parent.mkdir(parents=True, exist_ok=True) + + with open(output_path, 'w', encoding='utf-8') as f: + yaml.dump(self.config, f, default_flow_style=False, indent=2, allow_unicode=True) + + print(f"✅ Configuration saved to {output_path}") + return True + except Exception as e: + print(f"❌ Error saving config: {e}") + return False + + def print_config_summary(self): + """Print configuration summary""" + print("\n" + "="*50) + print("CONFIGURATION SUMMARY") + print("="*50) + + # Paths + paths = self.get_paths_config() + print(f"Input directory: {paths.get('input_dir', 'Not set')}") + print(f"Output directory: {paths.get('output_dir', 'Not set')}") + + # Processing + processing = self.get_processing_config() + print(f"Target size: {processing.get('target_size', 'Not set')}") + print(f"Number of augmentations: {processing.get('num_augmentations', 'Not set')}") + + # Augmentation + augmentation = self.get_augmentation_config() + enabled_augmentations = [] + for name, config in augmentation.items(): + if isinstance(config, dict) and config.get('enabled', False): + enabled_augmentations.append(name) + + print(f"Enabled augmentations: {', '.join(enabled_augmentations) if enabled_augmentations else 'None'}") + + print("="*50) \ No newline at end of file diff --git a/src/data_augmentation.py b/src/data_augmentation.py new file mode 100644 index 0000000..430ee79 --- /dev/null +++ b/src/data_augmentation.py @@ -0,0 +1,161 @@ +""" +Data augmentation class for image augmentation - ONLY ROTATION +""" +import cv2 +import numpy as np +from pathlib import Path +from typing import List, Tuple, Optional, Dict, Any +import random +import math +from image_processor import ImageProcessor +from utils import load_image, save_image, create_augmented_filename, print_progress + +class DataAugmentation: + """Class for image data augmentation - ONLY ROTATION""" + + def __init__(self, config: Dict[str, Any] = None): + """ + Initialize DataAugmentation + + Args: + config: Configuration dictionary for augmentation parameters + """ + self.config = config or {} + self.image_processor = ImageProcessor() + + def rotate_image(self, image: np.ndarray, angle: float) -> np.ndarray: + """ + Rotate image by given angle + + Args: + image: Input image + angle: Rotation angle in degrees + + Returns: + Rotated image + """ + height, width = image.shape[:2] + center = (width // 2, height // 2) + + # Create rotation matrix + rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0) + + # Perform rotation + rotated = cv2.warpAffine(image, rotation_matrix, (width, height), + borderMode=cv2.BORDER_REPLICATE) + + return rotated + + def augment_single_image(self, image: np.ndarray, num_augmentations: int = None) -> List[np.ndarray]: + """ + Apply rotation augmentation to a single image + + Args: + image: Input image + num_augmentations: Number of augmented versions to create + + Returns: + List of augmented images + """ + num_augmentations = num_augmentations or 3 # Default value + augmented_images = [] + + # Get rotation configuration + rotation_config = self.config.get("rotation", {}) + angles = rotation_config.get("angles", [30, 60, 120, 150, 180, 210, 240, 300, 330]) + + for i in range(num_augmentations): + augmented = image.copy() + + # Apply rotation with random angle from the specified list + if rotation_config.get("enabled", False): + angle = random.choice(angles) + augmented = self.rotate_image(augmented, angle) + + augmented_images.append(augmented) + + return augmented_images + + def augment_image_file(self, image_path: Path, output_dir: Path, num_augmentations: int = None) -> List[Path]: + """ + Augment a single image file and save results + + Args: + image_path: Path to input image + output_dir: Output directory for augmented images + num_augmentations: Number of augmented versions to create + + Returns: + List of paths to saved augmented images + """ + # Load image + image = load_image(image_path, self.image_processor.target_size) + if image is None: + return [] + + # Apply augmentations + augmented_images = self.augment_single_image(image, num_augmentations) + + # Save augmented images + saved_paths = [] + for i, aug_image in enumerate(augmented_images): + # Create output filename + output_filename = create_augmented_filename(image_path, i + 1) + output_path = output_dir / output_filename.name + + # Save image + if save_image(aug_image, output_path): + saved_paths.append(output_path) + + return saved_paths + + def batch_augment(self, input_dir: Path, output_dir: Path, num_augmentations: int = None) -> Dict[str, List[Path]]: + """ + Augment all images in a directory + + Args: + input_dir: Input directory containing images + output_dir: Output directory for augmented images + num_augmentations: Number of augmented versions per image + + Returns: + Dictionary mapping original images to their augmented versions + """ + from utils import get_image_files + + image_files = get_image_files(input_dir) + results = {} + + print(f"Found {len(image_files)} images to augment") + + for i, image_path in enumerate(image_files): + print_progress(i + 1, len(image_files), "Augmenting images") + + # Augment single image + augmented_paths = self.augment_image_file(image_path, output_dir, num_augmentations) + + if augmented_paths: + results[str(image_path)] = augmented_paths + + print(f"\nAugmented {len(results)} images successfully") + return results + + def get_augmentation_summary(self, results: Dict[str, List[Path]]) -> Dict[str, Any]: + """ + Get summary of augmentation results + + Args: + results: Results from batch_augment + + Returns: + Summary dictionary + """ + total_original = len(results) + total_augmented = sum(len(paths) for paths in results.values()) + + return { + "total_original_images": total_original, + "total_augmented_images": total_augmented, + "augmentation_ratio": total_augmented / total_original if total_original > 0 else 0, + "successful_augmentations": len([paths for paths in results.values() if paths]) + } \ No newline at end of file diff --git a/src/image_processor.py b/src/image_processor.py new file mode 100644 index 0000000..46c3111 --- /dev/null +++ b/src/image_processor.py @@ -0,0 +1,174 @@ +""" +Image processing class for basic image operations +""" +import cv2 +import numpy as np +from pathlib import Path +from typing import Tuple, Optional, List +from utils import load_image, save_image, validate_image, get_image_files + +class ImageProcessor: + """Class for basic image processing operations""" + + def __init__(self, target_size: Tuple[int, int] = None): + """ + Initialize ImageProcessor + + Args: + target_size: Target size for image resizing (width, height) + """ + self.target_size = target_size or (224, 224) # Default size + + def load_and_preprocess(self, image_path: Path) -> Optional[np.ndarray]: + """ + Load and preprocess image + + Args: + image_path: Path to image file + + Returns: + Preprocessed image as numpy array or None if failed + """ + if not validate_image(image_path): + print(f"Invalid image file: {image_path}") + return None + + image = load_image(image_path, self.target_size) + if image is None: + return None + + # Normalize pixel values + image = image.astype(np.float32) / 255.0 + + return image + + def resize_image(self, image: np.ndarray, target_size: Tuple[int, int]) -> np.ndarray: + """ + Resize image to target size + + Args: + image: Input image as numpy array + target_size: Target size (width, height) + + Returns: + Resized image + """ + return cv2.resize(image, target_size, interpolation=cv2.INTER_AREA) + + def normalize_image(self, image: np.ndarray) -> np.ndarray: + """ + Normalize image pixel values to [0, 1] + + Args: + image: Input image + + Returns: + Normalized image + """ + return image.astype(np.float32) / 255.0 + + def denormalize_image(self, image: np.ndarray) -> np.ndarray: + """ + Denormalize image pixel values to [0, 255] + + Args: + image: Input image (normalized) + + Returns: + Denormalized image + """ + return (image * 255).astype(np.uint8) + + def get_image_info(self, image_path: Path) -> dict: + """ + Get information about image + + Args: + image_path: Path to image file + + Returns: + Dictionary containing image information + """ + try: + image = cv2.imread(str(image_path)) + if image is None: + return {} + + height, width, channels = image.shape + file_size = image_path.stat().st_size / (1024 * 1024) # MB + + return { + "path": str(image_path), + "width": width, + "height": height, + "channels": channels, + "file_size_mb": round(file_size, 2), + "format": image_path.suffix + } + except Exception as e: + print(f"Error getting image info for {image_path}: {e}") + return {} + + def batch_process_images(self, input_dir: Path, output_dir: Path) -> List[Path]: + """ + Process all images in a directory + + Args: + input_dir: Input directory containing images + output_dir: Output directory for processed images + + Returns: + List of processed image paths + """ + image_files = get_image_files(input_dir) + processed_files = [] + + print(f"Found {len(image_files)} images to process") + + for i, image_path in enumerate(image_files): + print_progress(i + 1, len(image_files), "Processing images") + + # Load and preprocess image + image = self.load_and_preprocess(image_path) + if image is None: + continue + + # Create output path + output_path = output_dir / image_path.name + + # Denormalize for saving + image = self.denormalize_image(image) + + # Save processed image + if save_image(image, output_path): + processed_files.append(output_path) + + print(f"\nProcessed {len(processed_files)} images successfully") + return processed_files + + def create_thumbnail(self, image: np.ndarray, size: Tuple[int, int] = (100, 100)) -> np.ndarray: + """ + Create thumbnail of image + + Args: + image: Input image + size: Thumbnail size (width, height) + + Returns: + Thumbnail image + """ + return cv2.resize(image, size, interpolation=cv2.INTER_AREA) + + def convert_to_grayscale(self, image: np.ndarray) -> np.ndarray: + """ + Convert image to grayscale + + Args: + image: Input image (RGB) + + Returns: + Grayscale image + """ + if len(image.shape) == 3: + return cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) + return image \ No newline at end of file diff --git a/src/model/__init__.py b/src/model/__init__.py new file mode 100644 index 0000000..ecb9162 --- /dev/null +++ b/src/model/__init__.py @@ -0,0 +1,8 @@ +""" +Model module for YOLO-based ID card detection and cropping +""" + +from .yolo_detector import YOLODetector +from .id_card_processor import IDCardProcessor + +__all__ = ['YOLODetector', 'IDCardProcessor'] \ No newline at end of file diff --git a/src/model/__pycache__/__init__.cpython-39.pyc b/src/model/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..231b3dd8b9512918438d45b615e8fac91c63120f GIT binary patch literal 352 zcmYk1&q~BF5XPIf+g(bF_y#%JgW_2h5w@2FwTt4x5*TVG)?(U}q*w>soBiLN`5(SDXZD!`E9jLT67(f${n95c4`k4W zWoYvdXSH9>FuL}#`O`AKP$(pA6>g=V_e#^-;=DM#tBgRp&MB`nQWS&WL!&xcbx65Z zecyC-{6Xxnj7*6`@}C-|t>0VcD~oNk;=&lEwcAbyRxYz>nggKX08=_DsbJUpw`kgL z5Zp-U|6s>CsFD)k#nTJd%*od=KVdgUXv4aCI|`Ny^EjxU6{{;|G-viGAsIl^H3Ps= UKTeC~?$A@WxrwZVVL}rB8^woVj{pDw literal 0 HcmV?d00001 diff --git a/src/model/__pycache__/id_card_processor.cpython-39.pyc b/src/model/__pycache__/id_card_processor.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23bc067fee657c3ad34ecc93b4013fd3d5b3b80b GIT binary patch literal 8960 zcmbtZ-ESP%b)WCqZ&DOZQI=J?px6(z4p;S?f%Z4+1cfit^;%l z-kE#n-0$D{-E+=uR;wiiKlQ)=a_1M{Rh0jt%;Zl;<{kWpTS%DFP?*ZJKy9e1TkwjM zGW&_n3ap6ok}RdN(i4r9S>uWi^ylj+D)?S(H_el$~O8D66q=${!YukvD7?9_6*6h#V0m9%btbo0(&Ks9e+Ld#9+^{=kfK7 z{5Jmb0{VE7IT*vD+_YE4%-S{DvV6zlYb6K2|Hn>Jb=t`4elO^`Y_u`j#IDdg;d6yH z<~byZ+Al=PV5H3@7=rRg?mx2fA7}H#=VfY?sa<13-0;M zZDA|iOe=c)Qh%R1@G)oCokc^Tfkd$@YP>kLZI?!3!fx~CR%;N|bz%CQX4?~1$BTA) zOjz8D2E5y%l^^}+-Ne7#iTCl6Ksoj_K9?v%?8+h5A0-+}&Z9)lOLUaHi4rX4&O{y&Sc`EbO8kf0?iA2GM{wNetV(*Ky=iX;Gc0Z^uW0v3z_xWIlV_HV*8_99oH*Xo;QZiNTQH zwGV9-Jq`=0Iv_3VX;*Oovf`mWEV9y}dQd=}0!F8PLSsaVF_pIUv&y`3P)v$-L2id%PD4Gu&zRJ^l*X683h`YesK=Pw1`tms5(u z0L7{6EAHC*+NOJVZS#h^vAVJ;%m+Ij2WK_IXut0XyX!q@KJ-ITTz%N`q%Dx*)%0jY z_3{|4!ss-^y+(1X&0I;SjpFv`RaiTo-`Op|CgA!D6U4h`ViI6!V_Csc-mbQhRJsdZo^7XPcLlvpS$GviYY*aqGb zAngEh!&*{E>^8LLP$zWn&H%xLVe?acgxCall(?PU|JMoBo*K?E>jc!=8EWNIqSgj# z3-Oyd2EPuqIf0QmNTf?YK=nM9B~1vHPrpL)xhnKMZ~t?ZU#2y?kfSqi<3xO}*H80{ zs2~ch`_T%l6E12iumCpPtJgO!U*A~%Kvb?b2Vv+pyH|q&=Zhp<-Mq1}dgq3Fb!~ll z;{#F7vbQ$xyeBH#ewVGJM@A^B8>^ce?&|xSs~hXfH^u3d8_R3!?(*&18@Jwf@2vgk z_RUqHcbX4HIb-8x?lt+V)JaL|j2rlSo+xblL2$bl?2Gc2A3g9x&xO6r-+-LNKL1_3 z9YsD)hW&=UMPv4uP+^UvD}kAg1~(HldD{!4(Mceej(?vT{Q)IMX`Y`$**gBi7myJ2 zYkEz8PG8UsC_qV_)n`=;@UId4zkLk;`94e>6Qi^0hyqm5cuw>GflT~2->d)tssz6T zIiQ~y{IkeT43rm`4v>T5(Ml0iZystOUmH0nL1H7X1Nh}(MXEp%<<+E^6rciDrUEso z0!3y+1xiVYz@MwY%>KVY1&u@*h&!6(kcRJ!)rcD2zO>C79-OQVIC)Q^mKC z2y>;`?e6odR6yq~zedR_CGR2;#U^jv^tx>jXQkPPdHRshdF{4@F27E4Hz--7WD2I0 zF+uY`pvFv(0fziXB>RVy5Xyw@^+P}CbsJXN!8f3je1?zkA38`z=uELWdmNv~={Z}V zZ$hVq#&dN3Ju-1Hjm};2qd~|(ou0r323f-nh5?YKCLRCqWk&hD^a>dhC?_m~=n1bD zObtS;4%#*{(zW-01H{2Ym!?7ih6YIf7~-G?i+ym{zc!$#53sh{!f_L#z&$wk=s$qaT@EIjnwzOuP*k(8)`1z-bTa(IdcV*LVzY z!PUA9c$f*3oAfE2Be3R+=Q%}%YGY&09|BOYWhJ0W!!pwo z`vgGQk8o}SPAHvM8;s<5YE*H&&Oi$zGSghGmDfP<==_gmJsXyxx?9L=(otMS>t!6&f163zQxc_2!sP#)Nx+*r(ndKcCxm|ZO1qWa zDl6iz^r;5BRUOtagVJLaZ6J?hzI;|0&ZK4JQ?Vk-DpwGaz^rUmooxB8(tQhF*(|)W zO=w01nt@P)Si02&z(-Osxqb2>-Sl5=+l zQ6UgNnoUb5Oy(nXGPoV7k(Eu&hI+^`c|<-hJ)y<~RW3Q3k~dQ}q1f@L;!!@Kj`A^{ z7?R1>pC~2u?!sLDT|~x8awStWt7NN3Iu{B%9JDBQy&f;D^acUL2qFj0YqoPsPFx!W z=sb35x8IJ>N{5DenlQIyJS#4(=I38rS)ZV}&CfMq_nRC{8H$1j2RCwidpv-`@(z-? zluf{6PJAk*N!VR9FzV||s2071ye>j_I+@@M&8}J=Ud`fmHAr4KcQbvs; zWE-_?xE|t}SlGtQI4>#g#6x)`{Oe;Ufe@i6L$d4 zGGn+npo=621%nPLo-}{QT6Lsmo;i_St~R6A4lb~Ag`DP?P2L4A;+L<<+YeG%y76l9 zUcXP7Y`Mcrmo6a+Or;)+U${MPFDpADSj_r^=y%t9X=9)VsZBH&+^@YpIVKolxUIm1 zQ;GF@f$w8bDM#F=H+id3-J;2iI89iRj6!i{RN`i1N$s4FKaZMqJ7t6eYDF@_BfJSK zU2WW7rPJA_H!ANgH=;jpHiMS<|UvYC96%z9?Nfa|V zsrzXuE5Sf3lX(^uNeaFHIf8iwm|fI1Y9l~bIisL>Ox#uH4Y8T`#kj#73}BPmOU(TQOWDlm~w!E7>;isYCIn>(xU&ro{WRB+{3`97}f zt|%XUmW^UQV~x%GsP055nMsSdB__?JUxxP_$x|c*`(?mG|GpQD zZJ4)+cVk1IyT+w{w;h|8B7b`uMO&RdGP`}RE!2zgX=uo#W)P*tr==E#gRK!c`Cj;L zJToN5p|4^7g?}5kpPpD(M_>t4(k`9{4ilWyAJ97Bp`~l_5+#{Fm zzUSh0{UKG5B$D85WEYIp>tm`hE7#ri9!6d_1m%XJoR-iy5Xh&8v*8|BmV_9|7gyR69(zvY~Yq_ ziZHf%y?~QWaElV!75)h&v^EnwD_3rg%1TIx_e1I%OJ>ZM7N41aasFHL=jRul7jQB4 zrKK9`!k3k|m9H$d@>S(!S*JeBe?j#xA`vDq(%;X9PumsvTj_`_+}Z%OTqjvM+wtgV l2Al%E^599AfrfL^I7Q(~x|VNYgO!lh*U0GqIXVQi{{_={8O#6x literal 0 HcmV?d00001 diff --git a/src/model/__pycache__/yolo_detector.cpython-39.pyc b/src/model/__pycache__/yolo_detector.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a67f54a06b6400e5a637a096bd6ebf4e3efd2097 GIT binary patch literal 7104 zcmbtZOOG7ab*^_+S3l>$;mnYx%Z?JT6>B6h4?7crro@LOgBdd9D74yRw`=-V_f)gG zx_PUboas~xK`1S8 zb?UzFJ?Gx@opaB))v7Jwx%BJ5KluHoB>gvK79SOv>v*F+2$Nb8lbI68Em?e(mWr<$ zXal{a%T%WY#=vZu1FK~XO0CksZrOu!t1Rl2pwgDPbFqP zQ(6tCJdm2!{Z5>!?}bOymZs@Y&&9_?<~rW!A3;0` z{7Nl_$>8^#r+NB@&QejTmK&vjN|RgBPJ=h4dYmgb|d`g4sn*b-`&+4^g>%cxyptEgQO65B{C#K}xv zlk2vFhl<|XcO#EE>=j9NNRnCMOp=|MB&qAp|A2Ko|MmT*lG=kI^MdvWa!u{dFzouw z3p-wi&d2Pbw1)Qv-so>Z5+#w^QY;@Uv3epSttB$Q9_xwxOg=VFB&1DBqsB_4zII$< zaw2~reW{$tvJ|exCR0w;ViZ}5?SA>Vg5Hd2A$9=0t9fq&-|DPCGsQXjCm{!7=)=&D zeK+tYo)AkRkfK46c#rp@?IK5g=j?Agbhw;&m<=K-oUR{u&gHF9+&tB5E~stiMj0B3|)pT>x`p#I5;&CM}A@-4TIs6?}eCf^39JToPB4+y|-fS z`XRc}rlU4nbgFg6%uUtKXgsN9ZJGDP?|9qJPP4`<*r(JChP@s}NlQI1zJs*KQ>)z$ z-GSF`rv}w}oGaAS@D8C<5w|F#3+_|R54*!uhn98@coh|`iXXMzC$1m3`_Mcio3N>- zTI2;?UPco!@h>I5f%MBixv{7$CQH)i@?zo=^2e@0(L2F))J6f0!63{d8(3X5V7khICvyk&f*XX-}J4)6&$Q zmg6$&DpdDOo>t;2aCt+9xpRHUP-iP;E|ORxxWftHR_XS z#H%L7=RTK@ekx1ThAd4ki#{5wkW5%LY0O%aWDMzJ7>jA+lg5Mm`;PPuQc&31#Qgf>RNOp^)52^#W{FI`N@hn!|noDLVw3W z963*YK*9cS95T|u{o$b(6+`X3T)QzGL&?Ko>;zEgY~|Yy8683fjx!~ZmD)$gOcAY4 z;6_o~X8>7|A{9(EmwRr9TF)Lz6Wl()8#O=@NN54)R*u*p~_fvj&)YawRqV0T^bf8pWKGTBkbTXUO>~ z(wU%iA?ucQcyxHgYobXw9sm>cE<_#kZ&H!!hcNS}p z*(Q982KzR(DHWYl*?V<2>aR!}4w)2v4@9yxr6ybQvTUdgMMKKKYbl#@4Zn)IDJ$Yv z`;B4N)Me$8Y$y$RU7lS0ZJYUm(5>+s< z$he8KIQfrLxCm?$xcE+v0SmbJP-r=@4=tGb8BXh-GBtpE$^z2CwMwi+TbG_GQ!A!i zDbJPQFql9v?L5`l$QeMgELTo+v@3t6ps)RzLXOz9lBoRO6SZ&kOK{cTtPz@Ftu`=R zg=@1&9=S_O!<3Ph^{(f%Ok#uA8DO32X=`1deFv+5^RLNMSOQU#vb zZW{tBK!JSK*|~SRWXv-}8Ve)7h$KaqBJux3xDY$|3P=*f5y)&i?=N7L^W(#JzvCR? zJ-mYNm7f$dQTv6XFo(~`ed0O7ON5&tR(`5aeqyM-I698Y1s22y-YJwk)v3swiM_(P zpH1etkt5U`19SJN=()o$;Yhw9>n%Wa_m@hF z7o=$!v*6MJj3}R|z=sO6;D6mtG**JwWhG|j(J-5GSSLIuDtUZ6Q9qY?H_?g4B_N34 z0jsShI@Z!9-z{4QYRmAnELMg^wwOT{8Q+S?!AI*Vti~!QxG zTUtlUxYVz*WzpWmDt-lSz$JKF5!}=;LJcda0n4g{WqDd-t5+p<4qUBWCGSl+uCw!* zB-wi0Ksopmp4`TlI`Q{V3a`O~TLRww3V2rs-n}9EEwfF5ck4<06y90W`X}|!X5vd| zmuT=hF1U9V(e54}cDyKZT!-KdS0n)$+qnxMF7I@>X8JbwoeItV-7s%{Mz@{2!n8R! zuOgmsB<^0h4rS|_Pfz%F>P&At4-Y(`7^i!m0KwU5&$|9rN03G<7&;T=ayb$7gZxj45Jn4ZZ#EWy`md<;A(5*@#7@XQf1ffA5f>yiqH(wDA7+TYOT~of{HH|D zg6h?r>vpChZJ{E6ht?78g;X;#-|l}6Q&VjCeyijI%9d&nWUDpB0)(%s6;ODD@|!4K zlZ~$ojR3uZ+RN|+x0K1ZPC@!vb2hcKlpQqa2q$LKnoxpbPhUrNy6l(m5@XT7;9Mn2 zpD-P$PM+3*#MFO~ryB0P2>Ki&Ja|I5{RZ88j}?S)-bf6(IiCPMo={Fd26_v;MLG(A zK*q5NM5jBX@udj=Fp~mn&vC~QwH9urC9(%-fm^9)ClDwo0nf*d27MOYv-;xst_&cI%r>DQ6Hs5)!nf|)aM{Czj> z93T`Ytln8@zRN#^rd*@T6eqWAqHT@B-bCMu{cabT{lN&C{>ba4@|%+tzQ$$ToFz1LCTsJS^97=OWnP{y6ETK5zlT9mo%e#_K78}2O~G5<6@|g@ zDIyKRj1>O;f?fC-8fEsU%)P*keYm}ddfuht9tds;!mmG|0%0zEq|%1i#5T;27e>8k~MhzYcL9AL~P{}Oh^_sP<~^S8uDdCOl`S_8v+>>3Lk93tgK}}WwLq7fSj79 zsq6@MfWG7#H09g$x5_Dmn%eC);2NYtQBX_y EKQ$^5T>t<8 literal 0 HcmV?d00001 diff --git a/src/model/id_card_processor.py b/src/model/id_card_processor.py new file mode 100644 index 0000000..710f0cf --- /dev/null +++ b/src/model/id_card_processor.py @@ -0,0 +1,343 @@ +""" +ID Card Processor for background removal and preprocessing +""" +import cv2 +import numpy as np +from pathlib import Path +from typing import List, Optional, Dict, Any, Tuple +import logging +from .yolo_detector import YOLODetector + +class IDCardProcessor: + """ + ID Card Processor for background removal and preprocessing + """ + + def __init__(self, yolo_detector: Optional[YOLODetector] = None): + """ + Initialize ID Card Processor + + Args: + yolo_detector: YOLO detector instance + """ + self.yolo_detector = yolo_detector or YOLODetector() + self.logger = logging.getLogger(__name__) + + def remove_background(self, image: np.ndarray, method: str = 'grabcut') -> np.ndarray: + """ + Remove background from image + + Args: + image: Input image + method: Background removal method ('grabcut', 'threshold', 'contour') + + Returns: + Image with background removed + """ + if method == 'grabcut': + return self._grabcut_background_removal(image) + elif method == 'threshold': + return self._threshold_background_removal(image) + elif method == 'contour': + return self._contour_background_removal(image) + else: + self.logger.warning(f"Unknown method: {method}, using grabcut") + return self._grabcut_background_removal(image) + + def _grabcut_background_removal(self, image: np.ndarray) -> np.ndarray: + """ + Remove background using GrabCut algorithm + """ + try: + # Create mask + mask = np.zeros(image.shape[:2], np.uint8) + + # Create temporary arrays + bgd_model = np.zeros((1, 65), np.float64) + fgd_model = np.zeros((1, 65), np.float64) + + # Define rectangle (assuming ID card is in center) + height, width = image.shape[:2] + rect = (width//8, height//8, width*3//4, height*3//4) + + # Apply GrabCut + cv2.grabCut(image, mask, rect, bgd_model, fgd_model, 5, cv2.GC_INIT_WITH_RECT) + + # Create mask + mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8') + + # Apply mask + result = image * mask2[:, :, np.newaxis] + + return result + + except Exception as e: + self.logger.error(f"Error in grabcut background removal: {e}") + return image + + def _threshold_background_removal(self, image: np.ndarray) -> np.ndarray: + """ + Remove background using thresholding + """ + try: + # Convert to grayscale + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + + # Apply Gaussian blur + blurred = cv2.GaussianBlur(gray, (5, 5), 0) + + # Apply threshold + _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) + + # Find contours + contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + # Find largest contour (assumed to be the ID card) + if contours: + largest_contour = max(contours, key=cv2.contourArea) + + # Create mask + mask = np.zeros_like(gray) + cv2.fillPoly(mask, [largest_contour], 255) + + # Apply mask + result = cv2.bitwise_and(image, image, mask=mask) + return result + + return image + + except Exception as e: + self.logger.error(f"Error in threshold background removal: {e}") + return image + + def _contour_background_removal(self, image: np.ndarray) -> np.ndarray: + """ + Remove background using contour detection + """ + try: + # Convert to grayscale + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + + # Apply edge detection + edges = cv2.Canny(gray, 50, 150) + + # Find contours + contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + # Find largest contour + if contours: + largest_contour = max(contours, key=cv2.contourArea) + + # Approximate contour to get rectangle + epsilon = 0.02 * cv2.arcLength(largest_contour, True) + approx = cv2.approxPolyDP(largest_contour, epsilon, True) + + # Create mask + mask = np.zeros_like(gray) + cv2.fillPoly(mask, [approx], 255) + + # Apply mask + result = cv2.bitwise_and(image, image, mask=mask) + return result + + return image + + except Exception as e: + self.logger.error(f"Error in contour background removal: {e}") + return image + + def enhance_image(self, image: np.ndarray) -> np.ndarray: + """ + Enhance image quality for better OCR + """ + try: + # Convert to LAB color space + lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB) + + # Apply CLAHE to L channel + clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) + lab[:, :, 0] = clahe.apply(lab[:, :, 0]) + + # Convert back to BGR + enhanced = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR) + + # Apply slight Gaussian blur to reduce noise + enhanced = cv2.GaussianBlur(enhanced, (3, 3), 0) + + return enhanced + + except Exception as e: + self.logger.error(f"Error enhancing image: {e}") + return image + + def normalize_image(self, image: np.ndarray, target_size: Tuple[int, int] = (800, 600)) -> np.ndarray: + """ + Normalize image size and orientation + """ + try: + # Resize image + resized = cv2.resize(image, target_size, interpolation=cv2.INTER_AREA) + + # Convert to grayscale if needed + if len(resized.shape) == 3: + gray = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY) + else: + gray = resized + + # Apply histogram equalization + equalized = cv2.equalizeHist(gray) + + # Convert back to BGR for consistency + if len(image.shape) == 3: + result = cv2.cvtColor(equalized, cv2.COLOR_GRAY2BGR) + else: + result = equalized + + return result + + except Exception as e: + self.logger.error(f"Error normalizing image: {e}") + return image + + def process_id_card(self, image_path: Path, output_dir: Path, + remove_bg: bool = True, enhance: bool = True, + normalize: bool = True, target_size: Tuple[int, int] = (800, 600)) -> Dict[str, Any]: + """ + Process a single ID card image + + Args: + image_path: Path to input image + output_dir: Output directory + remove_bg: Whether to remove background + enhance: Whether to enhance image + normalize: Whether to normalize image + target_size: Target size for normalization + + Returns: + Processing results + """ + result = { + 'input_path': str(image_path), + 'output_paths': [], + 'success': False + } + + try: + # Load image + image = cv2.imread(str(image_path)) + if image is None: + self.logger.error(f"Could not load image: {image_path}") + return result + + # Create output filename + stem = image_path.stem + processed_path = output_dir / f"{stem}_processed.jpg" + + # Apply processing steps + processed_image = image.copy() + + if remove_bg: + self.logger.info(f"Removing background from {image_path.name}") + processed_image = self.remove_background(processed_image) + + if enhance: + self.logger.info(f"Enhancing {image_path.name}") + processed_image = self.enhance_image(processed_image) + + if normalize: + self.logger.info(f"Normalizing {image_path.name}") + processed_image = self.normalize_image(processed_image, target_size) + + # Save processed image + processed_path.parent.mkdir(parents=True, exist_ok=True) + cv2.imwrite(str(processed_path), processed_image) + result['output_paths'].append(str(processed_path)) + + result['success'] = True + self.logger.info(f"Processed {image_path.name}") + + except Exception as e: + self.logger.error(f"Error processing {image_path}: {e}") + + return result + + def batch_process_id_cards(self, input_dir: Path, output_dir: Path, + detect_first: bool = True, **kwargs) -> Dict[str, Any]: + """ + Process all ID card images in a directory + + Args: + input_dir: Input directory + output_dir: Output directory + detect_first: Whether to detect ID cards first using YOLO + **kwargs: Additional arguments for processing + + Returns: + Batch processing results + """ + # Create output directory + output_dir.mkdir(parents=True, exist_ok=True) + + if detect_first: + # First detect and crop ID cards + self.logger.info("Detecting and cropping ID cards...") + detection_results = self.yolo_detector.batch_process(input_dir, output_dir / "cropped") + + # Process cropped images + cropped_dir = output_dir / "cropped" + if cropped_dir.exists(): + self.logger.info("Processing cropped ID cards...") + return self._process_cropped_images(cropped_dir, output_dir / "processed", **kwargs) + else: + self.logger.warning("No cropped images found, processing original images") + return self._process_cropped_images(input_dir, output_dir / "processed", **kwargs) + else: + # Process original images directly + return self._process_cropped_images(input_dir, output_dir / "processed", **kwargs) + + def _process_cropped_images(self, input_dir: Path, output_dir: Path, **kwargs) -> Dict[str, Any]: + """ + Process cropped ID card images recursively + """ + # Get all image files recursively from input directory and subdirectories + image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'} + image_files = [] + + # Recursively find all image files + for file_path in input_dir.rglob('*'): + if file_path.is_file() and file_path.suffix.lower() in image_extensions: + image_files.append(file_path) + + if not image_files: + self.logger.error(f"No images found in {input_dir} and subdirectories") + return {'success': False, 'error': 'No images found'} + + self.logger.info(f"Processing {len(image_files)} images from {input_dir} and subdirectories") + + results = { + 'total_images': len(image_files), + 'processed_images': 0, + 'results': [] + } + + # Process each image + for i, image_path in enumerate(image_files): + self.logger.info(f"Processing {i+1}/{len(image_files)}: {image_path.name}") + + # Create subdirectory structure in output to match input structure + relative_path = image_path.relative_to(input_dir) + output_subdir = output_dir / relative_path.parent + output_subdir.mkdir(parents=True, exist_ok=True) + + result = self.process_id_card(image_path, output_subdir, **kwargs) + results['results'].append(result) + + if result['success']: + results['processed_images'] += 1 + + # Summary + self.logger.info(f"ID card processing completed:") + self.logger.info(f" - Total images: {results['total_images']}") + self.logger.info(f" - Processed: {results['processed_images']}") + + return results \ No newline at end of file diff --git a/src/model/yolo_detector.py b/src/model/yolo_detector.py new file mode 100644 index 0000000..b1974bf --- /dev/null +++ b/src/model/yolo_detector.py @@ -0,0 +1,266 @@ +""" +YOLO Detector for ID Card Detection and Cropping +""" +import cv2 +import numpy as np +from pathlib import Path +from typing import List, Tuple, Optional, Dict, Any +import logging +from ultralytics import YOLO +import torch + +class YOLODetector: + """ + YOLO-based detector for ID card detection and cropping + """ + + def __init__(self, model_path: Optional[str] = None, confidence: float = 0.5): + """ + Initialize YOLO detector + + Args: + model_path: Path to YOLO model file (.pt) + confidence: Confidence threshold for detection + """ + self.confidence = confidence + self.logger = logging.getLogger(__name__) + + # Initialize model + if model_path and Path(model_path).exists(): + self.model = YOLO(model_path) + self.logger.info(f"Loaded custom YOLO model from {model_path}") + else: + # Use pre-trained YOLO model for general object detection + self.model = YOLO('yolov8n.pt') + self.logger.info("Using pre-trained YOLOv8n model") + + # Set device + self.device = 'cuda' if torch.cuda.is_available() else 'cpu' + self.logger.info(f"Using device: {self.device}") + + def detect_id_cards(self, image_path: Path) -> List[Dict[str, Any]]: + """ + Detect ID cards in an image + + Args: + image_path: Path to image file + + Returns: + List of detection results with bounding boxes + """ + try: + # Load image + image = cv2.imread(str(image_path)) + if image is None: + self.logger.error(f"Could not load image: {image_path}") + return [] + + # Run detection + results = self.model(image, conf=self.confidence) + + detections = [] + for result in results: + boxes = result.boxes + if boxes is not None: + for box in boxes: + # Get coordinates + x1, y1, x2, y2 = box.xyxy[0].cpu().numpy() + confidence = float(box.conf[0]) + class_id = int(box.cls[0]) + class_name = self.model.names[class_id] + + detection = { + 'bbox': [int(x1), int(y1), int(x2), int(y2)], + 'confidence': confidence, + 'class_id': class_id, + 'class_name': class_name, + 'area': (x2 - x1) * (y2 - y1) + } + detections.append(detection) + + # Sort by confidence and area (prefer larger, more confident detections) + detections.sort(key=lambda x: (x['confidence'], x['area']), reverse=True) + + self.logger.info(f"Found {len(detections)} detections in {image_path.name}") + return detections + + except Exception as e: + self.logger.error(f"Error detecting ID cards in {image_path}: {e}") + return [] + + def crop_id_card(self, image_path: Path, bbox: List[int], + output_path: Optional[Path] = None, + padding: int = 10) -> Optional[np.ndarray]: + """ + Crop ID card from image using bounding box + + Args: + image_path: Path to input image + bbox: Bounding box [x1, y1, x2, y2] + output_path: Path to save cropped image + padding: Padding around the bounding box + + Returns: + Cropped image as numpy array + """ + try: + # Load image + image = cv2.imread(str(image_path)) + if image is None: + self.logger.error(f"Could not load image: {image_path}") + return None + + height, width = image.shape[:2] + x1, y1, x2, y2 = bbox + + # Add padding + x1 = max(0, x1 - padding) + y1 = max(0, y1 - padding) + x2 = min(width, x2 + padding) + y2 = min(height, y2 + padding) + + # Crop image + cropped = image[y1:y2, x1:x2] + + # Save if output path provided + if output_path: + output_path.parent.mkdir(parents=True, exist_ok=True) + cv2.imwrite(str(output_path), cropped) + self.logger.info(f"Saved cropped image to {output_path}") + + return cropped + + except Exception as e: + self.logger.error(f"Error cropping ID card from {image_path}: {e}") + return None + + def process_single_image(self, image_path: Path, output_dir: Path, + save_original: bool = False) -> Dict[str, Any]: + """ + Process a single image: detect and crop ID cards + + Args: + image_path: Path to input image + output_dir: Output directory for cropped images + save_original: Whether to save original image with bounding boxes + + Returns: + Processing results + """ + result = { + 'input_path': str(image_path), + 'detections': [], + 'cropped_paths': [], + 'success': False + } + + try: + # Detect ID cards + detections = self.detect_id_cards(image_path) + + if not detections: + self.logger.warning(f"No ID cards detected in {image_path.name}") + return result + + # Process each detection + for i, detection in enumerate(detections): + bbox = detection['bbox'] + + # Create output filename + stem = image_path.stem + suffix = f"_card_{i+1}.jpg" + output_path = output_dir / f"{stem}{suffix}" + + # Crop ID card + cropped = self.crop_id_card(image_path, bbox, output_path) + + if cropped is not None: + result['detections'].append(detection) + result['cropped_paths'].append(str(output_path)) + + # Save original with bounding boxes if requested + if save_original and detections: + image = cv2.imread(str(image_path)) + for detection in detections: + bbox = detection['bbox'] + cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2) + cv2.putText(image, f"{detection['confidence']:.2f}", + (bbox[0], bbox[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, + 0.5, (0, 255, 0), 2) + + annotated_path = output_dir / f"{image_path.stem}_annotated.jpg" + cv2.imwrite(str(annotated_path), image) + result['annotated_path'] = str(annotated_path) + + result['success'] = True + self.logger.info(f"Processed {image_path.name}: {len(result['cropped_paths'])} cards cropped") + + except Exception as e: + self.logger.error(f"Error processing {image_path}: {e}") + + return result + + def batch_process(self, input_dir: Path, output_dir: Path, + save_annotated: bool = False) -> Dict[str, Any]: + """ + Process all images in a directory and subdirectories + + Args: + input_dir: Input directory containing images + output_dir: Output directory for cropped images + save_annotated: Whether to save annotated images + + Returns: + Batch processing results + """ + # Create output directory + output_dir.mkdir(parents=True, exist_ok=True) + + # Get all image files recursively from input directory and subdirectories + image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'} + image_files = [] + + # Recursively find all image files + for file_path in input_dir.rglob('*'): + if file_path.is_file() and file_path.suffix.lower() in image_extensions: + image_files.append(file_path) + + if not image_files: + self.logger.error(f"No images found in {input_dir} and subdirectories") + return {'success': False, 'error': 'No images found'} + + self.logger.info(f"Processing {len(image_files)} images from {input_dir} and subdirectories") + + results = { + 'total_images': len(image_files), + 'processed_images': 0, + 'total_detections': 0, + 'total_cropped': 0, + 'results': [] + } + + # Process each image + for i, image_path in enumerate(image_files): + self.logger.info(f"Processing {i+1}/{len(image_files)}: {image_path.name}") + + # Create subdirectory structure in output to match input structure + relative_path = image_path.relative_to(input_dir) + output_subdir = output_dir / relative_path.parent + output_subdir.mkdir(parents=True, exist_ok=True) + + result = self.process_single_image(image_path, output_subdir, save_annotated) + results['results'].append(result) + + if result['success']: + results['processed_images'] += 1 + results['total_detections'] += len(result['detections']) + results['total_cropped'] += len(result['cropped_paths']) + + # Summary + self.logger.info(f"Batch processing completed:") + self.logger.info(f" - Total images: {results['total_images']}") + self.logger.info(f" - Processed: {results['processed_images']}") + self.logger.info(f" - Total detections: {results['total_detections']}") + self.logger.info(f" - Total cropped: {results['total_cropped']}") + + return results \ No newline at end of file diff --git a/src/utils.py b/src/utils.py new file mode 100644 index 0000000..fa7f80a --- /dev/null +++ b/src/utils.py @@ -0,0 +1,98 @@ +""" +Utility functions for data augmentation +""" +import os +import logging +from pathlib import Path +from typing import List, Tuple, Optional +import cv2 +import numpy as np +from PIL import Image + +def setup_logging(log_level: str = "INFO") -> logging.Logger: + """Setup logging configuration""" + logging.basicConfig( + level=getattr(logging, log_level.upper()), + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + handlers=[ + logging.FileHandler('data_augmentation.log'), + logging.StreamHandler() + ] + ) + return logging.getLogger(__name__) + +def get_image_files(directory: Path) -> List[Path]: + """Get all image files from directory""" + SUPPORTED_FORMATS = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff'] + + image_files = [] + if directory.exists(): + for ext in SUPPORTED_FORMATS: + image_files.extend(directory.glob(f"*{ext}")) + image_files.extend(directory.glob(f"*{ext.upper()}")) + return sorted(image_files) + +def validate_image(image_path: Path) -> bool: + """Validate if file is a valid image""" + try: + with Image.open(image_path) as img: + img.verify() + return True + except Exception: + return False + +def load_image(image_path: Path, target_size: Tuple[int, int] = None) -> Optional[np.ndarray]: + """Load and resize image""" + try: + # Load image using OpenCV + image = cv2.imread(str(image_path)) + if image is None: + return None + + # Convert BGR to RGB + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + + # Resize if target_size is provided + if target_size: + image = cv2.resize(image, target_size, interpolation=cv2.INTER_AREA) + + return image + except Exception as e: + print(f"Error loading image {image_path}: {e}") + return None + +def save_image(image: np.ndarray, output_path: Path, quality: int = 95) -> bool: + """Save image to file""" + try: + # Convert RGB to BGR for OpenCV + image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + + # Create output directory if it doesn't exist + output_path.parent.mkdir(parents=True, exist_ok=True) + + # Save image + cv2.imwrite(str(output_path), image_bgr, [cv2.IMWRITE_JPEG_QUALITY, quality]) + return True + except Exception as e: + print(f"Error saving image {output_path}: {e}") + return False + +def create_augmented_filename(original_path: Path, index: int, suffix: str = "aug") -> Path: + """Create filename for augmented image""" + stem = original_path.stem + suffix = f"_{suffix}_{index:02d}" + return original_path.parent / f"{stem}{suffix}{original_path.suffix}" + +def get_file_size_mb(file_path: Path) -> float: + """Get file size in MB""" + return file_path.stat().st_size / (1024 * 1024) + +def print_progress(current: int, total: int, prefix: str = "Progress"): + """Print progress bar""" + bar_length = 50 + filled_length = int(round(bar_length * current / float(total))) + percents = round(100.0 * current / float(total), 1) + bar = '=' * filled_length + '-' * (bar_length - filled_length) + print(f'\r{prefix}: [{bar}] {percents}% ({current}/{total})', end='') + if current == total: + print() \ No newline at end of file