Update app.py

app.py

@@ -9,7 +9,6 @@ import insightface
 import numpy as np
 import onnxruntime as ort
 from PIL import Image
-import io
 
 TITLE = "insightface Person Detection"
 DESCRIPTION = "https://github.com/deepinsight/insightface/tree/master/examples/person_detection"
@@ -21,7 +20,7 @@ def load_model():
     options.intra_op_num_threads = 8
     options.inter_op_num_threads = 8
     session = ort.InferenceSession(
-        path, sess_options=options, providers=["CPUExecutionProvider"]
+        path, sess_options=options, providers=["CPUExecutionProvider", "CUDAExecutionProvider"]
     )
     model = insightface.model_zoo.retinaface.RetinaFace(model_file=path, session=session)
     return model
@@ -29,10 +28,37 @@ def load_model():
 
 def detect_person(
     img: np.ndarray, detector: insightface.model_zoo.retinaface.RetinaFace
-) -> np.ndarray:
-    bboxes, _ = detector.detect(img)
+) -> tuple[np.ndarray, np.ndarray]:
+    bboxes, kpss = detector.detect(img)
     bboxes = np.round(bboxes[:, :4]).astype(int)
-    return bboxes
+    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]
+    return bboxes, vbboxes
+
+
+def visualize(image: np.ndarray, bboxes: np.ndarray, vbboxes: np.ndarray) -> np.ndarray:
+    res = image.copy()
+    for i in range(bboxes.shape[0]):
+        bbox = bboxes[i]
+        vbbox = vbboxes[i]
+        x1, y1, x2, y2 = bbox
+        vx1, vy1, vx2, vy2 = vbbox
+        cv2.rectangle(res, (x1, y1), (x2, y2), (0, 255, 0), 1)
+        alpha = 0.8
+        color = (255, 0, 0)
+        for c in range(3):
+            res[vy1:vy2, vx1:vx2, c] = res[vy1:vy2, vx1:vx2, c] * alpha + color[c] * (1.0 - alpha)
+        cv2.circle(res, (vx1, vy1), 1, color, 2)
+        cv2.circle(res, (vx1, vy2), 1, color, 2)
+        cv2.circle(res, (vx2, vy1), 1, color, 2)
+        cv2.circle(res, (vx2, vy2), 1, color, 2)
+    return res
 
 
 def extract_persons(image: np.ndarray, bboxes: np.ndarray) -> list[np.ndarray]:
@@ -53,11 +79,15 @@ detector = load_model()
 detector.prepare(-1, nms_thresh=0.5, input_size=(640, 640))
 
 
-def detect(image: np.ndarray) -> list[Image.Image]:
+def detect(image: np.ndarray) -> tuple[Image.Image, list[Image.Image]]:
     image = image[:, :, ::-1]  # RGB -> BGR
-    bboxes = detect_person(image, detector)
-    person_images = extract_persons(image, bboxes)  # Extract each person as a separate image
-    return [convert_to_pil_image(person_img[:, :, ::-1]) for person_img in person_images]  # BGR -> RGB
+    bboxes, vbboxes = detect_person(image, detector)
+    res = visualize(image, bboxes, vbboxes)
+    person_images = extract_persons(image, bboxes)
+    # Convert the images to PIL Image format
+    result_image = convert_to_pil_image(res[:, :, ::-1])  # BGR -> RGB
+    person_images_pil = [convert_to_pil_image(person_img[:, :, ::-1]) for person_img in person_images]  # BGR -> RGB
+    return result_image, person_images_pil
 
 
 examples = sorted(pathlib.Path("images").glob("*.jpg"))
@@ -65,9 +95,8 @@ examples = sorted(pathlib.Path("images").glob("*.jpg"))
 demo = gr.Interface(
     fn=detect,
     inputs=gr.Image(label="Input", type="numpy"),
-    outputs=gr.Gallery(label="Detected Persons"),  # Display multiple images in a gallery
+    outputs=[gr.Image(label="Processed Image"), gr.Gallery(label="Detected Persons")],
     examples=examples,
-    cache_examples=False,  # Disable caching of examples
     examples_per_page=30,
     title=TITLE,
     description=DESCRIPTION,