Spaces:
Running
Running
File size: 4,371 Bytes
54e742e 4eddd18 008760f 54e742e 988dde3 54e742e 008760f 988dde3 008760f 988dde3 008760f 988dde3 54e742e 988dde3 54e742e 008760f 988dde3 54e742e 988dde3 008760f 988dde3 008760f 988dde3 008760f 988dde3 008760f 988dde3 c3143ea 008760f 988dde3 008760f 988dde3 008760f 988dde3 008760f 988dde3 c3143ea 008760f 54e742e 988dde3 008760f 54e742e 988dde3 54e742e 988dde3 54e742e 988dde3 |
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 |
import cv2
import yaml
import numpy as np
import os
from typing import Tuple
from src.cv_utils import get_image
from transformers import ViTImageProcessor, ViTForImageClassification
import urllib3
with open("parameters.yml", "r") as stream:
try:
parameters = yaml.safe_load(stream)
except yaml.YAMLError as exc:
print(exc)
class GetFaceDemographics:
def __init__(self):
pass
@staticmethod
def preprocess_image_for_caffe_cnn(image: np.array):
model_mean = (
78.4263377603,
87.7689143744,
114.895847746,
) # taken from the model page on Caffe
blob = cv2.dnn.blobFromImage(image, 1.0, (227, 227), model_mean, swapRB=False)
return blob
@staticmethod
def get_age_cnn(blob) -> Tuple:
age_net = cv2.dnn.readNet(
parameters["face_age"]["config"], parameters["face_age"]["model"]
)
age_list = [
"(0-2)",
"(4-6)",
"(8-12)",
"(15-20)",
"(25-32)",
"(38-43)",
"(48-53)",
"(60-100)",
]
age_net.setInput(blob)
age_preds = age_net.forward()
i = age_preds[0].argmax()
age = age_list[i]
age_confidence_score = age_preds[0][i]
return age, age_confidence_score
@staticmethod
def get_gender_cnn(blob) -> Tuple:
gender_net = cv2.dnn.readNet(
parameters["face_gender"]["config"], parameters["face_gender"]["model"]
)
gender_list = ["Male", "Female"]
gender_net.setInput(blob)
gender_preds = gender_net.forward()
i = gender_preds[0].argmax()
gender = gender_list[i]
gender_confidence_score = gender_preds[0][i]
return gender, gender_confidence_score
@staticmethod
def get_age_vit(image: np.array) -> Tuple:
os.environ[
"CURL_CA_BUNDLE"
] = "" # fixes VPN issue when connecting to hugging face hub
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
id2label = {
0: "0-2",
1: "3-9",
2: "10-19",
3: "20-29",
4: "30-39",
5: "40-49",
6: "50-59",
7: "60-69",
8: "more than 70",
}
model = ViTForImageClassification.from_pretrained("nateraw/vit-age-classifier")
transforms = ViTImageProcessor.from_pretrained("nateraw/vit-age-classifier")
inputs = transforms(image, return_tensors="pt")
output = model(**inputs)
proba = output.logits.softmax(1)
preds = proba.argmax(1)
age_confidence_score = round(max(proba[0]).item(), 2)
age = id2label[int(preds)]
return age, age_confidence_score
@staticmethod
def get_gender_vit(image: np.array) -> Tuple:
os.environ[
"CURL_CA_BUNDLE"
] = "" # fixes VPN issue when connecting to hugging face hub
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
id2label = {
0: "female",
1: "male",
}
model = ViTForImageClassification.from_pretrained(
"rizvandwiki/gender-classification"
)
transforms = ViTImageProcessor.from_pretrained(
"rizvandwiki/gender-classification"
)
inputs = transforms(image, return_tensors="pt")
output = model(**inputs)
proba = output.logits.softmax(1)
preds = proba.argmax(1)
gender_confidence_score = round(max(proba[0]).item(), 2)
gender = id2label[int(preds)]
return gender, gender_confidence_score
def main(self, image_input) -> dict:
image = get_image(image_input)
age, age_confidence_score = self.get_age_vit(image)
gender, gender_confidence_score = self.get_gender_vit(image)
d = {
"age_range": age,
"age_confidence": age_confidence_score,
"gender": gender,
"gender_confidence": gender_confidence_score,
}
return d
if __name__ == "__main__":
path_to_images = "data/"
image_files = os.listdir(path_to_images)
for image in image_files:
print(image)
results = GetFaceDemographics().main(path_to_images + image)
print(results)
|