fxing gradio

app.py

@@ -32,13 +32,17 @@
 #     model.to(device)
 #     return model, processor, device
 
-import gradio as gr
-import cv2
 import torch
-from facenet_pytorch import MTCNN
+from mtcnn import MTCNN
+import cv2
+import numpy as np
+import gradio as gr
+from PIL import Image
 
-# Load MTCNN for face detection
+# Set the device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Initialize MTCNN
 mtcnn = MTCNN(keep_all=True, min_face_size=20, thresholds=[0.6, 0.7, 0.7], device=device)
 
 def align_faces(frame, mtcnn, device):
@@ -62,46 +66,20 @@ def draw_annotations(frame, detections, names=None):
             cv2.putText(frame, names[i], (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)
     return frame
 
-def capture_frames():
-    cap = cv2.VideoCapture(0)
-    
-    if not cap.isOpened():
-        raise RuntimeError("Error: Could not open video stream.")
-
-    while True:
-        ret, frame = cap.read()
-        
-        if not ret:
-            raise RuntimeError("Error: Failed to capture image")
-        
-        # Align faces using MTCNN
-        aligned_faces, boxes = align_faces(frame, mtcnn, device)
-        
-        # Draw annotations on the frame
-        annotated_frame = draw_annotations(frame, boxes)
-
-        _, buffer = cv2.imencode('.jpg', annotated_frame)
-        frame_bytes = buffer.tobytes()
-        
-        yield frame_bytes
-
-def video_frame_generator():
-    for frame in capture_frames():
-        yield frame
-
-def gradio_interface():
-    with gr.Blocks() as demo:
-        with gr.Row():
-            webcam_output = gr.Video(source=video_frame_generator, streaming=True, label="Webcam Output")
-        stop_button = gr.Button("Stop")
-
-        def stop_streaming():
-            # Placeholder for stopping streaming if necessary
-            return "Streaming stopped."
-
-        stop_button.click(fn=stop_streaming, inputs=None, outputs=None)
+def process_image(image):
+    frame = np.array(image)
+    aligned_faces, boxes = align_faces(frame, mtcnn, device)
+    annotated_image = draw_annotations(frame, boxes)
+    return annotated_image
 
-    demo.launch(share=True, debug=True)
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs="image",
+    title="Face Detection with MTCNN",
+    description="Upload an image and the model will detect and align faces in it."
+)
 
-if __name__ == "__main__":
-    gradio_interface()
+# Launch the interface
+iface.launch()