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)