Update app.py

app.py

@@ -1,9 +1,7 @@
 #!/usr/bin/env python
 
 from __future__ import annotations
-
 import pathlib
-
 import cv2
 import gradio as gr
 import huggingface_hub
@@ -43,34 +41,24 @@ def detect_person(
     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]
+def extract_persons(image: np.ndarray, bboxes: np.ndarray) -> list[np.ndarray]:
+    person_images = []
+    for bbox in bboxes:
         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
+        person_image = image[y1:y2, x1:x2]  # Crop the detected person
+        person_images.append(person_image)
+    return person_images
 
 
 detector = load_model()
 detector.prepare(-1, nms_thresh=0.5, input_size=(640, 640))
 
 
-def detect(image: np.ndarray) -> np.ndarray:
+def detect(image: np.ndarray) -> list[np.ndarray]:
     image = image[:, :, ::-1]  # RGB -> BGR
     bboxes, vbboxes = detect_person(image, detector)
-    res = visualize(image, bboxes, vbboxes)
-    return res[:, :, ::-1]  # BGR -> RGB
+    person_images = extract_persons(image, bboxes)  # Extract each person as a separate image
+    return [person_img[:, :, ::-1] for person_img in person_images]  # BGR -> RGB
 
 
 examples = sorted(pathlib.Path("images").glob("*.jpg"))
@@ -78,7 +66,7 @@ examples = sorted(pathlib.Path("images").glob("*.jpg"))
 demo = gr.Interface(
     fn=detect,
     inputs=gr.Image(label="Input", type="numpy"),
-    outputs=gr.Image(label="Output"),
+    outputs=gr.Gallery(label="Detected Persons"),  # Display a gallery of cropped images
     examples=examples,
     examples_per_page=30,
     title=TITLE,