Spaces:
Sleeping
Sleeping
File size: 4,739 Bytes
3d3f535 |
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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
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
# Initialize the Hugging Face pipeline for facial emotion detection using the "trpakov/vit-face-expression" model
emotion_pipeline = pipeline("image-classification", model="trpakov/vit-face-expression")
# Function to analyze sentiment
def analyze_sentiment(face):
# Convert face to RGB
rgb_face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
# Convert the face to a PIL image
pil_image = Image.fromarray(rgb_face)
# Analyze sentiment using the Hugging Face pipeline
results = emotion_pipeline(pil_image)
# Get the dominant emotion
dominant_emotion = max(results, key=lambda x: x['score'])['label']
return dominant_emotion
TEXT_SIZE = 3
# Function to detect faces, analyze sentiment, and draw a red box around them
def detect_and_draw_faces(frame):
# Convert frame to RGB
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Convert the frame to a PIL image
pil_image = Image.fromarray(rgb_frame)
# Analyze sentiment using the Hugging Face pipeline
results = emotion_pipeline(pil_image)
# Print the results to understand the structure
print(results)
# Draw on the PIL image
draw = ImageDraw.Draw(pil_image)
# Iterate through detected faces
for result in results:
box = result['box']
sentiment = result['label']
# Draw rectangle and text
x, y, w, h = box['left'], box['top'], box['width'], box['height']
draw.rectangle(((x, y), (x+w, y+h)), outline="red", width=3)
# Calculate position for the text background and the text itself
text_size = draw.textsize(sentiment)
background_tl = (x, y - text_size[1] - 5)
background_br = (x + text_size[0], y)
# Draw black rectangle as background
draw.rectangle([background_tl, background_br], fill="black")
# Draw white text on top
draw.text((x, y - text_size[1]), sentiment, fill="white")
# Convert back to OpenCV format
frame_with_boxes = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
return frame_with_boxes
# Function to capture video from webcam
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()
# Streamlit UI
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")
# Columns for input and output streams
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()
# Start video stream
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
# Detect faces, analyze sentiment, and draw red boxes with sentiment labels
frame_with_boxes = detect_and_draw_faces(frame)
# Display the input stream with the red box around the face
video_placeholder.image(frame_with_boxes, channels="BGR")
# Display the output stream (here it's the same as input, modify as needed)
output_placeholder.image(frame_with_boxes, channels="BGR")
# Add a short delay to control the frame rate
if cv2.waitKey(1) & 0xFF == ord('q'):
break
|