474 lines
16 KiB
Python
474 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import curses
|
|
import random
|
|
import shutil
|
|
import struct
|
|
import zlib
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
|
|
|
|
DEFAULT_SIZES = "320x240,640x480,1024x768,1600x900"
|
|
|
|
|
|
@dataclass
|
|
class GeneratorConfig:
|
|
output: Path
|
|
sizes: list[tuple[int, int]]
|
|
variants_per_size: int
|
|
folder_count: int
|
|
images_per_folder: int
|
|
folder_prefix: str
|
|
file_prefix: str
|
|
seed: int
|
|
clear_output: bool
|
|
|
|
|
|
def parse_sizes(value: str) -> list[tuple[int, int]]:
|
|
sizes: list[tuple[int, int]] = []
|
|
for item in value.split(","):
|
|
item = item.strip().lower()
|
|
if not item:
|
|
continue
|
|
if "x" not in item:
|
|
raise ValueError(f"Invalid size '{item}'. Use WIDTHxHEIGHT,WIDTHxHEIGHT")
|
|
width_text, height_text = item.split("x", 1)
|
|
try:
|
|
width = int(width_text)
|
|
height = int(height_text)
|
|
except ValueError as exc:
|
|
raise ValueError(
|
|
f"Invalid size '{item}'. Width and height must be integers."
|
|
) from exc
|
|
if width <= 0 or height <= 0:
|
|
raise ValueError(
|
|
f"Invalid size '{item}'. Width and height must be positive."
|
|
)
|
|
sizes.append((width, height))
|
|
if not sizes:
|
|
raise ValueError("At least one image size is required.")
|
|
return sizes
|
|
|
|
|
|
def build_png(width: int, height: int, variant_seed: int) -> bytes:
|
|
rng = random.Random(variant_seed)
|
|
base_r = rng.randint(20, 230)
|
|
base_g = rng.randint(20, 230)
|
|
base_b = rng.randint(20, 230)
|
|
accent_r = (base_r + rng.randint(40, 120)) % 256
|
|
accent_g = (base_g + rng.randint(40, 120)) % 256
|
|
accent_b = (base_b + rng.randint(40, 120)) % 256
|
|
block = max(12, min(width, height) // max(6, rng.randint(6, 12)))
|
|
raw = bytearray()
|
|
|
|
for y in range(height):
|
|
raw.append(0)
|
|
for x in range(width):
|
|
mix_x = x / max(1, width - 1)
|
|
mix_y = y / max(1, height - 1)
|
|
checker = ((x // block) + (y // block) + variant_seed) % 2
|
|
stripe = ((x // max(4, block // 2)) + variant_seed) % 3
|
|
noise = rng.randint(-12, 12)
|
|
|
|
r = int(base_r * (1 - mix_x) + accent_r * mix_x)
|
|
g = int(base_g * (1 - mix_y) + accent_g * mix_y)
|
|
b = int((base_b + accent_b) / 2)
|
|
|
|
if checker:
|
|
r = min(255, r + 18)
|
|
g = max(0, g - 10)
|
|
if stripe == 0:
|
|
b = min(255, b + 28)
|
|
elif stripe == 1:
|
|
r = max(0, r - 18)
|
|
|
|
raw.extend(
|
|
(
|
|
max(0, min(255, r + noise)),
|
|
max(0, min(255, g - noise)),
|
|
max(0, min(255, b + noise // 2)),
|
|
)
|
|
)
|
|
|
|
def chunk(chunk_type: bytes, data: bytes) -> bytes:
|
|
return (
|
|
struct.pack(">I", len(data))
|
|
+ chunk_type
|
|
+ data
|
|
+ struct.pack(">I", zlib.crc32(chunk_type + data) & 0xFFFFFFFF)
|
|
)
|
|
|
|
header = b"\x89PNG\r\n\x1a\n"
|
|
ihdr = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)
|
|
idat = zlib.compress(bytes(raw), level=6)
|
|
return header + chunk(b"IHDR", ihdr) + chunk(b"IDAT", idat) + chunk(b"IEND", b"")
|
|
|
|
|
|
def build_variants(
|
|
sizes: list[tuple[int, int]], variants_per_size: int, seed: int
|
|
) -> list[tuple[str, bytes]]:
|
|
variants: list[tuple[str, bytes]] = []
|
|
for size_index, (width, height) in enumerate(sizes):
|
|
for variant_index in range(variants_per_size):
|
|
variant_seed = seed + (size_index * 10_000) + variant_index
|
|
name = f"{width}x{height}-v{variant_index + 1}.png"
|
|
variants.append((name, build_png(width, height, variant_seed)))
|
|
return variants
|
|
|
|
|
|
def write_collection(
|
|
output_dir: Path,
|
|
folder_count: int,
|
|
images_per_folder: int,
|
|
folder_prefix: str,
|
|
file_prefix: str,
|
|
variants: list[tuple[str, bytes]],
|
|
progress: Callable[[int, int, int], None] | None = None,
|
|
) -> int:
|
|
total_files = 0
|
|
for folder_number in range(1, folder_count + 1):
|
|
folder = output_dir / f"{folder_prefix}-{folder_number:03d}"
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
|
|
for image_number in range(1, images_per_folder + 1):
|
|
variant_name, variant_bytes = variants[(image_number - 1) % len(variants)]
|
|
destination = folder / f"{file_prefix}-{image_number:05d}-{variant_name}"
|
|
destination.write_bytes(variant_bytes)
|
|
total_files += 1
|
|
|
|
if progress:
|
|
progress(folder_number, folder_count, total_files)
|
|
|
|
return total_files
|
|
|
|
|
|
def validate_config(config: GeneratorConfig) -> None:
|
|
if config.variants_per_size <= 0:
|
|
raise ValueError("Variants per size must be greater than 0.")
|
|
if config.folder_count <= 0:
|
|
raise ValueError("Folder count must be greater than 0.")
|
|
if config.images_per_folder <= 0:
|
|
raise ValueError("Images per folder must be greater than 0.")
|
|
if not config.folder_prefix.strip():
|
|
raise ValueError("Folder prefix cannot be empty.")
|
|
if not config.file_prefix.strip():
|
|
raise ValueError("File prefix cannot be empty.")
|
|
|
|
|
|
def generate_collection(
|
|
config: GeneratorConfig,
|
|
progress: Callable[[int, int, int], None] | None = None,
|
|
) -> dict[str, str | int]:
|
|
validate_config(config)
|
|
|
|
output_dir = config.output.expanduser().resolve()
|
|
if config.clear_output and output_dir.exists():
|
|
shutil.rmtree(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
variants = build_variants(config.sizes, config.variants_per_size, config.seed)
|
|
total_files = write_collection(
|
|
output_dir=output_dir,
|
|
folder_count=config.folder_count,
|
|
images_per_folder=config.images_per_folder,
|
|
folder_prefix=config.folder_prefix,
|
|
file_prefix=config.file_prefix,
|
|
variants=variants,
|
|
progress=progress,
|
|
)
|
|
|
|
return {
|
|
"output_dir": str(output_dir),
|
|
"unique_variants": len(variants),
|
|
"folders_created": config.folder_count,
|
|
"files_created": total_files,
|
|
"sizes_used": ", ".join(f"{w}x{h}" for w, h in config.sizes),
|
|
}
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate fake placeholder image collections for testing."
|
|
)
|
|
parser.add_argument("output", nargs="?", type=Path)
|
|
parser.add_argument("--sizes", default=DEFAULT_SIZES)
|
|
parser.add_argument("--variants-per-size", type=int, default=3)
|
|
parser.add_argument("--folder-count", type=int, default=5)
|
|
parser.add_argument("--images-per-folder", type=int, default=100)
|
|
parser.add_argument("--folder-prefix", default="batch")
|
|
parser.add_argument("--file-prefix", default="image")
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--clear-output", action="store_true")
|
|
parser.add_argument(
|
|
"--no-tui",
|
|
action="store_true",
|
|
help="Run in non-interactive mode using command-line options.",
|
|
)
|
|
return parser
|
|
|
|
|
|
def config_from_args(args: argparse.Namespace) -> GeneratorConfig:
|
|
if args.output is None:
|
|
raise ValueError("Output directory is required in --no-tui mode.")
|
|
return GeneratorConfig(
|
|
output=args.output,
|
|
sizes=parse_sizes(args.sizes),
|
|
variants_per_size=args.variants_per_size,
|
|
folder_count=args.folder_count,
|
|
images_per_folder=args.images_per_folder,
|
|
folder_prefix=args.folder_prefix,
|
|
file_prefix=args.file_prefix,
|
|
seed=args.seed,
|
|
clear_output=args.clear_output,
|
|
)
|
|
|
|
|
|
def cli_mode(args: argparse.Namespace) -> int:
|
|
try:
|
|
config = config_from_args(args)
|
|
summary = generate_collection(config)
|
|
except ValueError as exc:
|
|
raise SystemExit(str(exc)) from exc
|
|
|
|
print(f"Output directory: {summary['output_dir']}")
|
|
print(f"Unique variants: {summary['unique_variants']}")
|
|
print(f"Folders created: {summary['folders_created']}")
|
|
print(f"Files created: {summary['files_created']}")
|
|
print(f"Sizes used: {summary['sizes_used']}")
|
|
return 0
|
|
|
|
|
|
def truncate(text: str, width: int) -> str:
|
|
if width <= 0:
|
|
return ""
|
|
if len(text) <= width:
|
|
return text
|
|
if width <= 3:
|
|
return text[:width]
|
|
return text[: width - 3] + "..."
|
|
|
|
|
|
def prompt_text(
|
|
stdscr: curses.window, label: str, current_value: str, message: str
|
|
) -> str | None:
|
|
curses.curs_set(1)
|
|
height, width = stdscr.getmaxyx()
|
|
prompt = f"{label}: "
|
|
stdscr.move(height - 2, 0)
|
|
stdscr.clrtoeol()
|
|
stdscr.addstr(height - 2, 0, truncate(prompt + current_value, width - 1))
|
|
stdscr.move(height - 1, 0)
|
|
stdscr.clrtoeol()
|
|
stdscr.addstr(height - 1, 0, truncate(message, width - 1))
|
|
buffer = list(current_value)
|
|
position = len(buffer)
|
|
|
|
while True:
|
|
stdscr.move(height - 2, 0)
|
|
stdscr.clrtoeol()
|
|
display = prompt + "".join(buffer)
|
|
stdscr.addstr(height - 2, 0, truncate(display, width - 1))
|
|
cursor_x = min(len(prompt) + position, max(0, width - 2))
|
|
stdscr.move(height - 2, cursor_x)
|
|
key = stdscr.getch()
|
|
|
|
if key in (10, 13):
|
|
curses.curs_set(0)
|
|
return "".join(buffer).strip()
|
|
if key == 27:
|
|
curses.curs_set(0)
|
|
return None
|
|
if key in (curses.KEY_BACKSPACE, 127, 8):
|
|
if position > 0:
|
|
position -= 1
|
|
buffer.pop(position)
|
|
continue
|
|
if key == curses.KEY_LEFT and position > 0:
|
|
position -= 1
|
|
continue
|
|
if key == curses.KEY_RIGHT and position < len(buffer):
|
|
position += 1
|
|
continue
|
|
if 32 <= key <= 126:
|
|
buffer.insert(position, chr(key))
|
|
position += 1
|
|
|
|
|
|
def draw_screen(
|
|
stdscr: curses.window,
|
|
fields: list[tuple[str, str]],
|
|
selected: int,
|
|
message: str,
|
|
status_lines: list[str],
|
|
) -> None:
|
|
stdscr.erase()
|
|
height, width = stdscr.getmaxyx()
|
|
title = "Fake Media Generator"
|
|
subtitle = "Arrows to move, Enter to edit, Space to toggle clear-output, G to generate, Q to quit"
|
|
stdscr.addstr(0, 0, truncate(title, width - 1), curses.A_BOLD)
|
|
stdscr.addstr(1, 0, truncate(subtitle, width - 1))
|
|
|
|
for index, (label, value) in enumerate(fields):
|
|
y = 3 + index
|
|
if y >= height - 5:
|
|
break
|
|
prefix = ">" if index == selected else " "
|
|
line = f"{prefix} {label:<18} {value}"
|
|
attr = curses.A_REVERSE if index == selected else curses.A_NORMAL
|
|
stdscr.addstr(y, 0, truncate(line, width - 1), attr)
|
|
|
|
status_start = min(height - 4, 4 + len(fields))
|
|
for offset, line in enumerate(status_lines[-3:]):
|
|
y = status_start + offset
|
|
if y < height - 1:
|
|
stdscr.addstr(y, 0, truncate(line, width - 1))
|
|
|
|
stdscr.addstr(height - 1, 0, truncate(message, width - 1), curses.A_DIM)
|
|
stdscr.refresh()
|
|
|
|
|
|
def tui_mode(args: argparse.Namespace) -> int:
|
|
initial_output = (
|
|
str(args.output) if args.output else str(Path.cwd() / "fake-media-output")
|
|
)
|
|
state: dict[str, str | bool] = {
|
|
"Output Directory": initial_output,
|
|
"Sizes": args.sizes,
|
|
"Variants Per Size": str(args.variants_per_size),
|
|
"Folder Count": str(args.folder_count),
|
|
"Images Per Folder": str(args.images_per_folder),
|
|
"Folder Prefix": args.folder_prefix,
|
|
"File Prefix": args.file_prefix,
|
|
"Seed": str(args.seed),
|
|
"Clear Output": args.clear_output,
|
|
}
|
|
|
|
def run(stdscr: curses.window) -> int:
|
|
curses.curs_set(0)
|
|
stdscr.keypad(True)
|
|
|
|
field_order = [
|
|
"Output Directory",
|
|
"Sizes",
|
|
"Variants Per Size",
|
|
"Folder Count",
|
|
"Images Per Folder",
|
|
"Folder Prefix",
|
|
"File Prefix",
|
|
"Seed",
|
|
"Clear Output",
|
|
]
|
|
selected = 0
|
|
message = "Fill in the fields, then press G to generate."
|
|
status_lines: list[str] = []
|
|
|
|
while True:
|
|
fields = []
|
|
for key in field_order:
|
|
value = state[key]
|
|
if isinstance(value, bool):
|
|
display = "yes" if value else "no"
|
|
else:
|
|
display = value
|
|
fields.append((key, str(display)))
|
|
|
|
draw_screen(stdscr, fields, selected, message, status_lines)
|
|
key = stdscr.getch()
|
|
|
|
if key in (ord("q"), ord("Q")):
|
|
return 0
|
|
if key in (curses.KEY_UP, ord("k")):
|
|
selected = (selected - 1) % len(field_order)
|
|
continue
|
|
if key in (curses.KEY_DOWN, ord("j")):
|
|
selected = (selected + 1) % len(field_order)
|
|
continue
|
|
|
|
current_field = field_order[selected]
|
|
|
|
if key == ord(" ") and current_field == "Clear Output":
|
|
state[current_field] = not bool(state[current_field])
|
|
message = (
|
|
f"{current_field} set to {'yes' if state[current_field] else 'no'}."
|
|
)
|
|
continue
|
|
|
|
if key in (10, 13):
|
|
if current_field == "Clear Output":
|
|
state[current_field] = not bool(state[current_field])
|
|
message = f"{current_field} set to {'yes' if state[current_field] else 'no'}."
|
|
continue
|
|
|
|
current_value = str(state[current_field])
|
|
updated = prompt_text(
|
|
stdscr,
|
|
current_field,
|
|
current_value,
|
|
"Enter saves, Esc cancels",
|
|
)
|
|
if updated is not None and updated:
|
|
state[current_field] = updated
|
|
message = f"Updated {current_field}."
|
|
elif updated == "":
|
|
message = f"{current_field} cannot be empty."
|
|
else:
|
|
message = f"Canceled edit for {current_field}."
|
|
continue
|
|
|
|
if key in (ord("g"), ord("G")):
|
|
try:
|
|
config = GeneratorConfig(
|
|
output=Path(str(state["Output Directory"])),
|
|
sizes=parse_sizes(str(state["Sizes"])),
|
|
variants_per_size=int(str(state["Variants Per Size"])),
|
|
folder_count=int(str(state["Folder Count"])),
|
|
images_per_folder=int(str(state["Images Per Folder"])),
|
|
folder_prefix=str(state["Folder Prefix"]),
|
|
file_prefix=str(state["File Prefix"]),
|
|
seed=int(str(state["Seed"])),
|
|
clear_output=bool(state["Clear Output"]),
|
|
)
|
|
status_lines = ["Generating files..."]
|
|
|
|
def progress(
|
|
folder_number: int, folder_total: int, total_files: int
|
|
) -> None:
|
|
status_lines[:] = [
|
|
f"Output: {config.output.expanduser().resolve()}",
|
|
f"Progress: folder {folder_number}/{folder_total}",
|
|
f"Files written: {total_files}",
|
|
]
|
|
draw_screen(stdscr, fields, selected, message, status_lines)
|
|
|
|
summary = generate_collection(config, progress=progress)
|
|
status_lines = [
|
|
f"Output: {summary['output_dir']}",
|
|
f"Unique variants: {summary['unique_variants']}",
|
|
f"Files created: {summary['files_created']}",
|
|
]
|
|
message = "Generation complete. Press G to run again or Q to quit."
|
|
except ValueError as exc:
|
|
message = str(exc)
|
|
continue
|
|
|
|
return 0
|
|
|
|
return curses.wrapper(run)
|
|
|
|
|
|
def main() -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args()
|
|
if args.no_tui:
|
|
return cli_mode(args)
|
|
return tui_mode(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|