quick imp
So im trying to implement a pre-trained model for arousal and valence, this is the model
https://github.com/face-analysis/emonet
model = EmoNet(n_expression=5)
state_dict = torch.load('emonet\pretrained\emonet_5.pth', map_location='cpu')
model.load_state_dict(state_dict)
model.eval()
transform = transforms.Compose(#Simply transforming to tensor)
def preprocess_frame(frame):
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image = Image.fromarray(frame)
image = transform(image).unsqueeze(0)
return image
def analyze_emotion(frame):
preprocessed_frame = preprocess_frame(frame)
with torch.no_grad():
output = model(preprocessed_frame)
valence = output['valence'].item()
arousal = output['arousal'].item()
emotion_scores = output['expression'].squeeze(0)
print(f'Raw valence: {output["valence"].item()}, Raw arousal: {output["arousal"].item()}')
print(f'Emotion scores: {emotion_scores}')
emotion = emotion_scores.argmax().item()
print(f'Predicted Emotion: {emotion_labels[emotion]} (Index: {emotion})')
return valence, arousal, emotion
emotion_labels = ["Neutral", "Happy", "Sad", "Surprise", "Fear"]
# Emotion Plot
def setup_emotion_plot(ax):
#Character waste
video_source = 0
cap = cv2.VideoCapture(video_source)
if not cap.isOpened():
print("Error: Could not open video source.")
exit()
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
point = setup_emotion_plot(ax2)
plt.ion()
while True:
ret, frame = cap.read()
if not ret:
break
valence, arousal, emotion = analyze_emotion(frame)
update_emotion_plot(point, valence, arousal)
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
ax1.clear()
ax1.imshow(frame_rgb)
ax1.axis('off')
fig.canvas.draw()
plt.pause(0.001)
if plt.waitforbuttonpress(timeout=0.001):
break
cap.release()
plt.close()
red plover
