File size: 2,461 Bytes
bff58bd |
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 |
import os
import argparse
from tqdm import tqdm
from armor import ArmorGenerator, ArmorConfig
def batch_armor_images(input_dir, output_dir, strength=2.5):
"""
Applies Ghostprint armor to a directory of images.
Args:
input_dir (str): Directory containing the original images.
output_dir (str): Directory to save the armored images.
strength (float): The strength of the armor to apply (1.0 to 5.0).
"""
if not os.path.isdir(input_dir):
print(f"Error: Input directory '{input_dir}' not found.")
return
os.makedirs(output_dir, exist_ok=True)
# Initialize the ArmorGenerator
# Using default ArmorConfig but can be customized if needed
armor_generator = ArmorGenerator()
armor_config = ArmorConfig()
image_files = sorted([f for f in os.listdir(input_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
if not image_files:
print(f"No images found in '{input_dir}'.")
return
print(f"Applying armor with strength {strength} to {len(image_files)} images...")
for filename in tqdm(image_files, desc="Armoring images"):
input_path = os.path.join(input_dir, filename)
output_path = os.path.join(output_dir, filename)
try:
armor_generator.apply_to_image(
image_path=input_path,
out_path=output_path,
cfg=armor_config,
strength=strength,
return_metrics=False # We don't need metrics for batch processing
)
except Exception as e:
print(f"\nFailed to process {filename}: {e}")
print("Batch armoring complete.")
print(f"Armored images saved in '{output_dir}'.")
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="Batch apply Ghostprint armor to images.")
parser.add_argument(
'--input-dir',
type=str,
default='placeholder_images',
help='Directory containing the original images.'
)
parser.add_argument(
'--output-dir',
type=str,
default='armored_images',
help='Directory to save the armored images.'
)
parser.add_argument(
'--strength',
type=float,
default=2.5,
help='Armor strength (e.g., 1.0 to 5.0). 2.5 is a good starting point.'
)
args = parser.parse_args()
batch_armor_images(args.input_dir, args.output_dir, args.strength)
|