Update app.py

app.py

@@ -4,20 +4,19 @@ import numpy as np
 from PIL import Image
 import io
 import base64
-from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+from transformers import ViTFeatureExtractor, ViTForImageClassification
 import torch
 
 app = FastAPI()
 
-# Cargar el modelo preentrenado para clasificación de género
-model_name = "nateraw/bert-imagenet"  # Cambiar a un modelo adecuado si se encuentra uno más específico
-model = AutoModelForImageClassification.from_pretrained(model_name)
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+# Cargar 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(...)):
     try:
-        # Leer y preparar la imagen
+        # Leer y procesar la imagen cargada
         image_bytes = await file.read()
         image = Image.open(io.BytesIO(image_bytes))
         img_np = np.array(image)
@@ -25,7 +24,7 @@ async def detect_face(file: UploadFile = File(...)):
         if img_np.shape[2] == 4:
             img_np = cv2.cvtColor(img_np, cv2.COLOR_BGRA2BGR)
 
-        # Detección de rostros con OpenCV
+        # Cargar el clasificador Haar para detección de rostros
         face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
         gray = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
         faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
@@ -33,38 +32,45 @@ async def detect_face(file: UploadFile = File(...)):
         if len(faces) == 0:
             raise HTTPException(status_code=404, detail="No se detectaron rostros en la imagen.")
 
-        result_data = []
+        # Procesar cada rostro detectado
+        results = []
         for (x, y, w, h) in faces:
-            # Extraer cada rostro
+            # Extraer el rostro de la imagen
             face_img = img_np[y:y+h, x:x+w]
-            face_img_pil = Image.fromarray(cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB)).resize((224, 224))
+            pil_face_img = Image.fromarray(cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB))
 
-            # Clasificar el rostro
-            inputs = feature_extractor(images=face_img_pil, return_tensors="pt")
-            outputs = model(**inputs)
-            predicted_class = outputs.logits.argmax(dim=-1).item()
-            label = model.config.id2label[predicted_class]
+            # Realizar la predicción de edad
+            inputs = transforms(pil_face_img, return_tensors='pt')
+            output = model(**inputs)
+            proba = output.logits.softmax(1)
+            preds = proba.argmax(1)
 
-            # Dibujar rectángulo y agregar datos
+            # Asumimos que la predicción está representando un rango de edad (esto puede adaptarse más tarde)
+            predicted_age_range = str(preds.item())
+
+            # Dibujar un rectángulo alrededor del rostro y añadir la edad predicha
             cv2.rectangle(img_np, (x, y), (x+w, y+h), (255, 0, 0), 2)
-            result_data.append({
-                "coordenadas": [int(x), int(y), int(w), int(h)],
-                "sexo": label
+            cv2.putText(img_np, f"Edad: {predicted_age_range}", (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
+
+            results.append({
+                "edad_predicha": predicted_age_range,
+                "coordenadas_rostro": (x, y, w, h)
             })
 
-        # Preparar imagen resultante
+        # Convertir la imagen procesada a base64
         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()
 
-        # Respuesta
         return {
-            "message": "Rostros detectados",
-            "resultados": result_data,
+            "message": "Rostros detectados y edad predicha",
+            "rostros": len(faces),
+            "resultados": results,
             "imagen_base64": base64.b64encode(img_byte_arr).decode('utf-8')
         }
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+