Venket / scripts /generate_multiple_degradation_extreme_gallery.py

Add files using upload-large-folder tool

1a12965 verified 3 months ago

7.04 kB

#!/usr/bin/env python
"""Generate WORST-CASE samples: many degradations stacked at high strength on one template.
Like scripts/generate_multiple_degradation_gallery.py, but every augmentation
family's probability and strength range are pushed high (instead of using
configs/generate.yaml's mild "realistic batch" defaults), and placement/
rendering are also pushed toward their hardest settings (heavy misplacement,
small/faint signature). Most effects fire on most samples, so each sample
ends up with a large, clearly-visible stack of simultaneous degradations.
The caption band at the top of each image lists not just which effects
fired but their resolved *extent*: strength, plus any effect-specific
parameter (blob count, shadow type, JPEG quality, wrinkle count/amplitude,
etc.) alongside the resolved placement/rendering severity for that sample.
Usage:
    python scripts/generate_multiple_degradation_extreme_gallery.py
"""

from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from PIL import Image, ImageDraw, ImageFont

from src.config.loader import load_config
from src.generator import GeneratorContext
from src.visualization.visualizer import save_visualization
from scripts.generate_degradation_gallery import ALL_EFFECTS

ROOT = Path(__file__).resolve().parents[1]
OUT_DIR = ROOT / "multiple_degradation_extreme"
N_SAMPLES = 10
TEMPLATE_ID = "bank_doc_05"

EFFECT_PROBABILITY = 0.85   # near-certain per effect, but still leaves sample-to-sample variety
STRENGTH_MIN, STRENGTH_MAX = 0.7, 1.0


def _load_font(size: int):
    for candidate in (r"C:\Windows\Fonts\consola.ttf", r"C:\Windows\Fonts\arial.ttf"):
        if Path(candidate).exists():
            return ImageFont.truetype(candidate, size)
    return ImageFont.load_default()


def build_overrides() -> list[str]:
    overrides = [
        f"document.template={TEMPLATE_ID}",
        "generation.run_validation=false",
        # Worst-case placement: large rotation, small scale, frequent/strong misplacement.
        "placement.rotation_deg.min=-15", "placement.rotation_deg.max=15",
        "placement.scale_factor.min=0.3", "placement.scale_factor.max=0.6",
        "placement.margin=0.03",
        "placement.misplacement_probability=0.5",
        "placement.misplacement_strength.min=0.5", "placement.misplacement_strength.max=1.0",
        "placement.allow_partial_outside=true",
        # Worst-case rendering: faint, blurred, heavily feathered ink.
        "rendering.ink_opacity.min=0.3", "rendering.ink_opacity.max=0.55",
        "rendering.blur_sigma.min=0.6", "rendering.blur_sigma.max=1.4",
        "rendering.feather_px.min=1.5", "rendering.feather_px.max=3.0",
        "rendering.scanner_noise=0.05",
    ]
    for name in ALL_EFFECTS:
        overrides.append(f"augmentation.{name}.enabled=true")
        overrides.append(f"augmentation.{name}.probability={EFFECT_PROBABILITY}")
        if name == "jpeg_compression":
            overrides += [
                "augmentation.jpeg_compression.strength=0.9",
                "augmentation.jpeg_compression.min_quality=10",
                "augmentation.jpeg_compression.max_quality=30",
            ]
        else:
            overrides.append(f"augmentation.{name}.strength.min={STRENGTH_MIN}")
            overrides.append(f"augmentation.{name}.strength.max={STRENGTH_MAX}")
    return overrides


def _fmt_effect(name: str, rec: dict) -> str:
    params = rec.get("params") or {}
    if name == "jpeg_compression" and "quality" in params:
        return f"{name}=q{params['quality']}"
    tag = f"{name}={rec['strength']:.2f}"
    if "resolved_type" in params:
        tag += f"({params['resolved_type']})"
    elif "resolved_tone" in params:
        tag += f"({params['resolved_tone']})"
    elif name == "crumple" and "wrinkle_lines" in params:
        tag += f"({params['wrinkle_lines']}L,{params['displacement_amplitude_px']:.1f}px)"
    elif name == "coffee_stain" and "blob_count" in params:
        tag += f"({params['blob_count']}blobs)"
    elif name == "page_tilt" and "tilt_deg" in params:
        tag += f"({params['tilt_deg']:.1f}deg)"
    elif name == "dust" and "speck_count" in params:
        tag += f"({params['speck_count']}px)"
    elif name == "scanner_streaks" and "streak_count" in params:
        tag += f"({params['streak_count']}streaks)"
    return tag


def _annotate(path: Path, metadata: dict) -> list[str]:
    effects = metadata["augmentation"]["effects"]
    applied_names = [n for n, r in effects.items() if r["applied"]]
    applied_tags = [_fmt_effect(n, effects[n]) for n in applied_names]

    placement, rendering, difficulty = metadata["placement"], metadata["rendering"], metadata["difficulty"]
    header = (
        f"{metadata['document_id']}  |  difficulty={difficulty['score']:.3f} ({difficulty['tier']})  |  "
        f"placement: rot={placement['rotation_deg']:.1f}deg scale={placement['scale_factor']:.2f} "
        f"misplaced={placement['misplaced']}  |  "
        f"render: ink_opacity={rendering['ink_opacity']:.2f} blur_sigma={rendering['blur_sigma']:.2f}"
    )
    body = f"Applied ({len(applied_tags)}/{len(effects)}): " + (", ".join(applied_tags) if applied_tags else "none")

    img = Image.open(path).convert("RGB")
    font = _load_font(15)
    draw_probe = ImageDraw.Draw(img)
    max_w = img.width - 28

    # Pixel-accurate wrap of the (potentially long) applied-effects line.
    words = body.replace(", ", " ,SPLIT_ME ").split(" ,SPLIT_ME ")
    lines, current = [], ""
    for tok in words:
        candidate = (current + ", " + tok) if current else tok
        if draw_probe.textlength(candidate, font=font) <= max_w:
            current = candidate
        else:
            if current:
                lines.append(current)
            current = tok
    if current:
        lines.append(current)

    all_lines = [header] + lines
    line_h = 22
    band_h = 14 + line_h * len(all_lines)

    canvas = Image.new("RGB", (img.width, img.height + band_h), "white")
    canvas.paste(img, (0, band_h))
    draw = ImageDraw.Draw(canvas)
    draw.rectangle([0, 0, canvas.width, band_h], fill=(60, 18, 18))
    y = 8
    for i, line in enumerate(all_lines):
        draw.text((14, y), line, fill="white" if i == 0 else (255, 214, 194), font=font)
        y += line_h
    canvas.save(path)
    return applied_names


def main() -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    config = load_config(overrides=build_overrides())
    ctx = GeneratorContext(config)

    for i in range(1, N_SAMPLES + 1):
        metadata, final_image, _ = ctx.run_sample(i)
        out_path = OUT_DIR / f"sample_{i}.png"
        save_visualization(out_path, final_image, metadata)
        applied = _annotate(out_path, metadata)
        print(f"[sample_{i}] {len(applied)}/{len(metadata['augmentation']['effects'])} effects -> {', '.join(applied)}")

    print(f"\nWrote {N_SAMPLES} extreme samples -> {OUT_DIR}")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python
	"""Generate WORST-CASE samples: many degradations stacked at high strength on one template.

	Like scripts/generate_multiple_degradation_gallery.py, but every augmentation
	family's probability and strength range are pushed high (instead of using
	configs/generate.yaml's mild "realistic batch" defaults), and placement/
	rendering are also pushed toward their hardest settings (heavy misplacement,
	small/faint signature). Most effects fire on most samples, so each sample
	ends up with a large, clearly-visible stack of simultaneous degradations.

	The caption band at the top of each image lists not just which effects
	fired but their resolved extent: strength, plus any effect-specific
	parameter (blob count, shadow type, JPEG quality, wrinkle count/amplitude,
	etc.) alongside the resolved placement/rendering severity for that sample.

	Usage:
	python scripts/generate_multiple_degradation_extreme_gallery.py
	"""

	from __future__ import annotations

	import sys
	from pathlib import Path

	sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

	from PIL import Image, ImageDraw, ImageFont

	from src.config.loader import load_config
	from src.generator import GeneratorContext
	from src.visualization.visualizer import save_visualization
	from scripts.generate_degradation_gallery import ALL_EFFECTS

	ROOT = Path(__file__).resolve().parents[1]
	OUT_DIR = ROOT / "multiple_degradation_extreme"
	N_SAMPLES = 10
	TEMPLATE_ID = "bank_doc_05"

	EFFECT_PROBABILITY = 0.85 # near-certain per effect, but still leaves sample-to-sample variety
	STRENGTH_MIN, STRENGTH_MAX = 0.7, 1.0


	def _load_font(size: int):
	for candidate in (r"C:\Windows\Fonts\consola.ttf", r"C:\Windows\Fonts\arial.ttf"):
	if Path(candidate).exists():
	return ImageFont.truetype(candidate, size)
	return ImageFont.load_default()


	def build_overrides() -> list[str]:
	overrides = [
	f"document.template={TEMPLATE_ID}",
	"generation.run_validation=false",
	# Worst-case placement: large rotation, small scale, frequent/strong misplacement.
	"placement.rotation_deg.min=-15", "placement.rotation_deg.max=15",
	"placement.scale_factor.min=0.3", "placement.scale_factor.max=0.6",
	"placement.margin=0.03",
	"placement.misplacement_probability=0.5",
	"placement.misplacement_strength.min=0.5", "placement.misplacement_strength.max=1.0",
	"placement.allow_partial_outside=true",
	# Worst-case rendering: faint, blurred, heavily feathered ink.
	"rendering.ink_opacity.min=0.3", "rendering.ink_opacity.max=0.55",
	"rendering.blur_sigma.min=0.6", "rendering.blur_sigma.max=1.4",
	"rendering.feather_px.min=1.5", "rendering.feather_px.max=3.0",
	"rendering.scanner_noise=0.05",
	]
	for name in ALL_EFFECTS:
	overrides.append(f"augmentation.{name}.enabled=true")
	overrides.append(f"augmentation.{name}.probability={EFFECT_PROBABILITY}")
	if name == "jpeg_compression":
	overrides += [
	"augmentation.jpeg_compression.strength=0.9",
	"augmentation.jpeg_compression.min_quality=10",
	"augmentation.jpeg_compression.max_quality=30",
	]
	else:
	overrides.append(f"augmentation.{name}.strength.min={STRENGTH_MIN}")
	overrides.append(f"augmentation.{name}.strength.max={STRENGTH_MAX}")
	return overrides


	def _fmt_effect(name: str, rec: dict) -> str:
	params = rec.get("params") or {}
	if name == "jpeg_compression" and "quality" in params:
	return f"{name}=q{params['quality']}"
	tag = f"{name}={rec['strength']:.2f}"
	if "resolved_type" in params:
	tag += f"({params['resolved_type']})"
	elif "resolved_tone" in params:
	tag += f"({params['resolved_tone']})"
	elif name == "crumple" and "wrinkle_lines" in params:
	tag += f"({params['wrinkle_lines']}L,{params['displacement_amplitude_px']:.1f}px)"
	elif name == "coffee_stain" and "blob_count" in params:
	tag += f"({params['blob_count']}blobs)"
	elif name == "page_tilt" and "tilt_deg" in params:
	tag += f"({params['tilt_deg']:.1f}deg)"
	elif name == "dust" and "speck_count" in params:
	tag += f"({params['speck_count']}px)"
	elif name == "scanner_streaks" and "streak_count" in params:
	tag += f"({params['streak_count']}streaks)"
	return tag


	def _annotate(path: Path, metadata: dict) -> list[str]:
	effects = metadata["augmentation"]["effects"]
	applied_names = [n for n, r in effects.items() if r["applied"]]
	applied_tags = [_fmt_effect(n, effects[n]) for n in applied_names]

	placement, rendering, difficulty = metadata["placement"], metadata["rendering"], metadata["difficulty"]
	header = (
	f"{metadata['document_id']} \| difficulty={difficulty['score']:.3f} ({difficulty['tier']}) \| "
	f"placement: rot={placement['rotation_deg']:.1f}deg scale={placement['scale_factor']:.2f} "
	f"misplaced={placement['misplaced']} \| "
	f"render: ink_opacity={rendering['ink_opacity']:.2f} blur_sigma={rendering['blur_sigma']:.2f}"
	)
	body = f"Applied ({len(applied_tags)}/{len(effects)}): " + (", ".join(applied_tags) if applied_tags else "none")

	img = Image.open(path).convert("RGB")
	font = _load_font(15)
	draw_probe = ImageDraw.Draw(img)
	max_w = img.width - 28

	# Pixel-accurate wrap of the (potentially long) applied-effects line.
	words = body.replace(", ", " ,SPLIT_ME ").split(" ,SPLIT_ME ")
	lines, current = [], ""
	for tok in words:
	candidate = (current + ", " + tok) if current else tok
	if draw_probe.textlength(candidate, font=font) <= max_w:
	current = candidate
	else:
	if current:
	lines.append(current)
	current = tok
	if current:
	lines.append(current)

	all_lines = [header] + lines
	line_h = 22
	band_h = 14 + line_h * len(all_lines)

	canvas = Image.new("RGB", (img.width, img.height + band_h), "white")
	canvas.paste(img, (0, band_h))
	draw = ImageDraw.Draw(canvas)
	draw.rectangle([0, 0, canvas.width, band_h], fill=(60, 18, 18))
	y = 8
	for i, line in enumerate(all_lines):
	draw.text((14, y), line, fill="white" if i == 0 else (255, 214, 194), font=font)
	y += line_h
	canvas.save(path)
	return applied_names


	def main() -> None:
	OUT_DIR.mkdir(parents=True, exist_ok=True)
	config = load_config(overrides=build_overrides())
	ctx = GeneratorContext(config)

	for i in range(1, N_SAMPLES + 1):
	metadata, final_image, _ = ctx.run_sample(i)
	out_path = OUT_DIR / f"sample_{i}.png"
	save_visualization(out_path, final_image, metadata)
	applied = _annotate(out_path, metadata)
	print(f"[sample_{i}] {len(applied)}/{len(metadata['augmentation']['effects'])} effects -> {', '.join(applied)}")

	print(f"\nWrote {N_SAMPLES} extreme samples -> {OUT_DIR}")


	if __name__ == "__main__":
	main()