276 lines
8.5 KiB
Python
276 lines
8.5 KiB
Python
"""
|
|
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() |