Spaces:
Running
Running
File size: 3,802 Bytes
04f9904 |
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 |
import os
import tempfile
import torch
import cv2
from PIL import Image
from torchvision import transforms
from facenet_pytorch import MTCNN
def process_emotion(input_video):
try:
# Load the pre-trained CNN model
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
num_ftrs = model.fc.in_features
model.fc = torch.nn.Linear(num_ftrs, 7) # 7 classes for emotions
# Load the pre-trained weights
model.eval()
# Define emotions
emotions = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
# Initialize MTCNN for face detection
mtcnn = MTCNN()
# Open the video capture
cap = cv2.VideoCapture(input_video.name)
if not cap.isOpened():
return "Error: Could not open video capture."
# Create a temporary directory to store the frames
temp_dir = tempfile.mkdtemp()
frame_paths = []
# Define image transformations
preprocess = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# Process each frame of the video
frame_count = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
if frame is None:
print("Error: Empty frame.")
continue
frame_count += 1
pil_frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
face_boxes, _ = mtcnn.detect(pil_frame)
if face_boxes is None:
print(f"No face detected in frame {frame_count}.")
else:
for box in face_boxes:
x1, y1, x2, y2 = box.astype(int)
w = x2 - x1
h = y2 - y1
face_roi = frame[y1:y2, x1:x2]
pil_face = Image.fromarray(cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB))
input_tensor = preprocess(pil_face)
input_batch = input_tensor.unsqueeze(0)
with torch.no_grad():
output = model(input_batch)
predicted_emotion = output.argmax().item()
emotion_label = emotions[predicted_emotion]
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(frame, emotion_label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
frame_path = os.path.join(temp_dir, f"{frame_count:04d}.png")
cv2.imwrite(frame_path, frame)
frame_paths.append(frame_path)
cap.release()
cv2.destroyAllWindows()
if not frame_paths:
return "No faces detected in the video."
# Convert the frames to a video
output_path = "output_video.mp4"
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
# Get dimensions from the first frame
sample_frame = cv2.imread(frame_paths[0])
if sample_frame is None:
return "Error: Unable to read sample frame."
frame_height, frame_width, _ = sample_frame.shape
out = cv2.VideoWriter(output_path, fourcc, 30, (frame_width, frame_height))
for frame_path in frame_paths:
frame = cv2.imread(frame_path)
if frame is not None:
out.write(frame)
os.remove(frame_path)
else:
print(f"Warning: Unable to read frame {frame_path}")
out.release()
return output_path
except Exception as e:
return f"An error occurred: {str(e)}"
|