Spaces:
Sleeping
Sleeping
File size: 6,183 Bytes
959739d |
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 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
import cv2
import torch
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import timm
from tqdm import tqdm
import torch.nn as nn
import os
import matplotlib.pyplot as plt
import torch.nn.functional as F
import dlib
import pandas as pd
# dnn_net = cv2.dnn.readNetFromCaffe("models/deploy.prototxt", "models/res10_300x300_ssd_iter_140000.caffemodel")
# # Initialize dlib's facial landmark predictor
# predictor = dlib.shape_predictor("models/shape_predictor_68_face_landmarks.dat")
def extract_face(image, net, predictor):
# Prepare the image for DNN face detection
(h, w) = image.shape[:2]
blob = cv2.dnn.blobFromImage(cv2.resize(image, (300, 300)), 1.0, (300, 300), (104.0, 177.0, 123.0))
net.setInput(blob)
detections = net.forward()
# Loop over the detections
for i in range(0, detections.shape[2]):
confidence = detections[0, 0, i, 2]
# Filter out weak detections
if confidence > 0.5:
box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
(startX, startY, endX, endY) = box.astype("int")
# Convert bounding box to dlib rectangle format
dlib_rect = dlib.rectangle(int(startX), int(startY), int(endX), int(endY))
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
landmarks = predictor(gray, dlib_rect)
landmarks_np = np.array([[p.x, p.y] for p in landmarks.parts()])
x, y, w, h = cv2.boundingRect(landmarks_np)
x -= 25
y -= 25
w += 50
h += 50
x = max(0, x)
y = max(0, y)
w = min(w, image.shape[1] - x)
h = min(h, image.shape[0] - y)
# Crop and resize the face
try:
face_crop = cv2.resize(face_crop, (224, 224))
except:
face_crop = cv2.resize(image, (224, 224))
return face_crop
return None
class Model:
def __init__(self,fps,fer_model):
self.device="cuda" if torch.cuda.is_available() else "cpu"
self.transform = transforms.Compose([transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
)
self.fermodel= timm.create_model("tf_efficientnet_b0_ns", pretrained=False)
self.fermodel.classifier = torch.nn.Identity()
self.fermodel.classifier=nn.Sequential(
nn.Linear(in_features=1280, out_features=7)
)
self.fermodel = torch.load(
fer_model,
map_location=self.device)
self.fermodel.to(self.device)
self.class_labels = ["angry", "disgust", "fear", "happy", "neutral", "sad", "surprised"]
self.emotion_reorder = {
0: 6,
1: 5,
2: 4,
3: 1,
4: 0,
5: 2,
6: 3,
}
self.label_dict = {
0: "angry",
1: "disgust",
2: "fear",
3: "happy",
4: "neutral",
5: "sad",
6: "surprised",
}
self.class_wise_frame_count=None
self.emotion_count = [0] * 7
self.frame_count=0
self.fps=fps
self.df=None
self.faces_=0
def predict(self,frames):
emotion_list=[]
emt=[]
for frame in tqdm(frames):
if frame is not None:
frame=np.copy(frame)
face_pil = Image.fromarray(
cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
)
face_tensor = self.transform(face_pil).unsqueeze(0).to(self.device)
with torch.no_grad():
output = self.fermodel(face_tensor)
_, predicted = torch.max(output, 1)
emotion = self.emotion_reorder[predicted.item()]
if isinstance(emotion, np.ndarray):
emotion = (
emotion.astype(float).item()
if emotion.size == 1
else emotion.tolist()
)
emotion = torch.tensor(
[emotion], dtype=torch.float32
) # Ensures it's a tensor
emotion.to(self.device)
emt.append(emotion)
self.emotion_count[predicted.item()] += 1
label = f"{self.label_dict[predicted.item()]}"
emotion_list.append(label)
else:
emt.append('frame error')
emotion_list.append('frame error')
return emotion_list,emt
def get_data(self,emotion_list,emt):
self.class_wise_frame_count = dict(zip(self.class_labels, self.emotion_count))
return emotion_list,self.class_wise_frame_count,emt
def fer_predict(video_frames,fps,model):
emotion_list,emt=model.predict(video_frames)
return model.get_data(emotion_list,emt)
def filter(list1,list2):
filtered_list1 = [x for i, x in enumerate(list1) if list2[i]!='fnf']
filtered_list2 = [x for x in list2 if x!='fnf']
return [filtered_list1,filtered_list2]
def plot_graph(x,y,var,path):
y = [value if isinstance(value, (int, float)) else np.nan for value in y]
print(len(y))
plt.plot(range(len(x)), y, linestyle='-')
plt.xlabel('Frame')
plt.ylabel(var)
plt.title(f'{var} Values vs Frame')
plt.savefig(path)
plt.clf()
# def save_frames(frames,folder_path):
# for i in tqdm(range(len(frames))):
# frame_filename = os.path.join(folder_path, f'frame_{i+1:04d}.jpg')
# # Save the frame as a .jpg file
# frame=cv2.cvtColor(frames[i],cv2.COLOR_BGR2RGB)
# cv2.imwrite(frame_filename, frame)
|