import time
import gradio as gr
import pandas as pd
from PIL import Image, ImageDraw
from transformers import pipeline


print("Завантаження моделі детекції для Safety Counter...")
detector = pipeline("object-detection", model="facebook/detr-resnet-50")


COLORS = ["#e74c3c", "#3498db", "#2ecc71", "#f39c12", "#9b59b6", "#1abc9c"]

def get_color(label):
    return COLORS[hash(label) % len(COLORS)]

def safety_monitor(image, threshold, max_allowed_people):
    if image is None:
        return None, 0, "⚠️ Будь ласка, завантажте зображення.", pd.DataFrame()

    th = float(threshold)
    
    results = detector(image, threshold=th)
    
    img_draw = image.copy()
    draw = ImageDraw.Draw(img_draw)
    
    people_count = 0
    counts = {}

    for obj in results:
        box = obj["box"]
        label = obj["label"]
        score = obj["score"]
        
        counts[label] = counts.get(label, 0) + 1
        
        
        if label == "person":
            people_count += 1
            color = "#ff3333" 
        else:
            color = get_color(label)
        
        draw.rectangle([box["xmin"], box["ymin"], box["xmax"], box["ymax"]], outline=color, width=3)
        text = f"{label} {score:.0%}"
        draw.rectangle([box["xmin"], box["ymin"] - 20, box["xmin"] + len(text)*7, box["ymin"]], fill=color)
        draw.text((box["xmin"] + 3, box["ymin"] - 18), text, fill="white")

    
    if people_count == 0:
        safety_status = "🟢 Скупчень людей не виявлено (0 осіб)."
    elif people_count <= max_allowed_people:
        safety_status = f"🟢 Безпечно. Кількість людей ({people_count}) в межах норми (макс. {max_allowed_people})."
    else:
        safety_status = f"🚨 ПОПЕРЕДЖЕННЯ! Перевищено ліміт людей! Виявлено {people_count} осіб (Дозволено: {max_allowed_people})."

    
    df_details = pd.DataFrame([
        {
            "Клас": r["label"],
            "Впевненість": f"{r['score']:.1%}",
            "Координати (xmin, ymin)": f"({int(r['box']['xmin'])}, {int(r['box']['ymin'])})"
        } for r in results
    ])

    return img_draw, people_count, safety_status, df_details



with gr.Blocks(title="Safety People Counter", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚨 AI Safety Counter & Crowd Monitor")
    gr.Markdown(
        "Спеціалізований застосунок для моніторингу безпеки та підрахунку кількості людей у приміщеннях або на вулиці."
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            img_in = gr.Image(type="pil", label="Камера моніторингу")
            
            
            threshold_sl = gr.Slider(0.1, 0.9, value=0.4, step=0.05, label="Поріг детекції")
            max_people_sl = gr.Slider(1, 20, value=5, step=1, label="Макс. дозволена кількість людей")
            
            btn = gr.Button("Просканувати зону ▶", variant="primary")
            
        with gr.Column(scale=1):
            img_out = gr.Image(label="Візуалізація детекції")
            
            with gr.Row():
                people_counter = gr.Number(label="ЗНАЙДЕНО ЛЮДЕЙ", value=0, precision=0)
            
            safety_msg = gr.Markdown("### Статус: Очікування сканування...")

    gr.Markdown("---")
    with gr.Row():
        detail_df = gr.DataFrame(label="Повний звіт системи детекції", interactive=False)


    btn.click(
        fn=safety_monitor,
        inputs=[img_in, threshold_sl, max_people_sl],
        outputs=[img_out, people_counter, safety_msg, detail_df],
    )

demo.launch()