File size: 2,070 Bytes
9de35b0 18a9dce 8cc5e9d 18a9dce 8cc5e9d 18a9dce 8eb51e6 18a9dce 9de35b0 18a9dce |
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 |
import torch
import os
from torchvision import transforms
import numpy as np
# Checking for all types of devices available
if torch.backends.mps.is_available():
device = "mps"
elif torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
print(f"Using device: {device}")
image_arr = []
for foldername, subfolders, filenames in os.walk("../FER/Images/"):
for filename in filenames:
# Construct the full path to the file
file_path = os.path.join(foldername, filename)
image_arr.append(f"{file_path}")
def predict(model, image_path):
from face_detection import face_detection
with torch.no_grad():
transform = transforms.Compose(
[
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
),
transforms.RandomErasing(p=1, scale=(0.05, 0.05)),
]
)
face = face_detection(image_path)
image_tensor = transform(face).unsqueeze(0)
image_tensor = image_tensor.to(device)
model.eval()
img_pred = model(image_tensor)
topk = (3,)
with torch.no_grad():
maxk = max(topk)
# batch_size = target.size(0)
_, pred = img_pred.topk(maxk, 1, True, True)
pred = pred.t()
img_pred = pred
img_pred = img_pred.squeeze().cpu().numpy()
im_pre_label = np.array(img_pred)
y_pred = im_pre_label.flatten()
emotions = {
0: "Surprise",
1: "Fear",
2: "Disgust",
3: "Happy",
4: "Sad",
5: "Angry",
6: "Neutral",
}
labels = []
for i in y_pred:
labels.append(emotions.get(i))
print(
f"-->Image Path {image_path} [!] The predicted labels are {y_pred} and the label is {labels}"
)
return
|