MoinulwithAI commited on
Commit
42532cc
·
verified ·
1 Parent(s): 8653c6c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +58 -0
app.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ import torch.nn as nn
4
+ import torch.optim as optim
5
+ from torchvision import transforms, models
6
+ from PIL import Image
7
+ import gradio as gr
8
+ import matplotlib.pyplot as plt
9
+ import numpy as np
10
+
11
+ # Load the pre-trained model (ensure to use the saved model checkpoint)
12
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
13
+
14
+ # Model: EfficientNet-B0 with dropout added to reduce overfitting
15
+ model = models.efficientnet_b0(pretrained=True)
16
+ model.classifier = nn.Sequential(
17
+ nn.Dropout(0.4),
18
+ nn.Linear(model.classifier[1].in_features, 7) # num_classes = 7 (angry, disgust, fear, happy, neutral, sad, surprise)
19
+ )
20
+ model.load_state_dict(torch.load("D:/Dataset/MMAFEDB/best_mood_classifier.pth"))
21
+ model = model.to(device)
22
+ model.eval()
23
+
24
+ # Define the image transformations for the uploaded image
25
+ transform = transforms.Compose([
26
+ transforms.Resize((224, 224)),
27
+ transforms.ToTensor(),
28
+ transforms.Normalize([0.485, 0.456, 0.406],
29
+ [0.229, 0.224, 0.225])
30
+ ])
31
+
32
+ # Class names (same order as in your dataset)
33
+ class_names = ['angry', 'disgust', 'fear', 'happy', 'neutral', 'sad', 'surprise']
34
+
35
+ # Function to predict the mood from the uploaded image
36
+ def predict_mood(image):
37
+ image = Image.fromarray(image)
38
+ image = transform(image).unsqueeze(0).to(device)
39
+
40
+ with torch.no_grad():
41
+ outputs = model(image)
42
+ _, preds = torch.max(outputs, 1)
43
+ predicted_class = class_names[preds.item()]
44
+
45
+ return predicted_class
46
+
47
+ # Gradio interface
48
+ iface = gr.Interface(
49
+ fn=predict_mood,
50
+ inputs=gr.Image(type="numpy"),
51
+ outputs="text",
52
+ live=True,
53
+ title="Mood Classifier",
54
+ description="Upload an image of a face and the model will predict the mood."
55
+ )
56
+
57
+ # Launch the app
58
+ iface.launch()