gregmerritt commited on
Commit
05c0ede
·
verified ·
1 Parent(s): ab2e879

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +108 -0
app.py ADDED
@@ -0,0 +1,108 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import torch
3
+ import torch.nn as nn
4
+ from torchvision import transforms
5
+ from PIL import Image
6
+ import gradio as gr
7
+ from transformers import pipeline
8
+
9
+ # Load emotion classes
10
+ classes = ['Angry', 'Disgust', 'Fear', 'Happy', 'Sad', 'Surprise', 'Neutral']
11
+
12
+ # Set device
13
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14
+
15
+ # Define the actual architecture used for training
16
+ class EmotionModel(nn.Module):
17
+ def __init__(self):
18
+ super(EmotionModel, self).__init__()
19
+ self.model = nn.Sequential(
20
+ nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1),
21
+ nn.ReLU(),
22
+ nn.MaxPool2d(2, 2),
23
+ nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
24
+ nn.ReLU(),
25
+ nn.MaxPool2d(2, 2),
26
+ nn.Flatten(),
27
+ nn.Linear(32 * 12 * 12, 128),
28
+ nn.ReLU(),
29
+ nn.Linear(128, 7)
30
+ )
31
+
32
+ def forward(self, x):
33
+ return self.model(x)
34
+
35
+ # Load the model
36
+ model = EmotionModel().to(device)
37
+ model.load_state_dict(torch.load("emotion_model.pth", map_location=device))
38
+ model.eval()
39
+
40
+ # Transformation
41
+ transform = transforms.Compose([
42
+ transforms.Grayscale(),
43
+ transforms.Resize((48, 48)),
44
+ transforms.ToTensor(),
45
+ transforms.Normalize((0.5,), (0.5,))
46
+ ])
47
+
48
+ # NLP pipeline
49
+ gen = pipeline("text-generation", model="distilgpt2")
50
+
51
+ # Icebreaker templates
52
+ templates = {
53
+ "Friendly": {
54
+ "Happy": "Looks like you're in a great mood! 🌞 Here's a cheerful icebreaker: ",
55
+ "Sad": "Hey, if you're feeling down, maybe this will lift your spirits: ",
56
+ "Angry": "Let’s turn that frown around! Here's something light-hearted: ",
57
+ "Fear": "Everything okay? Here's a calming question to ease the mood: ",
58
+ "Disgust": "Here's a quirky icebreaker to distract from the ick: ",
59
+ "Surprise": "Caught off guard? Here's a fun fact to break the ice: ",
60
+ "Neutral": "Let's get the conversation started with this: "
61
+ },
62
+ "Professional": {
63
+ "Happy": "You seem upbeat! Here's a smart opener for professional settings: ",
64
+ "Sad": "Here's a thoughtful question to get engagement going: ",
65
+ "Angry": "Let's channel focus with this topic: ",
66
+ "Fear": "Start with a grounded tone: ",
67
+ "Disgust": "Try this conversation starter to refocus: ",
68
+ "Surprise": "Introduce with curiosity using this: ",
69
+ "Neutral": "Here's a balanced opener: "
70
+ },
71
+ "Funny": {
72
+ "Happy": "This joke might make your smile bigger: ",
73
+ "Sad": "Here’s something silly to flip the mood: ",
74
+ "Angry": "Try this pun to vent the steam: ",
75
+ "Fear": "Scared? Laugh it off with this: ",
76
+ "Disgust": "Try this gross-but-funny line: ",
77
+ "Surprise": "Unexpected? This might top that: ",
78
+ "Neutral": "Break the silence with this joke: "
79
+ }
80
+ }
81
+
82
+ # Prediction logic
83
+ def predict_emotion_and_icebreaker(image, tone):
84
+ image = Image.fromarray(image).convert("RGB")
85
+ image = transform(image).unsqueeze(0).to(device)
86
+ with torch.no_grad():
87
+ output = model(image)
88
+ pred = output.argmax(dim=1).item()
89
+ emotion = classes[pred]
90
+
91
+ prompt = templates[tone][emotion] + " Here's an idea to open the chat:"
92
+ response = gen(prompt, max_length=40, num_return_sequences=1)[0]['generated_text']
93
+
94
+ return f"🧠 Emotion Detected: {emotion}\n💬 Icebreaker ({tone}):\n{response}"
95
+
96
+ # Interface
97
+ webcam_input = gr.Image(type="numpy", label="Upload or Take a Photo")
98
+ tone_dropdown = gr.Dropdown(choices=["Friendly", "Professional", "Funny"], value="Friendly", label="Tone")
99
+
100
+ demo = gr.Interface(
101
+ fn=predict_emotion_and_icebreaker,
102
+ inputs=[webcam_input, tone_dropdown],
103
+ outputs="text",
104
+ title="Emotion + Icebreaker Generator",
105
+ description="Upload or capture a face photo. AI will predict the emotion and generate a tone-specific conversation starter."
106
+ )
107
+
108
+ demo.launch()