Spaces:
Build error
Build error
File size: 2,383 Bytes
fd6967b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import streamlit as st
from PIL import Image
import torch
from torch import nn as nn
from torchvision.transforms import transforms
class CNNModel(nn.Module):
def __init__(self):
super(CNNModel, self).__init__()
self.conv_layers = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=0),
nn.ReLU(),
nn.BatchNorm2d(32),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(32, 64, kernel_size=3, padding=0),
nn.ReLU(),
nn.BatchNorm2d(64),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(64, 128, kernel_size=3, padding=0),
nn.ReLU(),
nn.BatchNorm2d(128),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.fc_layers = nn.Sequential(
nn.Flatten(),
nn.Linear(128 * 30 * 30, 128),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(128, 64),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(64, 1),
nn.Sigmoid()
)
def forward(self, x):
x = self.conv_layers(x)
x = self.fc_layers(x)
return x
def load_checkpoint(checkpoint, model):
print("=> Loading checkpoint")
model.load_state_dict(checkpoint["state_dict"])
model = CNNModel()
load_checkpoint(torch.load("emotion.pth.tar"), model)
model.eval()
class_names = ["angry", "disgust", "fear", "happy", "neutral", "sad", "surprize"]
st.title("Emotion Detector")
st.write("Upload an image and let the model predict your emotion!")
uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
if uploaded_image is not None:
image = Image.open(uploaded_image)
st.image(image, caption="Uploaded Image", use_column_width=True)
# Define the transformation to convert the image to a tensor
transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# Apply the transformation to the image
tensor_image = transform(image)
# Predict
predictions = model(tensor_image.unsqueeze(0))
predicted_class_index = torch.argmax(predictions).item()
predicted_class = class_names[predicted_class_index]
st.write(f"Prediction: {predicted_class}")
|