chromadb-api / app.py
Saad0KH's picture
Update app.py
061dadc verified
raw
history blame
2.18 kB
from PIL import Image
import numpy as np
import io
import gradio as gr
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 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 = []
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).convert('RGB') # Convert to RGB
person_images.append(person_pil_image)
return person_images
def detect(image: np.ndarray) -> tuple[Image.Image, list[Image.Image]]:
if image is None:
return None, []
image = image[:, :, ::-1] # RGB -> BGR
bboxes, vbboxes = detect_person(image, detector)
res = visualize(image, bboxes, vbboxes)
person_images = extract_persons(res, bboxes)
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])
return demo
# Example usage
if __name__ == "__main__":
demo = build_gradio_interface()
demo.launch()