Update app.py

app.py

@@ -4,62 +4,87 @@ import numpy as np
 from PIL import Image
 import io
 from transformers import ViTFeatureExtractor, ViTForImageClassification
-import torch
 
 app = FastAPI()
 
-# Cargar el modelo y el extractor de características
-model = ViTForImageClassification.from_pretrained("nateraw/vit-age-classifier")
-transforms = ViTFeatureExtractor.from_pretrained("nateraw/vit-age-classifier")
+# Inicializamos el modelo de clasificación de edad y el extractor
+model = ViTForImageClassification.from_pretrained('nateraw/vit-age-classifier')
+transforms = ViTFeatureExtractor.from_pretrained('nateraw/vit-age-classifier')
 
-@app.post("/detect/")
-async def detect_face(file: UploadFile = File(...)):
-    # Validar tipo de archivo
-    if not file.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="El archivo no es una imagen válida.")
+# Mapeo de índices de clase a rangos de edad
+age_mapping = [
+    "0-2", "3-6", "7-9", "10-12", "13-15",
+    "16-19", "20-24", "25-29", "30-34", "35-39",
+    "40-44", "45-49", "50-54", "55-59", "60-64",
+    "65-69", "70+"
+]
 
+# Endpoint para predecir la edad de los rostros detectados en una imagen
+@app.post("/predict/")
+async def predict_age(file: UploadFile = File(...)):
+    """
+    Endpoint para predecir el rango de edad de los rostros detectados en una imagen.
+    """
     try:
-        # Leer la imagen y convertirla a numpy array
+        # Leer la imagen cargada
         image_bytes = await file.read()
-        image = Image.open(io.BytesIO(image_bytes))
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")  # Convertimos a RGB si es necesario
+
+        # Convertir la imagen a formato NumPy para usar OpenCV
         img_np = np.array(image)
 
-        # Convertir a formato adecuado (BGR) si tiene transparencia (4 canales)
-        if img_np.shape[2] == 4:
-            img_np = cv2.cvtColor(img_np, cv2.COLOR_BGRA2BGR)
+        # Cargar el clasificador Haar para detección de rostros
+        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
 
-        # Detectar rostros con Haar
-        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
-        gray = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
+        # Convertir la imagen a escala de grises para la detección de rostros
+        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
         faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
 
+        # Verificamos si se detectaron rostros
         if len(faces) == 0:
-            return {"message": "No se detectaron rostros en la imagen."}
+            raise HTTPException(status_code=404, detail="No se detectaron rostros en la imagen.")
 
+        # Lista para almacenar los resultados de predicción de cada rostro
         results = []
+
         for (x, y, w, h) in faces:
-            # Extraer el rostro y convertirlo a RGB
+            # Extraer la región del rostro
             face_img = img_np[y:y+h, x:x+w]
             pil_face_img = Image.fromarray(cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB))
 
-            # Predicción del rango de edad
-            inputs = transforms(pil_face_img, return_tensors="pt")
-            output = model(**inputs)
-            proba = output.logits.softmax(1)
-            predicted_class = proba.argmax(1).item()
-            predicted_age_range = str(predicted_class)  # Mapeo de clases a rangos
+            # Aplicar la transformación y hacer la predicción de edad
+            inputs = transforms(pil_face_img, return_tensors='pt')
+            outputs = model(**inputs)
+
+            # Calcular probabilidades y predicción
+            proba = outputs.logits.softmax(1)
+            preds = proba.argmax(1).item()  # Índice de la clase predicha
+            predicted_age_range = age_mapping[preds]
 
-            # Añadir datos al resultado
+            # Dibujar un rectángulo alrededor del rostro y agregar la edad predicha
+            cv2.rectangle(img_np, (x, y), (x + w, y + h), (255, 0, 0), 2)
+            cv2.putText(img_np, f"Edad: {predicted_age_range}", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
+
+            # Guardar el resultado de la predicción de edad y las coordenadas del rostro
             results.append({
                 "edad_predicha": predicted_age_range,
-                "coordenadas_rostro": {"x": x, "y": y, "w": w, "h": h}
+                "coordenadas_rostro": (x, y, w, h)
             })
 
+        # Convertir la imagen procesada a base64 para la respuesta
+        result_image = Image.fromarray(cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB))
+        img_byte_arr = io.BytesIO()
+        result_image.save(img_byte_arr, format='JPEG')
+        img_byte_arr = img_byte_arr.getvalue()
+
+        # Devolver los resultados
         return {
-            "message": "Rostros detectados y edad predicha.",
-            "cantidad_rostros": len(faces),
-            "detalles": results
+            "message": "Rostros detectados y edad predicha",
+            "rostros_detectados": len(faces),
+            "resultados": results,
+            "imagen_base64": base64.b64encode(img_byte_arr).decode('utf-8')
         }
 
     except Exception as e:
+        # Manejo de errores generales
         raise HTTPException(status_code=500, detail=f"Error procesando la imagen: {str(e)}")