Merge branch 'timotheschmidt/add-weight-vector' of https://github.com/EmoTim/stylegan2-ada-pytorch into timotheschmidt/add-weight-vector

Author: EmoTim

age_estimator.py

@@ -0,0 +1,32 @@
+import insightface
+from insightface.app import FaceAnalysis
+import numpy as np
+from PIL import Image
+
+
+class AgeEstimator:
+    def __init__(
+        self, ctx_id: int = 0, det_size: tuple[int, int] = (1024, 1024)
+    ) -> None:
+        self.app = FaceAnalysis(name="buffalo_l")  # SOTA Model including age
+        self.app.prepare(ctx_id=ctx_id, det_size=det_size)  # 0 = gpu
+
+    def estimate_age(self, pil_img: Image.Image) -> float | None:
+        """
+        Estimate age from a PIL Image.
+
+        Args:
+            pil_img: PIL Image object
+
+        Returns:
+            Estimated age as float, or None if no face detected
+        """
+        img = np.array(pil_img)
+        faces = self.app.get(img)
+        if len(faces) == 0:
+            return None
+        return faces[0].age
+
+    def __call__(self, pil_img: Image.Image) -> float | None:
+        """Allow instance to be called directly."""
+        return self.estimate_age(pil_img)

alpha_patch_work.png

@@ -20,6 +20,7 @@
 import matplotlib.pyplot as plt
 
 import legacy
+from age_estimator import AgeEstimator
 
 # ----------------------------------------------------------------------------
 
@@ -224,6 +225,12 @@ def generate_images(
             f"Applying weight vector to style blocks {start_idx} to {end_idx} (inclusive)"
         )
 
+    # Initialize age estimator if using weight vector
+    age_estimator = None
+    if weight_vector is not None:
+        print("Initializing age estimator...")
+        age_estimator = AgeEstimator(ctx_id=0, det_size=(1024, 1024))
+
     # Generate images.
     all_images = []  # Store images for composite: list of lists (one per seed)
 
@@ -233,9 +240,31 @@ def generate_images(
         w = G.mapping(z, label)
 
         if weight_vec is not None:
-            # Generate images with weight modulation
-            seed_images = []
             start_idx, end_idx = style_range
+
+            # First, check age with alpha=0 if age estimator is available
+            if age_estimator is not None:
+                img = G.synthesis(w, noise_mode=noise_mode)
+                img = (
+                    (img.permute(0, 2, 3, 1) * 127.5 + 128)
+                    .clamp(0, 255)
+                    .to(torch.uint8)
+                )
+                img_array = img[0].cpu().numpy()
+                pil_img = PIL.Image.fromarray(img_array, "RGB")
+                estimated_age = age_estimator(pil_img)
+
+                if estimated_age is None:
+                    print(f"  ⚠️  Seed {seed}: No face detected, skipping...")
+                    continue
+                elif estimated_age <= 20:
+                    print(f"  ⚠️  Seed {seed}: Age {estimated_age:.1f} <= 20, skipping...")
+                    continue
+                else:
+                    print(f"  ✓ Seed {seed}: Age {estimated_age:.1f} > 20, proceeding...")
+
+            # If age check passed (or not applicable), generate all alphas
+            seed_images = []
             for alpha in alphas:
                 # Clone w to avoid modifying the original
                 w_modified = w.clone()
@@ -244,7 +273,7 @@ def generate_images(
                     start_idx : end_idx + 1, :
                 ].unsqueeze(0)
                 assert w_modified.shape[1:] == (G.num_ws, G.w_dim)
-                img = G.synthesis(w_modified, noise_mode=noise_mode)
+                img = G.synthesis(w_modified, truncation_psi=truncation_psi, noise_mode=noise_mode)
                 img = (
                     (img.permute(0, 2, 3, 1) * 127.5 + 128)
                     .clamp(0, 255)
@@ -264,7 +293,7 @@ def generate_images(
             all_images.append(seed_images)
         else:
             # Generate image without weight modulation
-            img = G.synthesis(w, noise_mode=noise_mode)
+            img = G.synthesis(w, truncation_psi=truncation_psi, noise_mode=noise_mode)
             img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
             PIL.Image.fromarray(img[0].cpu().numpy(), "RGB").save(
                 f"{outdir}/seed{seed:04d}.png"

generate.py

@@ -8,6 +8,7 @@ dependencies = [
     "click>=8.3.0",
     "imageio>=2.37.2",
     "imageio-ffmpeg>=0.6.0",
+    "insightface>=0.7.3",
     "matplotlib>=3.10.7",
     "ninja>=1.13.0",
     "numpy>=2.3.4",