Update app.py

app.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import pathlib
+import io
 import cv2
 import gradio as gr
 import huggingface_hub
@@ -11,6 +12,7 @@ from PIL import Image
 TITLE = "insightface Person Detection"
 DESCRIPTION = "https://github.com/deepinsight/insightface/tree/master/examples/person_detection"
 
+
 def load_model():
     path = huggingface_hub.hf_hub_download("public-data/insightface", "models/scrfd_person_2.5g.onnx")
     options = ort.SessionOptions()
@@ -22,6 +24,7 @@ def load_model():
     model = insightface.model_zoo.retinaface.RetinaFace(model_file=path, session=session)
     return model
 
+
 def detect_person(
     img: np.ndarray, detector: insightface.model_zoo.retinaface.RetinaFace
 ) -> tuple[np.ndarray, np.ndarray]:
@@ -37,6 +40,7 @@ def detect_person(
     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]):
@@ -55,18 +59,21 @@ def visualize(image: np.ndarray, bboxes: np.ndarray, vbboxes: np.ndarray) -> np.
         cv2.circle(res, (vx2, vy2), 1, color, 2)
     return res
 
+
 def extract_persons(image: np.ndarray, bboxes: np.ndarray) -> list[Image.Image]:
     person_images = []
     for bbox in bboxes:
         x1, y1, x2, y2 = bbox
         person_image = image[y1:y2, x1:x2]  # Crop the detected person
-        person_pil_image = Image.fromarray(person_image)
+        person_pil_image = Image.fromarray(person_image).convert('RGB')  # Convert to RGB
         person_images.append(person_pil_image)
     return person_images
 
+
 detector = load_model()
 detector.prepare(-1, nms_thresh=0.5, input_size=(640, 640))
 
+
 def detect(image: np.ndarray) -> tuple[Image.Image, list[Image.Image]]:
     image = image[:, :, ::-1]  # RGB -> BGR
     bboxes, vbboxes = detect_person(image, detector)
@@ -74,6 +81,7 @@ def detect(image: np.ndarray) -> tuple[Image.Image, list[Image.Image]]:
     person_images = extract_persons(res, bboxes)
     return Image.fromarray(res[:, :, ::-1], 'RGB'), person_images  # BGR -> RGB
 
+
 examples = sorted(pathlib.Path("images").glob("*.jpg"))
 
 demo = gr.Interface(