Spaces:
No application file
No application file
File size: 3,538 Bytes
0a9d067 |
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 |
from ultralytics import YOLO
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array
import cv2
import numpy as np
import sys
import os
import base64
# Función para preprocesar imágenes desde un arreglo
def preprocess_image_from_array(image, target_size=(256, 256)):
img = cv2.resize(image, target_size) # Redimensionar la imagen
img_array = img_to_array(img) # Convertir a arreglo
img_array = np.expand_dims(img_array, axis=0) # Expandir dimensiones
img_array = img_array / 255.0 # Normalizar
return img_array
# Función para detectar y recortar el rostro
def detect_and_crop_face(image, yolo_model, padding=0.1):
results = yolo_model(image)
boxes = results[0].boxes.xyxy.cpu().numpy() # Obtener las coordenadas de las cajas
confidences = results[0].boxes.conf.cpu().numpy().tolist() # Obtener las confianzas
if len(boxes) > 0:
max_index = confidences.index(max(confidences)) # Índice de mayor confianza
x1, y1, x2, y2 = map(int, boxes[max_index]) # Coordenadas de la caja
# Calcular el tamaño del padding
h, w, _ = image.shape
box_width = x2 - x1
box_height = y2 - y1
pad_x = int(box_width * padding) # Padding horizontal
pad_y = int(box_height * padding) # Padding vertical
# Ajustar las coordenadas con padding, asegurándose de no salir de la imagen
x1 = max(0, x1 - pad_x)
y1 = max(0, y1 - pad_y)
x2 = min(w, x2 + pad_x)
y2 = min(h, y2 + pad_y)
# Recortar la cara con padding
cropped_face = image[y1:y2, x1:x2]
return cropped_face, max(confidences)
return None, 0 # No se detectó ninguna cara
# Función principal
def predict_face(img_base64):
# Carga del modelo YOLO para detección de caras
yolo_model = YOLO(os.path.join(os.path.dirname(__file__), "yolov8m-face.pt"))
# Carga del modelo de clasificación
emotion_model = load_model(os.path.join(os.path.dirname(__file__), './model13.h5'))
# class_names = ["sadness", "happiness", "love", "anger", "worry", "neutral"]
class_names = ["anger", "fear", "happy", "neutral", "sad", "surprise"]
image_data = base64.b64decode(img_base64)
np_array = np.frombuffer(image_data, np.uint8)
# Paso 3: Leer la imagen con OpenCV desde el array NumPy
image = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
if image is None:
print("Error: no se pudo cargar la imagen.")
return
# Detectar y recortar el rostro
face, confidence = detect_and_crop_face(image, yolo_model)
if face is not None:
print(f"Rostro detectado con confianza: {confidence*100:.2f}%")
preprocessed_face = preprocess_image_from_array(face)
# Predicción de emoción
prediction = emotion_model.predict(preprocessed_face)
print(prediction)
predicted_class = np.argmax(prediction)
confidence = np.max(prediction)
# Mostrar la imagen con la predicción
# plt.imshow(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))
# plt.title(f"Emoción: {class_names[predicted_class]} ({confidence*100:.2f}%)")
# plt.axis("off")
# plt.show()
return str(class_names[predicted_class])
else:
print("No se detectó ningún rostro en la imagen.")
# if __name__ == "__main__":
# predict_face() |