Update app.py

app.py

@@ -2,14 +2,14 @@ from flask import Flask, request, jsonify
 from PIL import Image
 import base64
 from io import BytesIO
-import io 
-from SegCloth import segment_clothing
 import numpy as np
 import cv2
 import insightface
 import onnxruntime as ort
 import huggingface_hub
 
+app = Flask(__name__)
+
 # Charger le modèle
 def load_model():
     path = huggingface_hub.hf_hub_download("public-data/insightface", "models/scrfd_person_2.5g.onnx")
@@ -25,12 +25,27 @@ def load_model():
 detector = load_model()
 detector.prepare(-1, nms_thresh=0.5, input_size=(640, 640))
 
+# Fonction pour décoder une image encodée en base64 en objet PIL.Image.Image
+def decode_image_from_base64(image_data):
+    image_data = base64.b64decode(image_data)
+    image = Image.open(BytesIO(image_data)).convert("RGB")  # Convertir en RGB pour éviter les problèmes de canal alpha
+    return image
+
+# Fonction pour encoder une image PIL en base64
+def encode_image_to_base64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode('utf-8')
+
 # Détecter les personnes et extraire les images
 def detect_person(image):
     img = np.array(image)
     img = img[:, :, ::-1]  # RGB -> BGR
 
     bboxes, kpss = detector.detect(img)
+    if bboxes.shape[0] == 0:  # Aucun visage détecté
+        return []
+
     bboxes = np.round(bboxes[:, :4]).astype(int)
     kpss = np.round(kpss).astype(int)
     kpss[:, :, 0] = np.clip(kpss[:, :, 0], 0, img.shape[1])
@@ -54,49 +69,18 @@ def detect_person(image):
 
     return person_images
 
-
-
-# Créer une instance FastAPI
-app = Flask(__name__)
-
-def segment_image(img, clothes):
-    img = decode_image_from_base64(img)
-    return segment_clothing(img, clothes)
-
-# Fonction pour décoder une image encodée en base64 en objet PIL.Image.Image
-def decode_image_from_base64(image_data):
-    image_data = base64.b64decode(image_data)
-    image = Image.open(io.BytesIO(image_data))
-    return image
-
-def encode_image_to_base64(image):
-    buffered = BytesIO()
-    image.save(buffered, format="PNG")
-    return base64.b64encode(buffered.getvalue()).decode('utf-8')
-
-@app.get("/")
-def root():
-    return "Welcome to  Fashion Clothing Detectiony API!"
-
-# Route pour l'API REST
-@app.route('/api/classify', methods=['POST'])
-def classify():
-    data = request.json
-    print(data)
-    clothes = ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress"]
-    image = data['image']
-    result = segment_image(image,clothes)
-    return jsonify({'result': result})
-
 @app.route('/api/detect', methods=['POST'])
 def detect():
-    data = request.json
-    image_base64 = data['image']
-    image = decode_image_from_base64(image_base64)
-    
-    person_images_base64 = detect_person(image)
-    
-    return jsonify({'images': person_images_base64})
+    try:
+        data = request.json
+        image_base64 = data['image']
+        image = decode_image_from_base64(image_base64)
+        
+        person_images_base64 = detect_person(image)
+        
+        return jsonify({'images': person_images_base64})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
 
 if __name__ == "__main__":
-    app.run(debug=True, host="0.0.0.0", port=7860)
+    app.run(debug=True, host="0.0.0.0", port=7860)