Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,522 Bytes
4450790 |
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 |
from pathlib import Path
from PIL import Image
from PIL.PngImagePlugin import PngImageFile, PngInfo
import json
from pprint import pprint
import argparse
from rich.console import Console
from rich.progress import Progress
from rich_argparse import RichHelpFormatter
def parse_a111(params, verbose=False):
# params = [p.split(": ") for p in params.split("\n")]
params = params.split("\n")
prompt = params[0].strip()
neg = params[1].split(":")[1].strip()
settings = {}
try:
settings = {
s.split(":")[0].strip(): s.split(":")[1].strip()
for s in params[2].split(",")
}
except IndexError:
settings = {"raw": params[2].strip()}
if verbose:
print(f"PROMPT: {prompt}")
print(f"NEG: {neg}")
print("SETTINGS:")
pprint(settings, indent=4)
return {"prompt": prompt, "negative": neg, "settings": settings}
import glob
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Crude metadata extractor from A111 pngs",
formatter_class=RichHelpFormatter
)
parser.add_argument("inputs", nargs="*", help="Input image files")
parser.add_argument("--output", help="Output JSON file")
parser.add_argument("-v", "--verbose", action="store_true", help="Verbose mode")
parser.add_argument(
"--glob", help="Enable glob pattern matching", metavar="PATTERN"
)
args = parser.parse_args()
# - checks
if not args.glob and not args.inputs:
parser.error("Either --glob flag or inputs must be provided.")
if args.glob:
glob_pattern = args.glob
try:
pattern_path = str(Path(glob_pattern).expanduser().resolve())
if not any(glob.glob(pattern_path)):
raise ValueError(f"No files found for glob pattern: {glob_pattern}")
except Exception as e:
console = Console()
console.print(
f"[bold red]Error: Invalid glob pattern '{glob_pattern}': {e}[/bold red]"
)
exit(1)
else:
glob_pattern = None
input_files = []
if glob_pattern:
input_files = list(glob.glob(str(Path(glob_pattern).expanduser().resolve())))
else:
input_files = [Path(p) for p in args.inputs]
console = Console()
console.print("Input Files:", style="bold", end=" ")
console.print(f"{len(input_files):03d} files", style="cyan")
# for input_file in args.inputs:
# console.print(f"- {input_file}", style="cyan")
console.print("\nOutput File:", style="bold", end=" ")
console.print(f"{Path(args.output).resolve().absolute()}", style="cyan")
with Progress(console=console, auto_refresh=True) as progress:
# files = Path(pth).rglob("*.png")
unique_info = {}
last = None
task = progress.add_task("[cyan]Extracting meta...", total=len(input_files) + 1)
for p in input_files:
im = Image.open(p)
parsed = parse_a111(im.info["parameters"], args.verbose)
if parsed != last:
unique_info[Path(p).stem] = parsed
last = parsed
progress.update(task, advance=1)
progress.refresh()
unique_info = json.dumps(unique_info, indent=4)
with open(args.output, "w") as f:
f.write(unique_info)
progress.update(task, advance=1)
progress.refresh()
console.print("\nProcessing completed!", style="bold green")
|