test.py

import gradio as gr
from ultralytics import YOLO
import torch
import cv2
import numpy as np
import os
import math
import pandas as pd
from datetime import datetime

# Caminho para o modelo treinado
modelo_path = "runs/segment/meu_yolo_seg20/weights/best.pt"
model_cache = {}

def carregar_modelo(device):
    if device not in model_cache:
        model = YOLO(modelo_path)
        model.to(device)
        model_cache[device] = model
    return model_cache[device]

def contar_pixels_conexos(mask, minimo):
    bin_mask = (mask > 0).astype(np.uint8)
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(bin_mask)
    return sum(stat[cv2.CC_STAT_AREA] for stat in stats[1:] if stat[cv2.CC_STAT_AREA] >= minimo)

def calcular_distancia(p1, p2):
    return math.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)

def detectar_escala(img_bgr):
    h, w = img_bgr.shape[:2]
    roi = img_bgr[h-100:h-10, int(w*0.6):w]
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    lines = cv2.HoughLinesP(edges, rho=1, theta=np.pi/180, threshold=30,
                            minLineLength=30, maxLineGap=5)
    longest = 0
    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            if abs(y2 - y1) < 5:
                length = abs(x2 - x1)
                if length > longest:
                    longest = length
    return 1.0 / longest if longest > 0 else None

def processar_uma_imagem(imagem_path, dispositivo, conf_min, pixels_minimos):
    device = "cuda" if dispositivo == "GPU" and torch.cuda.is_available() else "cpu"
    model = carregar_modelo(device)
    results = model.predict(source=imagem_path, device=device, conf=conf_min, verbose=False)
    r = results[0]
    img_bgr = cv2.imread(imagem_path)
    if img_bgr is None:
        return None, None, "Erro ao carregar imagem."
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    h, w = img_rgb.shape[:2]
    mm_por_pixel = detectar_escala(img_bgr)
    mascaras_binarias = []
    if r.masks is not None and r.masks.data is not None:
        for mask in r.masks.data:
            m = mask.cpu().numpy().astype(np.uint8) * 255
            m = cv2.resize(m, (w, h), interpolation=cv2.INTER_NEAREST)
            if contar_pixels_conexos(m, pixels_minimos) > 0:
                mascaras_binarias.append((m > 0).astype(np.uint8))
    mascaras_fundidas = []
    for m in mascaras_binarias:
        fundida = False
        for i in range(len(mascaras_fundidas)):
            if np.count_nonzero(cv2.bitwise_and(m, mascaras_fundidas[i])) > 0:
                mascaras_fundidas[i] = cv2.bitwise_or(mascaras_fundidas[i], m)
                fundida = True
                break
        if not fundida:
            mascaras_fundidas.append(m)
    centroides = []
    for m in mascaras_fundidas:
        M = cv2.moments(m)
        if M["m00"] != 0:
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
            centroides.append((cx, cy))
            cv2.circle(img_rgb, (cx, cy), 12, (0, 255, 0), 2)
    distancia_px = None
    distancia_mm = None
    if len(centroides) == 2:
        cv2.line(img_rgb, centroides[0], centroides[1], (0, 0, 255), 2)
        for c in centroides:
            cv2.circle(img_rgb, c, 6, (0, 0, 255), -1)
        distancia_px = calcular_distancia(centroides[0], centroides[1])
        if mm_por_pixel:
            distancia_mm = distancia_px * mm_por_pixel
    nome = os.path.basename(imagem_path)
    output_path = os.path.splitext(nome)[0] + "_out.png"
    cv2.imwrite(output_path, cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR))
    return output_path, nome, distancia_px, distancia_mm

def processar_varias(imagens, dispositivo, conf_min, pixels_minimos):
    imagens_resultado = []
    nomes, pxs, mms = [], [], []
    for imagem in imagens:
        out_path, nome, dist_px, dist_mm = processar_uma_imagem(imagem, dispositivo, conf_min, pixels_minimos)
        imagens_resultado.append(out_path)
        nomes.append(nome)
        pxs.append(f"{dist_px:.1f}" if dist_px else "-")
        mms.append(f"{dist_mm:.3f}" if dist_mm else "-")
    valores_mm = [float(x) for x in mms if x != "-"]
    media = np.mean(valores_mm) if valores_mm else 0
    desvio = np.std(valores_mm) if valores_mm else 0
    tabela = pd.DataFrame({"Imagem": nomes, "Distância (px)": pxs, "Distância (mm)": mms})
    resumo_df = pd.DataFrame({"Métrica": ["Média", "Desvio padrão"], "Valor (mm)": [f"{media:.3f}", f"{desvio:.3f}"]})
    agora = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    resultados_xlsx_path = f"resultados_medidas_{agora}.xlsx"
    with pd.ExcelWriter(resultados_xlsx_path) as writer:
        tabela.to_excel(writer, sheet_name="Resultados", index=False)
        resumo_df.to_excel(writer, sheet_name="Resumo", index=False)
    # Criar pasta para imagens
    pasta_saida = f"imagens_processadas_{agora}"
    os.makedirs(pasta_saida, exist_ok=True)
    for img_path in imagens_resultado:
        novo_caminho = os.path.join(pasta_saida, os.path.basename(img_path))
        os.replace(img_path, novo_caminho)
    imagens_resultado = [os.path.join(pasta_saida, os.path.basename(p)) for p in imagens_resultado]
    legendas = [f"{n} - {m} mm" for n, m in zip(nomes, mms)]
    return list(zip(imagens_resultado, legendas)), tabela, resumo_df, resultados_xlsx_path

# Interface Gradio
with gr.Blocks(title="Medição de asas múltiplas") as demo:
    gr.Markdown("## Medição automática de asas de abelha com normalização à escala impressa")
    imagens_input = gr.File(label="Subir imagens", file_types=[".png", ".jpg", ".jpeg"], file_count="multiple")
    dispositivo = gr.Radio(["CPU", "GPU"], label="Dispositivo", value="CPU", visible=False)
    conf_slider = gr.Slider(0.01, 1.0, value=0.1, step=0.01, label="Confiança mínima", visible=False)
    pixel_slider = gr.Slider(1, 1000, value=10, step=1, label="Pixels conectados mínimos por máscara", visible=False)
    imagens_saida = gr.Gallery(label="Resultados das imagens segmentadas", columns=4, preview=True, show_label=True)
    tabela_resultados = gr.Dataframe(label="Resultados", interactive=False, headers=["Imagem", "Distância (px)", "Distância (mm)"])
    tabela_resumo = gr.Dataframe(label="Resumo estatístico", interactive=False, headers=["Métrica", "Valor (mm)"])
    download_btn = gr.File(label="Download XLSX")
    imagens_input.change(
        fn=lambda imagens: processar_varias(imagens, "CPU", 0.1, 10),
        inputs=imagens_input,
        outputs=[imagens_saida, tabela_resultados, tabela_resumo, download_btn]
    )
demo.launch(share=False)