Update app.py

app.py

@@ -1,56 +1,59 @@
 import cv2
 import numpy as np
-import gradio as gr
 import onnxruntime as ort
-from facenet_pytorch import MTCNN
+import gradio as gr
 from PIL import Image
 
-# Load GFPGAN ONNX
-gfp_session = ort.InferenceSession("models/gfpgan_1.4.onnx", providers=["CPUExecutionProvider"])
+# Load GFPGAN ONNX model
+onnx_path = "models/gfpgan_1.4.onnx"
+session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
 
-# Load MTCNN for face detection and alignment
-face_detector = MTCNN(keep_all=True, device='cpu')
+def preprocess_face(img, box, size=512):
+    x1, y1, x2, y2 = [int(v) for v in box]
+    face = img[y1:y2, x1:x2]
+    face = cv2.resize(face, (size, size), interpolation=cv2.INTER_LINEAR)
+    face = face.astype(np.float32) / 255.0
+    face = face.transpose(2, 0, 1)[np.newaxis, :]
+    return face, (x1, y1, x2, y2)
 
-def restore_face_from_crop(crop):
-    """Run GFPGAN on a single aligned crop"""
-    crop = crop.resize((512, 512), Image.BICUBIC)
-    np_crop = np.array(crop).astype(np.float32) / 255.0
-    np_crop = np_crop.transpose(2, 0, 1)[None, ...]
-    inputs = {gfp_session.get_inputs()[0].name: np_crop}
-    output = gfp_session.run(None, inputs)[0]
+def postprocess_face(output, face_box, original_shape):
     restored = output.squeeze().transpose(1, 2, 0)
     restored = np.clip(restored * 255, 0, 255).astype(np.uint8)
-    return Image.fromarray(restored)
-
-def enhance_faces(input_image):
-    input_pil = input_image.convert("RGB")
-    boxes, probs = face_detector.detect(input_pil)
-
-    if boxes is None or len(boxes) == 0:
-        return input_image, "❌ No faces detected."
-
-    restored_faces = []
-    for box in boxes:
-        x1, y1, x2, y2 = map(int, box)
-        face_crop = input_pil.crop((x1, y1, x2, y2))
-        restored = restore_face_from_crop(face_crop)
-        restored_faces.append(restored)
-
-    # Show only restored faces separately
-    if restored_faces:
-        return input_pil, restored_faces
-    else:
-        return input_pil, "❌ Could not restore faces."
-
+    restored = cv2.resize(restored, (face_box[2] - face_box[0], face_box[3] - face_box[1]))
+    return restored
+
+def detect_faces(img):
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
+    boxes = face_cascade.detectMultiScale(gray, 1.3, 5)
+    return [[x, y, x + w, y + h] for (x, y, w, h) in boxes]
+
+def enhance_faces(image_pil):
+    image = np.array(image_pil.convert("RGB"))
+    img = image.copy()
+    faces = detect_faces(img)
+    
+    if not faces:
+        return image_pil, "No faces detected."
+
+    for box in faces:
+        face_input, face_coords = preprocess_face(img, box)
+        output = session.run(None, {session.get_inputs()[0].name: face_input})[0]
+        restored = postprocess_face(output, box, img.shape)
+
+        # Paste restored face back
+        x1, y1, x2, y2 = box
+        img[y1:y2, x1:x2] = restored
+
+    return Image.fromarray(img), f"{len(faces)} face(s) enhanced."
+
+# Gradio UI
 iface = gr.Interface(
     fn=enhance_faces,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=[
-        gr.Image(type="pil", label="Original"),
-        gr.Gallery(label="Restored Faces").style(grid=3, height=200),
-    ],
-    title="🧠 GFPGAN 1.4 Face Restoration (Improved)",
-    description="Uses MTCNN face detection + GFPGAN ONNX for high-quality restoration.",
+    inputs=gr.Image(type="pil", label="Input Image"),
+    outputs=[gr.Image(type="pil", label="Restored Image"), gr.Textbox(label="Log")],
+    title="GFPGAN v1.4 (ONNX) - Face Restoration",
+    description="Run GFPGAN face restoration using the ONNX model (no PyTorch dependencies)."
 )
 
 if __name__ == "__main__":