Update app.py

app.py

@@ -1,63 +1,16 @@
-from __future__ import annotations
+from PIL import Image
+import numpy as np
 import io
-import cv2
 import gradio as gr
-import huggingface_hub
-import insightface
-import numpy as np
-import onnxruntime as ort
-from PIL import Image
-
-TITLE = "insightface Person Detection"
-DESCRIPTION = "https://github.com/deepinsight/insightface/tree/master/examples/person_detection"
 
+def encode_pil_to_bytes(pil_image: Image.Image, format: str, **params) -> bytes:
+    with io.BytesIO() as output_bytes:
+        pil_image.save(output_bytes, format=format, **params)
+        output_bytes.seek(0)  # Rewind the BytesIO object to the beginning
+        return output_bytes.read()
 
-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"]
-    )
-    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]:
-    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]
-    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 save_pil_to_cache(pil_image: Image.Image, format: str) -> bytes:
+    return encode_pil_to_bytes(pil_image, format)
 
 def extract_persons(image: np.ndarray, bboxes: np.ndarray) -> list[Image.Image]:
     person_images = []
@@ -65,18 +18,9 @@ def extract_persons(image: np.ndarray, bboxes: np.ndarray) -> list[Image.Image]:
         x1, y1, x2, y2 = bbox
         person_image = image[y1:y2, x1:x2]  # Crop the detected person
         person_pil_image = Image.fromarray(person_image).convert('RGB')  # Convert to RGB
-        with io.BytesIO() as output:
-            person_pil_image.save(output, format='PNG')  # Save as PNG
-            output.seek(0)  # Move to the start of the BytesIO buffer
-            person_pil_image = Image.open(output)  # Reopen to ensure format
-            person_images.append(person_pil_image)
+        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]]:
     if image is None:
         return None, []
@@ -85,16 +29,32 @@ def detect(image: np.ndarray) -> tuple[Image.Image, list[Image.Image]]:
     bboxes, vbboxes = detect_person(image, detector)
     res = visualize(image, bboxes, vbboxes)
     person_images = extract_persons(res, bboxes)
-    return Image.fromarray(res[:, :, ::-1], 'RGB'), person_images  # BGR -> RGB
+    processed_image = Image.fromarray(res[:, :, ::-1], 'RGB')  # BGR -> RGB
+    return processed_image, person_images
+
+def process_image(image: Image.Image) -> tuple[Image.Image, list[Image.Image]]:
+    try:
+        np_image = np.array(image)
+        processed_image, person_images = detect(np_image)
+        return processed_image, person_images
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return None, []
+
+def build_gradio_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("## Person Detection App")
+
+        with gr.Row():
+            input_image = gr.Image(type="pil", label="Upload an Image")
+            output_image = gr.Image(type="pil", label="Processed Image")
+            gallery = gr.Gallery(label="Detected Persons")
 
+        input_image.change(fn=process_image, inputs=input_image, outputs=[output_image, gallery])
 
-demo = gr.Interface(
-    fn=detect,
-    inputs=gr.Image(label="Input", type="numpy"),
-    outputs=[gr.Image(label="Processed Image", type="numpy"), gr.Gallery(label="Detected Persons", type="numpy")],
-    title=TITLE,
-    description=DESCRIPTION,
-)
+    return demo
 
+# Example usage
 if __name__ == "__main__":
-    demo.queue(max_size=10).launch()
+    demo = build_gradio_interface()
+    demo.launch()