Create app.py

app.py

@@ -0,0 +1,103 @@
+import gradio as gr
+import cv2
+import numpy as np
+import clip
+import torch
+from PIL import Image
+import torch.nn.functional as F
+from facenet_pytorch import MTCNN
+
+# Global variables
+input_labels_X = "Happy Face, Sad Face, Angry Face, Fear Face, Disgust Face, Contempt Face, Nervous Face, Curious Face, Flirtatious Face, Ashamed Face, Bored Face, Confused Face, Calm Face, Proud Face, Guilty Face, Annoyed Face, Desperate Face, Jealous Face, Embarrassed Face, Impatient Face, Uncomfortable Face, Bitter Face, Helpless Face, Shy Face, Infatuated Face, Betrayed Face, Shocked Face, Relaxed Face, Apathetic Face, Neutral Face"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model, preprocess = clip.load("ViT-L/14", device=device)
+
+# Initialize MTCNN for face detection
+mtcnn = MTCNN(keep_all=True, device=device)
+
+def process_image(input_image, selected_model):
+    global model, preprocess
+    
+    # Load the selected model if it's different from the current one
+    if selected_model != model.model_name:
+        model, preprocess = clip.load(selected_model, device=device)
+    
+    # Convert input_image to numpy array
+    cv2_frame = np.array(input_image)
+    cv2_frame = cv2.cvtColor(cv2_frame, cv2.COLOR_RGB2BGR)
+
+    # Detect faces
+    frame_pil = Image.fromarray(cv2.cvtColor(cv2_frame, cv2.COLOR_BGR2RGB))
+    boxes, _ = mtcnn.detect(frame_pil)
+
+    # Find the largest face detected
+    largest_face = None
+    if boxes is not None and len(boxes) > 0:
+        largest_face = max(boxes, key=lambda box: (box[2] - box[0]) * (box[3] - box[1]))
+
+    # Process the largest face
+    if largest_face is not None:
+        x, y, w, h = map(int, largest_face)
+        cv2.rectangle(cv2_frame, (x, y), (w, h), (0, 255, 0), 2)
+        cropped_face = cv2_frame[y:h, x:w]
+
+        # Convert the cropped face to grayscale
+        frame_gray = cv2.cvtColor(cropped_face, cv2.COLOR_BGR2GRAY)
+        frame_gray_resized = cv2.resize(frame_gray, (224, 224))
+
+        # Convert the resized grayscale image to a tensor
+        frame_tensor = preprocess(Image.fromarray(frame_gray_resized)).unsqueeze(0).to(device)
+
+        # Tokenize input labels and prepare for model
+        input_labels = input_labels_X.split(", ")
+        text = clip.tokenize(input_labels).to(device)
+
+        with torch.no_grad():
+            # Encode the frame and text
+            image_features = model.encode_image(frame_tensor)
+            text_features = model.encode_text(text)
+
+            # Calculate logit
+            logit_per_image, logit_per_text = model(frame_tensor, text)
+
+            # Apply softmax to convert logits to probabilities
+            probabilities = F.softmax(logit_per_image[0], dim=0)
+
+        # Combine labels with probabilities and sort
+        combined_labels_probs = list(zip(input_labels, probabilities.tolist()))
+        combined_labels_probs.sort(key=lambda x: x[1], reverse=True)
+        top_five_labels_probs = combined_labels_probs[:5]
+
+        # Prepare results
+        results = []
+        for label, prob in top_five_labels_probs:
+            results.append(f"{label.strip()}: {prob * 100:.1f}%")
+
+        # Draw results on the image
+        for idx, result in enumerate(results):
+            cv2.putText(cv2_frame, result, (10, 30 + idx * 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
+
+        # Convert back to RGB for display
+        output_image = cv2.cvtColor(cv2_frame, cv2.COLOR_BGR2RGB)
+        return output_image, "\n".join(results)
+    else:
+        return cv2_frame, "No face detected"
+
+# Define the Gradio interface
+iface = gr.Interface(
+    fn=process_image,
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Dropdown(choices=["ViT-B/32", "ViT-B/16", "ViT-L/14"], label="Model", value="ViT-L/14")
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Processed Image"),
+        gr.Textbox(label="Emotion Probabilities")
+    ],
+    title="Emotion Recognition",
+    description="Upload an image to detect faces and recognize emotions."
+)
+
+# Launch the app
+iface.launch()