Update app.py

app.py

@@ -46,51 +46,77 @@ def upscale_bg_tiled(frame, tile_size=128, overlap=16):
     return upscaled_img
 
 def restore_face_core(img, fidelity=0.5):
-    """CodeFormer Face Restoration with Dynamic Input & Type Support"""
-    # 1. Pre-process (CodeFormer expects 512x512)
+    """CodeFormer Face Restoration with Correct Math Normalization"""
     img_512 = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
-    img_in = np.transpose(cv2.cvtColor(img_512, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0, (2, 0, 1))[np.newaxis, :]
+    img_rgb = cv2.cvtColor(img_512, cv2.COLOR_BGR2RGB).astype(np.float32)
     
-    # 2. Map Inputs and Match Data Types Dynamically
-    inputs_info = face_session.get_inputs()
-    input_feed = {}
-    
-    # Map the main image input (index 0)
-    input_feed[inputs_info[0].name] = img_in
+    # --- MATH FIX: Map colors from [0, 255] to [-1.0, 1.0] ---
+    img_in = (img_rgb / 127.5) - 1.0
+    img_in = np.transpose(img_in, (2, 0, 1))[np.newaxis, :]
     
-    # Handle the auxiliary inputs (like 'weight')
+    # Map Inputs
+    inputs_info = face_session.get_inputs()
+    input_feed = {inputs_info[0].name: img_in}
     for i in range(1, len(inputs_info)):
-        input_name = inputs_info[i].name
-        input_type = inputs_info[i].type
-        
-        # Check if the model explicitly wants a double-precision float
-        dt = np.float64 if 'double' in input_type else np.float32
-        
-        input_feed[input_name] = np.array([fidelity], dtype=dt)
+        dt = np.float64 if 'double' in inputs_info[i].type else np.float32
+        input_feed[inputs_info[i].name] = np.array([fidelity], dtype=dt)
     
-    # 3. AI Inference
+    # AI Inference
     out = face_session.run(None, input_feed)[0]
     
-    # 4. Post-process
-    res = cv2.cvtColor((np.clip(np.squeeze(out), 0, 1).transpose(1, 2, 0) * 255.0).astype(np.uint8), cv2.COLOR_RGB2BGR)
+    # --- MATH FIX: Map colors back from [-1.0, 1.0] to [0, 1] ---
+    out = (np.squeeze(out) + 1.0) / 2.0
+    res = cv2.cvtColor((np.clip(out, 0, 1).transpose(1, 2, 0) * 255.0).astype(np.uint8), cv2.COLOR_RGB2BGR)
+    
     return cv2.resize(res, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_LANCZOS4)
 
+
 def hybrid_enhancer(img_data, mode, face_strength, progress=gr.Progress()):
     if img_data is None: return None
     img = img_data["composite"]
     if img.shape[2] == 4: img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
     
+    # Stage 1: Background
     if mode == "Full Image (BG + Face)":
         progress(0, desc="Stage 1: Upscaling Background...")
-        bg_upscaled = upscale_bg_tiled(img)
-        progress(0.5, desc="Stage 2: Restoring Face Details...")
-        # Use the strength slider as the internal model weight too
-        face_restored = restore_face_core(bg_upscaled, fidelity=face_strength)
-        return cv2.addWeighted(face_restored, face_strength, bg_upscaled, 1 - face_strength, 0)
+        final_img = upscale_bg_tiled(img)
     else:
-        # Just Face Restoration
-        return restore_face_core(img, fidelity=face_strength)
+        final_img = img.copy()
 
+    # Stage 2: Smart Face Detection
+    progress(0.5, desc="Stage 2: Scanning for Faces...")
+    
+    # Use OpenCV's built-in AI to find where the face actually is
+    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
+    gray = cv2.cvtColor(final_img, cv2.COLOR_BGR2GRAY)
+    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(64, 64))
+    
+    if len(faces) == 0:
+        print("No face detected! Running on whole image as fallback.")
+        restored = restore_face_core(final_img, fidelity=face_strength)
+        final_img = cv2.addWeighted(restored, face_strength, final_img, 1 - face_strength, 0)
+    else:
+        # Process each face found
+        for (x, y, w, h) in faces:
+            # Add some padding around the face so the AI sees the jawline/hair
+            pad_x = int(w * 0.2)
+            pad_y = int(h * 0.2)
+            x1 = max(0, x - pad_x)
+            y1 = max(0, y - int(pad_y * 1.5)) # Extra room at the top for hair
+            x2 = min(final_img.shape[1], x + w + pad_x)
+            y2 = min(final_img.shape[0], y + h + pad_y)
+            
+            # Extract just the face box
+            face_crop = final_img[y1:y2, x1:x2]
+            
+            # Restore the face perfectly
+            restored_face = restore_face_core(face_crop, fidelity=face_strength)
+            
+            # Paste it seamlessly back into the high-res background
+            final_img[y1:y2, x1:x2] = cv2.addWeighted(restored_face, face_strength, face_crop, 1 - face_strength, 0)
+            
+    progress(1.0, desc="Done!")
+    return final_img
 # --- UI ---
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo")) as demo:
     gr.Markdown("# 🏆 Hybrid Face & Background Enhancer")