import gradio as gr
import time
from counting.counting import Counting


counting = Counting("counting/apgcc.onnx")


def filter_with_threshold(scores, points, threshold):
    filtered_scores = []
    filtered_points = []
    for score, point in zip(scores, points):
        if score > threshold:
            filtered_scores.append(score)
            filtered_points.append(point)
    return filtered_scores, filtered_points


def pred(img, threshold):
    # 计算处理时间
    start_at = time.time()
    
    processed_image, processed_image_original = counting.preprocess_image(img, True)

    scores, points = counting.run_inference(processed_image)

    scores, points = filter_with_threshold(scores, points, threshold)

    draw = counting.draw_pred(processed_image_original, scores, points)

    elapsed_time = time.time() - start_at
    use_time = f"use: {elapsed_time:.3f}s"


    total = len(points)

    return draw, total, use_time


model_description = """
# APGCC People Counting

APGCC (Adaptive Perspective Guidance for Crowd Counting) 

### based on

- [APGCC](https://github.com/AaronCIH/APGCC)
"""

demo = gr.Interface(
    description=model_description,
    fn=pred, 
    inputs=["image", 
            gr.Slider(0, 1, 0.5, label="Threshold")], 
    outputs=[
        "image",
        gr.Number(label="Count"),
        gr.Textbox(label="useTime"),
    ],
    examples=[
        ["examples/crowd-001.jpg", 0.5],
        ["examples/crowd-002.jpg", 0.5],
        ["examples/image.png", 0.5],
        ["examples/image2.png", 0.5],
        ["examples/few-001.png", 0.5],
    ])
    
demo.launch()