File size: 3,288 Bytes
059f915 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """
Batch Processing for VintageGAN
Specification Reference: Day 14
Process multiple images in a folder with vintage effects.
Author: VintageGAN Project
Date: 2024
"""
import argparse
from pathlib import Path
from tqdm import tqdm
import sys
sys.path.insert(0, str(Path(__file__).parent.parent))
from inference.apply_filter import VintageFilter
def process_folder(
input_dir: str,
output_dir: str,
checkpoint: str,
conditions: str = "soft_fade",
extensions: tuple = (".jpg", ".jpeg", ".png"),
device: str = "cuda",
image_size: int = 256,
):
"""
Process all images in a folder.
Args:
input_dir: Input directory
output_dir: Output directory
checkpoint: Generator checkpoint path
conditions: Condition preset or custom values
extensions: Image file extensions to process
device: Device to use
"""
input_path = Path(input_dir)
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
# Find all images
image_files = []
for ext in extensions:
image_files.extend(input_path.glob(f"*{ext}"))
image_files.extend(input_path.glob(f"*{ext.upper()}"))
if len(image_files) == 0:
print(f"No images found in {input_dir}")
return
print(f"Found {len(image_files)} images")
# Initialize filter
print("Loading model...")
filter = VintageFilter(checkpoint=checkpoint, device=device, image_size=image_size)
# Process images
print(f"\nProcessing with preset: {conditions}")
for img_path in tqdm(image_files, desc="Processing"):
try:
# Apply filter
output = filter.apply(img_path, conditions)
# Save with same filename
output_file = output_path / img_path.name
output.save(output_file, quality=95)
except Exception as e:
print(f"\nError processing {img_path.name}: {e}")
continue
print(f"\n✓ Processed {len(image_files)} images")
print(f"✓ Saved to: {output_dir}")
def main():
"""Command-line interface for batch processing."""
parser = argparse.ArgumentParser(description="Batch process images with VintageGAN")
parser.add_argument("input_dir", type=str, help="Input directory")
parser.add_argument("output_dir", type=str, help="Output directory")
parser.add_argument(
"--checkpoint", type=str, required=True, help="Path to generator checkpoint"
)
parser.add_argument(
"--preset", type=str, default="soft_fade", help="Preset condition name"
)
parser.add_argument(
"--device",
type=str,
default="cuda" if __import__("torch").cuda.is_available() else "cpu",
)
parser.add_argument("--image-size", type=int, default=256)
args = parser.parse_args()
print("=" * 60)
print("VINTAGEGAN BATCH PROCESSING")
print("=" * 60)
process_folder(
args.input_dir,
args.output_dir,
args.checkpoint,
args.preset,
device=args.device,
image_size=args.image_size,
)
print("=" * 60)
if __name__ == "__main__":
main()
|