diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0001.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0001.png new file mode 100644 index 0000000..d903f9b Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0001.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0002.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0002.png new file mode 100644 index 0000000..bb008d8 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0002.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0003.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0003.png new file mode 100644 index 0000000..aeda926 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0003.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0004.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0004.png new file mode 100644 index 0000000..998c27d Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0004.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0005.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0005.png new file mode 100644 index 0000000..68a19d4 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0005.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0006.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0006.png new file mode 100644 index 0000000..09a58a3 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0006.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0007.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0007.png new file mode 100644 index 0000000..663f3c9 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0007.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0008.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0008.png new file mode 100644 index 0000000..768b96e Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0008.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0009.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0009.png new file mode 100644 index 0000000..5013a14 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0009.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0010.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0010.png new file mode 100644 index 0000000..1ee7baf Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0010.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0011.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0011.png new file mode 100644 index 0000000..c6dce37 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0011.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0012.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0012.png new file mode 100644 index 0000000..bed469b Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0012.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0013.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0013.png new file mode 100644 index 0000000..37eac78 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0013.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0014.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0014.png new file mode 100644 index 0000000..8c2afc7 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0014.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0015.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0015.png new file mode 100644 index 0000000..ad446ad Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0015.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0016.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0016.png new file mode 100644 index 0000000..b79343f Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0016.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0017.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0017.png new file mode 100644 index 0000000..303b27e Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0017.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0018.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0018.png new file mode 100644 index 0000000..f535467 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0018.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0019.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0019.png new file mode 100644 index 0000000..995d3bc Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0019.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/grid_0020.png b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0020.png new file mode 100644 index 0000000..0e6bc5d Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/grid_0020.png differ diff --git a/generator/outputs/samples/final_comparison/p5_ddpm/training_sample.png b/generator/outputs/samples/final_comparison/p5_ddpm/training_sample.png new file mode 100644 index 0000000..ae3a2a4 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_ddpm/training_sample.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0001.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0001.png new file mode 100644 index 0000000..3a51111 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0001.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0002.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0002.png new file mode 100644 index 0000000..66f778b Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0002.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0003.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0003.png new file mode 100644 index 0000000..606bc27 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0003.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0004.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0004.png new file mode 100644 index 0000000..d03cc6b Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0004.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0005.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0005.png new file mode 100644 index 0000000..9ecf2d3 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0005.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0006.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0006.png new file mode 100644 index 0000000..60b5fa3 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0006.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0007.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0007.png new file mode 100644 index 0000000..3d80c6b Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0007.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0008.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0008.png new file mode 100644 index 0000000..44811b4 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0008.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0009.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0009.png new file mode 100644 index 0000000..27172ed Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0009.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0010.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0010.png new file mode 100644 index 0000000..ed22f62 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0010.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0011.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0011.png new file mode 100644 index 0000000..7ecdd6b Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0011.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0012.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0012.png new file mode 100644 index 0000000..3fe018c Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0012.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0013.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0013.png new file mode 100644 index 0000000..9d17011 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0013.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0014.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0014.png new file mode 100644 index 0000000..7c27f94 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0014.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0015.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0015.png new file mode 100644 index 0000000..2e06622 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0015.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0016.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0016.png new file mode 100644 index 0000000..5bf3eaf Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0016.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0017.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0017.png new file mode 100644 index 0000000..ed6f435 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0017.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0018.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0018.png new file mode 100644 index 0000000..adc9cc1 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0018.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0019.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0019.png new file mode 100644 index 0000000..f7624ac Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0019.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/grid_0020.png b/generator/outputs/samples/final_comparison/p5_gan/grid_0020.png new file mode 100644 index 0000000..1d48301 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/grid_0020.png differ diff --git a/generator/outputs/samples/final_comparison/p5_gan/training_sample.png b/generator/outputs/samples/final_comparison/p5_gan/training_sample.png new file mode 100644 index 0000000..1e4d001 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_gan/training_sample.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0001.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0001.png new file mode 100644 index 0000000..e6f96be Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0001.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0002.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0002.png new file mode 100644 index 0000000..650b4d7 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0002.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0003.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0003.png new file mode 100644 index 0000000..1079f82 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0003.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0004.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0004.png new file mode 100644 index 0000000..da4ab32 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0004.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0005.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0005.png new file mode 100644 index 0000000..6b6bb56 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0005.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0006.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0006.png new file mode 100644 index 0000000..cbbfe7c Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0006.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0007.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0007.png new file mode 100644 index 0000000..509a39e Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0007.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0008.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0008.png new file mode 100644 index 0000000..3ed1017 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0008.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0009.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0009.png new file mode 100644 index 0000000..f3e9a3b Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0009.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0010.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0010.png new file mode 100644 index 0000000..80c293d Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0010.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0011.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0011.png new file mode 100644 index 0000000..e1f6164 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0011.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0012.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0012.png new file mode 100644 index 0000000..c8d1450 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0012.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0013.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0013.png new file mode 100644 index 0000000..6517943 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0013.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0014.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0014.png new file mode 100644 index 0000000..68464b2 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0014.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0015.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0015.png new file mode 100644 index 0000000..dc8a7fe Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0015.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0016.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0016.png new file mode 100644 index 0000000..b2d0360 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0016.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0017.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0017.png new file mode 100644 index 0000000..7e9c3f8 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0017.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0018.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0018.png new file mode 100644 index 0000000..ed91313 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0018.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0019.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0019.png new file mode 100644 index 0000000..72afc23 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0019.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/grid_0020.png b/generator/outputs/samples/final_comparison/p5_vae/grid_0020.png new file mode 100644 index 0000000..d3dad23 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/grid_0020.png differ diff --git a/generator/outputs/samples/final_comparison/p5_vae/training_sample.png b/generator/outputs/samples/final_comparison/p5_vae/training_sample.png new file mode 100644 index 0000000..ac29b52 Binary files /dev/null and b/generator/outputs/samples/final_comparison/p5_vae/training_sample.png differ diff --git a/generator/tools/facecrop.py b/generator/tools/facecrop.py index 0773e01..cce8473 100644 --- a/generator/tools/facecrop.py +++ b/generator/tools/facecrop.py @@ -1,24 +1,7 @@ -#!/usr/bin/env python3 -""" -Pre-align face images using MTCNN landmarks + similarity transform. - -This is the generator-side counterpart to classifier/tools/facecrop.py. -Difference: classifier uses bbox crop+resize; the generator wants landmark-based -alignment so the eyes, nose, and mouth land at fixed pixel positions in every -training image — structurally consistent training data for the generator. - -Output mirrors the source layout exactly: - data/wiki/14/37591914.jpg -> cropped/generator/wiki/14/37591914.jpg - -Resumable: already-aligned images are skipped by default. - -Usage: - python generator/tools/facecrop.py - python generator/tools/facecrop.py --data-dir data --output-dir cropped/generator - python generator/tools/facecrop.py --sources wiki --device cpu - python generator/tools/facecrop.py --size 128 - python generator/tools/facecrop.py --no-skip-existing # reprocess everything -""" +# Pre-align face images using MTCNN landmarks + similarity transform. +# Generator-side counterpart to classifier/tools/facecrop.py — uses landmark-based alignment +# (not bbox crop) so eyes, nose, and mouth land at fixed pixel positions in every image. +# Usage: python generator/tools/facecrop.py [--data-dir data] [--output-dir cropped/generator] [--size 128] import argparse import sys import warnings @@ -32,7 +15,7 @@ sys.path.insert(0, str(ROOT)) # Generator trains on real images only (wiki). The other sources are AI-generated # and aren't used as training targets for the generator, so we don't align them # by default. Pass --sources to override. -SOURCES = ["wiki"] +SOURCES = ["wiki"] ALL_SOURCES = ["wiki", "inpainting", "text2img", "insight"] # Reference landmark positions for a 128px aligned face. @@ -50,44 +33,37 @@ _DETECTORS: dict[str, object] = {} def parse_args(): - p = argparse.ArgumentParser( - description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - p.add_argument("--data-dir", default="data", help="Source dataset root (default: data)") + p = argparse.ArgumentParser() + p.add_argument("--data-dir", default="data", help="Source dataset root (default: data)") p.add_argument("--output-dir", default="cropped/generator", help="Output root (default: cropped/generator)") - p.add_argument("--size", type=int, default=128, help="Output image size in px, square (default: 128)") - p.add_argument("--device", default=None, help="'cpu' or 'cuda'. Default: auto-detect") + p.add_argument("--size", type=int, default=128, help="Output image size in px, square (default: 128)") + p.add_argument("--device", default=None, help="'cpu' or 'cuda'. Default: auto-detect") p.add_argument("--sources", nargs="+", default=None, metavar="SOURCE", - help=f"Sources to process. Default: {', '.join(SOURCES)} (real images only). " - f"All available: {', '.join(ALL_SOURCES)}") + help=f"Sources to process. Default: {', '.join(SOURCES)}. All: {', '.join(ALL_SOURCES)}") p.add_argument("--skip-existing", dest="skip_existing", action="store_true", default=True, - help="Skip images already present in output-dir (default: on, resumable)") + help="Skip images already present in output-dir (default: on)") p.add_argument("--no-skip-existing", dest="skip_existing", action="store_false", help="Re-process all images even if already aligned") return p.parse_args() -# ── alignment helpers ───────────────────────────────────────────────────────── +# ── Alignment helpers ───────────────────────────────────────────────────────── +# Scale the 128px reference template to match the target size def _ref_landmarks(size: int): - """Reference landmarks scaled from the 128px template to `size`.""" import numpy as np scale = size / 128.0 - return np.asarray( - [(x * scale, y * scale) for x, y in REF_LANDMARKS_128], - dtype=np.float32, - ) + return np.asarray([(x * scale, y * scale) for x, y in REF_LANDMARKS_128], dtype=np.float32) +# Apply similarity transform so detected landmarks map to reference positions def _align_from_landmarks(img, landmarks, size: int): - """Apply similarity transform so detected landmarks map to reference positions.""" import numpy as np from PIL import Image as PILImage from skimage.transform import SimilarityTransform, warp - src = np.asarray(landmarks, dtype=np.float32) # 5x2 detected - dst = _ref_landmarks(size) # 5x2 reference + src = np.asarray(landmarks, dtype=np.float32) + dst = _ref_landmarks(size) try: tform = SimilarityTransform.from_estimate(src, dst) @@ -105,35 +81,28 @@ def _align_from_landmarks(img, landmarks, size: int): def _center_crop(img, size: int): from PIL import Image as PILImage - w, h = img.size - side = min(w, h) + w, h = img.size + side = min(w, h) left, top = (w - side) // 2, (h - side) // 2 return img.crop((left, top, left + side, top + side)).resize((size, size), PILImage.BILINEAR) +# Returns (standard, relaxed) MTCNN detectors; cached per device def _get_detectors(device: str): - """Return (standard, relaxed) MTCNN detectors with landmarks enabled.""" if device in _DETECTORS: return _DETECTORS[device] from facenet_pytorch import MTCNN - standard = MTCNN( - keep_all=False, select_largest=True, - min_face_size=15, - device=device, post_process=False, - ) - relaxed = MTCNN( - keep_all=False, select_largest=True, - min_face_size=10, - thresholds=[0.5, 0.6, 0.6], - device=device, post_process=False, - ) + standard = MTCNN(keep_all=False, select_largest=True, min_face_size=15, + device=device, post_process=False) + relaxed = MTCNN(keep_all=False, select_largest=True, min_face_size=10, + thresholds=[0.5, 0.6, 0.6], device=device, post_process=False) _DETECTORS[device] = (standard, relaxed) return standard, relaxed -# ── main ────────────────────────────────────────────────────────────────────── +# ── Main ────────────────────────────────────────────────────────────────────── def main(): args = parse_args() @@ -158,7 +127,7 @@ def main(): try: import facenet_pytorch # noqa: F401 - import skimage # noqa: F401 + import skimage # noqa: F401 except ImportError as exc: print(f"Error: missing dependency ({exc}).") print(" Run: pip install facenet-pytorch scikit-image") @@ -182,9 +151,7 @@ def main(): print(f"\nTotal images: {len(all_paths):,}\n") n_processed = n_skipped = n_error = 0 - src_stats: dict[str, dict] = { - s: {"aligned": 0, "retry": 0, "fallback": 0} for s in sources - } + src_stats: dict[str, dict] = {s: {"aligned": 0, "retry": 0, "fallback": 0} for s in sources} for img_path in tqdm(all_paths, desc="Aligning", unit="img"): rel = img_path.relative_to(data_dir) @@ -215,8 +182,8 @@ def main(): # Pass 2: upscale 2x and retry with relaxed thresholds if aligned is None: - w, h = img.size - img2x = img.resize((w * 2, h * 2), Image.BILINEAR) + w, h = img.size + img2x = img.resize((w * 2, h * 2), Image.BILINEAR) _, _, landmarks2 = relaxed.detect(img2x, landmarks=True) if landmarks2 is not None and len(landmarks2) > 0: lm_orig = [(x / 2, y / 2) for x, y in landmarks2[0]] @@ -253,9 +220,9 @@ def main(): print(f" {'Source':<12} {'pass-1':>8} {'pass-2':>8} {'fallback':>8} {'fallback%':>10}") print(f" {'─'*12} {'─'*8} {'─'*8} {'─'*8} {'─'*10}") for src in sources: - s = src_stats[src] + s = src_stats[src] total_src = s["aligned"] + s["retry"] + s["fallback"] - fb_pct = s["fallback"] / max(total_src, 1) + fb_pct = s["fallback"] / max(total_src, 1) print(f" {src:<12} {s['aligned']:>8,} {s['retry']:>8,} {s['fallback']:>8,} {fb_pct:>9.1%}") print(f"{'─' * 55}") print(f" Output: {output_dir.resolve()}") diff --git a/generator/tools/sampling.py b/generator/tools/sampling.py new file mode 100644 index 0000000..1a002e1 --- /dev/null +++ b/generator/tools/sampling.py @@ -0,0 +1,120 @@ +# Generate 4x4 sample grids from Phase 5 EMA checkpoints, matching training visualization. +# Usage: python generator/tools/sampling.py [--samples N] [--models p5_gan p5_vae p5_ddpm] +import argparse +import json +import sys +from pathlib import Path + +import torch +from torchvision.utils import save_image + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +MODELS_DIR = ROOT / "outputs" / "models" +CFG_DIR = ROOT / "configs" / "phase5" + +# final_ema first: best_ema is the lowest-FID snapshot, which for slowly-converging models +# (e.g. DDPM) can be saved while the EMA shadow is still close to random init. +PHASE5_RUNS = { + "p5_gan": {"type": "gan", "config": "p5_gan.json", "checkpoints": ["p5_gan_final_ema.pt", "p5_gan_best_ema.pt"]}, + "p5_vae": {"type": "vae", "config": "p5_vae.json", "checkpoints": ["p5_vae_final_ema.pt", "p5_vae_best_ema.pt"]}, + "p5_ddpm": {"type": "ddpm", "config": "p5_ddpm.json", "checkpoints": ["p5_ddpm_final_ema.pt", "p5_ddpm_best_ema.pt"]}, +} + + +def _load_cfg(name: str) -> dict: + with open(CFG_DIR / PHASE5_RUNS[name]["config"]) as f: + return json.load(f) + + +def _load_model(name: str, cfg: dict, device: torch.device): + from src.models import get_model + + result, _ = get_model(cfg) + # GAN returns (generator, critic); all others return the model directly + model = result[0] if isinstance(result, tuple) else result + + for ckpt_name in PHASE5_RUNS[name]["checkpoints"]: + ckpt_path = MODELS_DIR / ckpt_name + if not ckpt_path.exists(): + print(f" [{name}] {ckpt_name} not found, trying next") + continue + + state_dict = torch.load(ckpt_path, map_location=device, weights_only=True) + missing, unexpected = model.load_state_dict(state_dict, strict=False) + if missing or unexpected: + print(f" [{name}] {ckpt_name}: missing={len(missing)} unexpected={len(unexpected)} — trying next") + continue + print(f" [{name}] Loaded {ckpt_name}") + return model.to(device).eval() + + raise FileNotFoundError(f"No usable EMA checkpoint found for {name}") + + +# Returns a (16, C, H, W) tensor in [0, 1] ready for save_image with nrow=4 +@torch.no_grad() +def _generate_grid( + name: str, model, cfg: dict, device: torch.device, + *, truncation: float | None = None, +) -> torch.Tensor: + kind = PHASE5_RUNS[name]["type"] + image_size = cfg.get("image_size", 64) + + if kind == "gan": + z = torch.randn(16, cfg.get("latent_dim", 128), 1, 1, device=device) + if truncation is not None and truncation > 0: + z = z.clamp(-truncation, truncation) + imgs = model(z) + elif kind == "vae": + imgs = model.sample(16, device) + elif kind == "ddpm": + from src.training.diffusion import cosine_betas, linear_betas, make_alpha_bars, ddim_sample + T = cfg.get("T", 1000) + betas = (cosine_betas(T) if cfg.get("noise_schedule", "cosine") == "cosine" else linear_betas(T)).to(device) + alpha_bars = make_alpha_bars(betas) + # n_steps=50 mirrors the training-time preview; cfg's ddim_steps is for FID only + imgs = ddim_sample( + model, 16, image_size, alpha_bars, + n_steps=50, pred_type=cfg.get("pred_type", "eps"), + device=str(device), batch_size=16, + ) + else: + raise ValueError(f"Unknown model type: {kind}") + + return (imgs.clamp(-1, 1) + 1.0) / 2.0 + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--samples", type=int, default=10, help="Number of 4x4 grids per model") + p.add_argument("--output-dir", type=Path, + default=ROOT / "outputs" / "samples" / "final_comparison") + p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + p.add_argument("--truncation", type=float, default=None, + help="Optional GAN latent truncation (lower=less diversity but sharper)") + p.add_argument("--models", nargs="+", choices=list(PHASE5_RUNS.keys()), + default=list(PHASE5_RUNS.keys())) + args = p.parse_args() + device = torch.device(args.device) + args.output_dir.mkdir(parents=True, exist_ok=True) + + for name in args.models: + print(f"[{name}] Loading checkpoint...") + cfg = _load_cfg(name) + model = _load_model(name, cfg, device) + + out_dir = args.output_dir / name + out_dir.mkdir(exist_ok=True) + + print(f"[{name}] Generating {args.samples} grids -> {out_dir}") + for i in range(1, args.samples + 1): + grid = _generate_grid(name, model, cfg, device, truncation=args.truncation) + save_image(grid, out_dir / f"grid_{i:04d}.png", nrow=4) + + print("Done.") + + +if __name__ == "__main__": + main()