Update app.py

app.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-
 from __future__ import annotations
 import pathlib
 import cv2
@@ -9,7 +7,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 +18,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", "CUDAExecutionProvider"]
+        path, sess_options=options, providers=["CPUExecutionProvider"]
     )
     model = insightface.model_zoo.retinaface.RetinaFace(model_file=path, session=session)
     return model
@@ -71,32 +68,16 @@ def extract_persons(image: np.ndarray, bboxes: np.ndarray) -> list[np.ndarray]:
     return person_images
 
 
-def convert_to_pil_image(image: np.ndarray) -> Image.Image:
-    """Convert a NumPy image array to a PIL Image."""
-    return Image.fromarray(image)
-
-
-def convert_to_png(image: np.ndarray) -> bytes:
-    """Convert a NumPy image array to PNG bytes."""
-    pil_image = convert_to_pil_image(image)
-    buffer = io.BytesIO()
-    pil_image.save(buffer, format="PNG")
-    return buffer.getvalue()
-
-
 detector = load_model()
 detector.prepare(-1, nms_thresh=0.5, input_size=(640, 640))
 
 
-def detect(image: np.ndarray) -> tuple[bytes, list[bytes]]:
+def detect(image: np.ndarray) -> tuple[np.ndarray, list[np.ndarray]]:
     image = image[:, :, ::-1]  # RGB -> BGR
     bboxes, vbboxes = detect_person(image, detector)
     res = visualize(image, bboxes, vbboxes)
-    person_images = extract_persons(image, bboxes)
-    # Convert the images to PNG bytes
-    result_image_bytes = convert_to_png(res[:, :, ::-1])  # BGR -> RGB
-    person_images_bytes = [convert_to_png(person_img[:, :, ::-1]) for person_img in person_images]  # BGR -> RGB
-    return result_image_bytes, person_images_bytes
+    person_images = extract_persons(res, bboxes)  # Extract each person as a separate image
+    return res[:, :, ::-1], person_images  # BGR -> RGB
 
 
 examples = sorted(pathlib.Path("images").glob("*.jpg"))
@@ -104,7 +85,7 @@ examples = sorted(pathlib.Path("images").glob("*.jpg"))
 demo = gr.Interface(
     fn=detect,
     inputs=gr.Image(label="Input", type="numpy"),
-    outputs=[gr.Image(label="Processed Image", type="bytes"), gr.Gallery(label="Detected Persons", type="bytes")],
+    outputs=[gr.Image(label="Processed Image", type="numpy"), gr.Gallery(label="Detected Persons", type="numpy")],
     examples=examples,
     examples_per_page=30,
     title=TITLE,