Update app.py

app.py

@@ -6,54 +6,111 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 import gradio as gr
 import numpy as np
 import torch
+from PIL import Image
+
+# Backend libraries for StyleGAN and face detection
 import legacy
 import dnnlib
+from training.networks import SynthesisNetwork
+from projector import project
+from facenet_pytorch import MTCNN
 
-print("Loading and converting TensorFlow model...")
+# --- Load All Models ---
+print("Loading all models...")
 device = torch.device("cpu")
+
+# Load StyleGAN2 Generator
 with open("stylegan2-ffhq-config-f.pkl", "rb") as f:
-    # This line uses legacy.py to load the TF model and convert it to PyTorch
-    G = legacy.load_network_pkl(f)["G_ema"].to(device)
-print("Model loaded and converted successfully.")
+    G = legacy.load_network_pkl(f)['G_ema'].to(device)
+print("StyleGAN2 model loaded.")
+
+# Load Face Detector (MTCNN)
+mtcnn = MTCNN(keep_all=False, device=device) # keep_all=False finds only the best face
+print("Face detector model loaded.")
 
-# Load gender direction vector
-print("Loading gender direction vector...")
+# Load Gender Direction Vector
 gender_direction = np.load("stylegan2directions/gender.npy")
 gender_direction = torch.from_numpy(gender_direction).to(torch.float32).to(device)
-print("Vector loaded successfully.")
+print("All models and vectors loaded successfully.")
+# -----------------------------------
 
-# This is the function that will be called by the Gradio interface
-def edit_gender(seed, strength):
-    seed = int(seed)  # Ensure seed is an integer
-    
-    # Generate latent code from the seed
-    rnd = np.random.RandomState(seed)
-    z = torch.from_numpy(rnd.randn(1, G.z_dim)).to(torch.float32).to(device)
+
+def edit_uploaded_face(uploaded_image, strength):
+    """
+    This function detects a face in an uploaded image, projects it, edits it, and returns the result.
+    """
+    if uploaded_image is None:
+        raise gr.Error("No image uploaded. Please upload an image containing a face.")
+
+    print("Detecting face in the uploaded image...")
+    # The 'uploaded_image' from Gradio is a PIL Image, which is what MTCNN needs.
+    # We need to convert it to RGB if it has an alpha channel (like PNGs)
+    input_image = uploaded_image.convert("RGB")
     
-    # Map the latent code to the intermediate space (W)
-    w = G.mapping(z, None, truncation_psi=0.7)
+    # Detect face and get bounding box
+    boxes, _ = mtcnn.detect(input_image)
+
+    # Handle case where no face is detected
+    if boxes is None:
+        raise gr.Error("Could not detect a face. Please try a clearer picture or one where the face is more prominent.")
     
-    # Apply the gender direction transformation
-    w = w + gender_direction * strength
+    # --- Crop the image to the detected face ---
+    # boxes[0] contains the coordinates [x1, y1, x2, y2]
+    face_box = boxes[0]
     
-    # Synthesize the image from the modified W
-    img = G.synthesis(w, noise_mode="const")
+    # Add some padding to the crop to ensure the whole head is included
+    padding_x = (face_box[2] - face_box[0]) * 0.2
+    padding_y = (face_box[3] - face_box[1]) * 0.2
+    face_box[0] = max(0, face_box[0] - padding_x)
+    face_box[1] = max(0, face_box[1] - padding_y)
+    face_box[2] = min(input_image.width, face_box[2] + padding_x)
+    face_box[3] = min(input_image.height, face_box[3] + padding_y)
+
+    cropped_face = input_image.crop(face_box)
+    print("Face detected and cropped.")
+
+    # --- Run GAN Inversion on the CROPPED face ---
+    print("Projecting the face into the model's latent space...")
+    # This can be slow, especially on CPU. num_steps=100 is a good compromise for web apps.
+    projected_w = project(
+        G,
+        cropped_face, # Use the cropped face here
+        num_steps=100,
+        device=device,
+        verbose=False # Set to True for more detailed projection logs
+    )
+    print("Image projected successfully.")
+
+    # --- Apply Edit and Synthesize New Face ---
+    w_to_edit = projected_w[0]
+    w_edited = w_to_edit + gender_direction * strength
+    w_edited = w_edited.unsqueeze(0)
+
+    print("Synthesizing new image...")
+    img_out = G.synthesis(w_edited, noise_mode='const')
     
-    # Post-process the image for display
-    img = (img.clamp(-1, 1) + 1) * 127.5
-    img = img.permute(0, 2, 3, 1)[0].cpu().numpy().astype(np.uint8)
+    # Post-process for display
+    img_out = (img_out.clamp(-1, 1) + 1) * 127.5
+    img_out = img_out.permute(0, 2, 3, 1)[0].cpu().numpy().astype(np.uint8)
+    print("Processing complete.")
     
-    return img
+    return img_out
 
-# Create and launch the Gradio interface
+# --- Create the Gradio Interface ---
+# This interface now has a robust backend ready to be used as an API.
 gr.Interface(
-    fn=edit_gender,
+    fn=edit_uploaded_face,
     inputs=[
-        gr.Slider(0, 10000, step=1, value=1234, label="Random Seed"),
+        gr.Image(label="Upload Image With Face", type="pil"),
         gr.Slider(-5, 5, step=0.1, value=0, label="Gender Strength (← Feminine | Masculine →)")
     ],
-    outputs=gr.Image(label="Generated Face"),
-    title="StyleGAN2 Gender Editor",
-    description="Move the slider to change the gender expression of the generated face. Change the seed to get a new face.",
-    allow_flagging="never"
+    outputs=gr.Image(label="Edited Face"),
+    title="Face Editor Backend",
+    description="This engine detects a face in the uploaded image, then edits its gender expression. It is ready to be used as an API.",
+    allow_flagging="never",
+    examples=[
+        ["stylegan2directions/obama.jpg", 0],
+        ["stylegan2directions/obama.jpg", 3.5],
+        ["stylegan2directions/obama.jpg", -3.5],
+    ]
 ).launch()