fixing recognition integration

app.py

@@ -1,14 +1,14 @@
 import torch
-from facenet_pytorch import MTCNN
-import pickle
+from torch import nn
+import torchvision.transforms as transforms
 import cv2
+import numpy as np
 import gradio as gr
 from PIL import Image
-import numpy as np
+from facenet_pytorch import MTCNN
 from transformers import ViTImageProcessor, ViTModel
-import torch.nn as nn
-from torchvision import transforms
-
+import pickle
+import time
 
 # Define the ViT class
 class ViT(nn.Module):
@@ -87,11 +87,13 @@ def draw_annotations(frame, detections, names=None):
     return frame
 
 def process_image(image):
+    start_time = time.time()
+    
     frame = np.array(image)
     aligned_faces, boxes = align_faces(frame, mtcnn, device)
     
+    names = []
     if aligned_faces is not None:
-        names = []
         for face in aligned_faces:
             face = transform(face)
             face = face.unsqueeze(0).to(device)
@@ -108,27 +110,20 @@ def process_image(image):
         annotated_image = frame
         result = "No faces detected."
     
+    end_time = time.time()
+    inference_time = end_time - start_time
+    result += f" Inference time: {inference_time:.2f} seconds"
+    
     return annotated_image, result
 
-def capture_and_process_image(webcam_image):
-    captured_img, result = process_image(webcam_image)
-    return captured_img, result
-
 # Create the Gradio interface
-with gr.Blocks() as demo:
-    with gr.Row():
-        # Webcam input component
-        webcam_input = gr.Image(source="webcam", streaming=True, label="Webcam Input", height=483)
-        # Captured image display
-        captured_image = gr.Image(label="Captured Image", height=483)
-    # Capture button
-    capture_button = gr.Button("Capture Image")
-    # Result output textbox
-    result_output = gr.Textbox(label="Inference Result")
-
-    # Define the button click action
-    capture_button.click(fn=capture_and_process_image, inputs=webcam_input, outputs=[captured_image, result_output])
-
-if __name__ == "__main__":
-    # Launch the interface with share=True to create a public link
-    demo.launch(share=True, debug=True)
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="numpy"),
+    title="Face Detection and Recognition with MTCNN and ViT",
+     description="Upload an image and the model will detect and align faces in it."
+)
+
+# Launch the interface
+iface.launch(share=True, debug=True)