Update app.py

app.py

@@ -4,6 +4,55 @@ 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
+
+# Charger le modèle
+def load_model():
+    path = huggingface_hub.hf_hub_download("public-data/insightface", "models/scrfd_person_2.5g.onnx")
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 8
+    options.inter_op_num_threads = 8
+    session = ort.InferenceSession(
+        path, sess_options=options, providers=["CPUExecutionProvider", "CUDAExecutionProvider"]
+    )
+    model = insightface.model_zoo.retinaface.RetinaFace(model_file=path, session=session)
+    return model
+
+detector = load_model()
+detector.prepare(-1, nms_thresh=0.5, input_size=(640, 640))
+
+# 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)
+    bboxes = np.round(bboxes[:, :4]).astype(int)
+    kpss = np.round(kpss).astype(int)
+    kpss[:, :, 0] = np.clip(kpss[:, :, 0], 0, img.shape[1])
+    kpss[:, :, 1] = np.clip(kpss[:, :, 1], 0, img.shape[0])
+    vbboxes = bboxes.copy()
+    vbboxes[:, 0] = kpss[:, 0, 0]
+    vbboxes[:, 1] = kpss[:, 0, 1]
+    vbboxes[:, 2] = kpss[:, 4, 0]
+    vbboxes[:, 3] = kpss[:, 4, 1]
+
+    person_images = []
+    for i in range(bboxes.shape[0]):
+        bbox = bboxes[i]
+        x1, y1, x2, y2 = bbox
+        person_img = img[y1:y2, x1:x2]
+        
+        # Convert numpy array to PIL Image and encode to base64
+        pil_img = Image.fromarray(person_img[:, :, ::-1])  # BGR -> RGB
+        base64_img = encode_image_to_base64(pil_img)
+        person_images.append(base64_img)
+
+    return person_images
 
 
 
@@ -39,5 +88,15 @@ def classify():
     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})
+
 if __name__ == "__main__":
     app.run(debug=True, host="0.0.0.0", port=7860)