Spaces:
Sleeping
Sleeping
File size: 3,431 Bytes
6c9c423 |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import cv2
import numpy as np
import tempfile
import time
import streamlit as st
from PIL import Image
from io import BytesIO
import plotly.graph_objs as go
from transformers import CLIPProcessor, CLIPModel
from torch.cuda import is_available
MODEL_ID = "openai/clip-vit-base-patch32"
DEMO_IMAGE = 'demo.jpg'
EMOTION_DICT = {
0: ['Angry', '😡'],
1: ['Disgusted', '🤢'],
2: ['Fearful', '😨'],
3: ['Happy', '😃'],
4: ['Neutral', '😐'],
5: ['Sad', '☹️'],
6: ['Surprised', '😮']
}
device = 'cuda' if is_available() else 'cpu'
@st.cache_data
def load_model():
processor = CLIPProcessor.from_pretrained(MODEL_ID)
model = CLIPModel.from_pretrained(MODEL_ID)
return processor, model
@st.cache_data
def load_token_embds():
emotions = list(EMOTION_DICT.values())
desc = [f'a photo of a {emotion[0]} person' for emotion in emotions]
tok = processor(text = desc, return_tensors = 'pt', images = None, padding = True).to(device)
tok_emb = model.get_text_features(**tok)
tok_emb = tok_emb.detach().cpu().numpy() / np.linalg.norm(tok_emb.detach().cpu().numpy(), axis=0)
return tok_emb
st.set_page_config(page_title="Mood Scope", page_icon="🎭")
st.title('Mood-Scope')
st.sidebar.title('Options')
app_mode = st.sidebar.selectbox('Choose Page', ['About the App', 'Run Mood Scope'])
st.markdown(
"""
<style>
[data-testid = 'stSidebar'][aria-expanded = 'true'] > div:first-child{
width: 350px
}
[data-testid = 'stSidebar'][aria-expanded = 'false'] > div:first-child{
width: 350px
margin-left: -350px
}
</style>
""", unsafe_allow_html = True
)
if app_mode == 'About the App':
st.markdown('Will edit this later!!')
elif app_mode == 'Run Mood Scope':
processor, model = load_model()
st.sidebar.markdown('---')
with st.columns(3)[1]:
kpi = st.markdown('**Dominant Detected Emotion**')
emotion_emoji = st.markdown('-')
#emotion_text = st.markdown('-')
img_file_buffer = st.sidebar.file_uploader('Upload an Image', type = ['jpg', 'png', 'jpeg'])
if img_file_buffer:
buffer = BytesIO(img_file_buffer.read())
data = np.frombuffer(buffer.getvalue(), dtype=np.uint8)
image = cv2.imdecode(data, cv2.IMREAD_COLOR)
else:
demo_image = DEMO_IMAGE
image = cv2.imread(demo_image, cv2.IMREAD_COLOR)
st.sidebar.text('Original Image')
st.sidebar.image(image, channels = 'BGR')
im_proc = processor(images=image, return_tensors='pt')['pixel_values']
im_emb = model.to(device).get_image_features(im_proc.to(device))
im_emb = im_emb.detach().cpu().numpy()
tok_emb = load_token_embds()
score = np.dot(im_emb, tok_emb.T)
output_emoji = EMOTION_DICT[score.argmax(axis = 1).item()][1]
output_text = EMOTION_DICT[score.argmax(axis = 1).item()][0]
emotion_emoji.write(f'<h1> {output_emoji} </h1>', unsafe_allow_html = True)
categories = [emotion[0] for emotion in EMOTION_DICT.values()]
data = list(map(int, (100 * (score / score.sum())).squeeze()))
trace = go.Scatterpolar(r = data, theta = categories, fill = 'toself', name = 'Emotions')
layout = go.Layout(
polar = dict(
radialaxis = dict(
visible = False,
range = [0, 50]
)
),
)
fig = go.Figure(data=[trace], layout=layout)
st.plotly_chart(fig, use_container_width=True)
#emotion_text.write(f'**{output_text}**')
|