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()