|
import os |
|
os.environ['OPENCV_AVFOUNDATION_SKIP_AUTH'] = '1' |
|
|
|
import streamlit as st |
|
import cv2 |
|
import numpy as np |
|
from transformers import pipeline |
|
from PIL import Image, ImageDraw |
|
|
|
|
|
emotion_pipeline = pipeline("image-classification", model="trpakov/vit-face-expression") |
|
|
|
|
|
def analyze_sentiment(face): |
|
|
|
rgb_face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB) |
|
|
|
pil_image = Image.fromarray(rgb_face) |
|
|
|
results = emotion_pipeline(pil_image) |
|
|
|
dominant_emotion = max(results, key=lambda x: x['score'])['label'] |
|
return dominant_emotion |
|
|
|
TEXT_SIZE = 3 |
|
|
|
|
|
def detect_and_draw_faces(frame): |
|
|
|
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
|
|
|
pil_image = Image.fromarray(rgb_frame) |
|
|
|
results = emotion_pipeline(pil_image) |
|
|
|
|
|
print(results) |
|
|
|
|
|
draw = ImageDraw.Draw(pil_image) |
|
|
|
|
|
for result in results: |
|
box = result['box'] |
|
sentiment = result['label'] |
|
|
|
|
|
x, y, w, h = box['left'], box['top'], box['width'], box['height'] |
|
draw.rectangle(((x, y), (x+w, y+h)), outline="red", width=3) |
|
|
|
|
|
text_size = draw.textsize(sentiment) |
|
background_tl = (x, y - text_size[1] - 5) |
|
background_br = (x + text_size[0], y) |
|
|
|
|
|
draw.rectangle([background_tl, background_br], fill="black") |
|
|
|
draw.text((x, y - text_size[1]), sentiment, fill="white") |
|
|
|
|
|
frame_with_boxes = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR) |
|
|
|
return frame_with_boxes |
|
|
|
|
|
def video_stream(): |
|
video_capture = cv2.VideoCapture(0) |
|
if not video_capture.isOpened(): |
|
st.error("Error: Could not open video capture device.") |
|
return |
|
|
|
while True: |
|
ret, frame = video_capture.read() |
|
if not ret: |
|
st.error("Error: Failed to read frame from video capture device.") |
|
break |
|
yield frame |
|
|
|
video_capture.release() |
|
|
|
|
|
st.markdown( |
|
""" |
|
<style> |
|
.main { |
|
background-color: #FFFFFF; |
|
} |
|
.reportview-container .main .block-container{ |
|
padding-top: 2rem; |
|
} |
|
h1 { |
|
color: #E60012; |
|
font-family: 'Arial Black', Gadget, sans-serif; |
|
} |
|
h2 { |
|
color: #E60012; |
|
font-family: 'Arial', sans-serif; |
|
} |
|
h3 { |
|
color: #333333; |
|
font-family: 'Arial', sans-serif; |
|
} |
|
.stButton button { |
|
background-color: #E60012; |
|
color: white; |
|
border-radius: 5px; |
|
font-size: 16px; |
|
} |
|
</style> |
|
""", |
|
unsafe_allow_html=True |
|
) |
|
|
|
st.title("Computer Vision Test Lab") |
|
st.subheader("Facial Sentiment") |
|
|
|
|
|
col1, col2 = st.columns(2) |
|
|
|
with col1: |
|
st.header("Input Stream") |
|
st.subheader("Webcam") |
|
video_placeholder = st.empty() |
|
|
|
with col2: |
|
st.header("Output Stream") |
|
st.subheader("Analysis") |
|
output_placeholder = st.empty() |
|
|
|
sentiment_placeholder = st.empty() |
|
|
|
|
|
video_capture = cv2.VideoCapture(0) |
|
if not video_capture.isOpened(): |
|
st.error("Error: Could not open video capture device.") |
|
else: |
|
while True: |
|
ret, frame = video_capture.read() |
|
if not ret: |
|
st.error("Error: Failed to read frame from video capture device.") |
|
break |
|
|
|
|
|
frame_with_boxes = detect_and_draw_faces(frame) |
|
|
|
|
|
video_placeholder.image(frame_with_boxes, channels="BGR") |
|
|
|
|
|
output_placeholder.image(frame_with_boxes, channels="BGR") |
|
|
|
|
|
if cv2.waitKey(1) & 0xFF == ord('q'): |
|
break |
|
|