import glob import os import torch import torchvision.transforms as T import matplotlib.pyplot as plt import numpy as np from PIL import Image import cv2 import depthai as dai # === CONFIGURAÇÕES === model_path = 'backup/ruasModel_best.pth' num_classes = 4 labelmap_path = 'dataset/labelmap.txt' image_path = 'dataset/split/test/images/' USE_CAMERA = False device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # === FUNÇÕES AUXILIARES === def carregar_labelmap_completo(caminho): cor_para_id = {} id_para_nome = {} cores_bgr = [] with open(caminho, 'r') as arquivo: for linha in arquivo: if linha.startswith("#") or not linha.strip(): continue partes = linha.strip().split(':') if len(partes) >= 2: nome_classe, cor_rgb_str = partes[0], partes[1] r, g, b = map(int, cor_rgb_str.split(',')) cor_bgr = (b, g, r) # Corrige para BGR classe_id = len(cor_para_id) cor_para_id[cor_bgr] = classe_id cores_bgr.append(cor_bgr) id_para_nome[classe_id] = nome_classe return cor_para_id, cores_bgr, id_para_nome def segment_image(image_path, model): image = Image.open(image_path).convert("RGB") transform = T.Compose([ T.Resize((512, 512)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) input_tensor = transform(image).unsqueeze(0).to(device) model.eval() with torch.no_grad(): output = model(input_tensor)['out'] prediction = torch.argmax(output.squeeze(), dim=0).cpu().numpy() return prediction def display_segmentation(original_path, prediction): original = cv2.imread(original_path) original = cv2.resize(original, (512, 512)) _, cores_bgr, id_para_nome = carregar_labelmap_completo(labelmap_path) seg_color = np.zeros_like(original) for class_id, color in enumerate(cores_bgr): seg_color[prediction == class_id] = color overlay = cv2.addWeighted(original, 0.5, seg_color, 0.5, 0) # Legenda legenda_inicio_y = 20 for i, cor in enumerate(cores_bgr): nome = id_para_nome.get(i, f"Classe {i}") pos_y = legenda_inicio_y + i * 30 cv2.rectangle(overlay, (10, pos_y - 15), (30, pos_y + 5), cor, -1) cv2.putText(overlay, nome, (40, pos_y + 2), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 1, cv2.LINE_AA) cv2.imshow("Segmentacao - Overlay", overlay) # Remove: cv2.waitKey() e destroyAllWindows # === CARREGA MODELO === import torchvision.models.segmentation as models model = models.deeplabv3_resnet50(weights=None, num_classes=num_classes) model.load_state_dict(torch.load(model_path, map_location=device)) model.to(device) # === EXECUÇÃO EM TEMPO REAL COM OAK-D === if USE_CAMERA: print("[📷] Iniciando segmentação em tempo real com a OAK-D Lite...") # Setup da câmera pipeline = dai.Pipeline() cam_rgb = pipeline.createColorCamera() cam_rgb.setPreviewSize(640, 640) cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P) cam_rgb.setInterleaved(False) cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR) xout = pipeline.createXLinkOut() xout.setStreamName("rgb") cam_rgb.preview.link(xout.input) _, cores_bgr, id_para_nome = carregar_labelmap_completo(labelmap_path) with dai.Device(pipeline) as oak_device: queue = oak_device.getOutputQueue(name="rgb", maxSize=1, blocking=False) while True: in_rgb = queue.get() frame = in_rgb.getCvFrame() frame_resized = cv2.resize(frame, (512, 512)) # Pré-processa o frame image_pil = Image.fromarray(cv2.cvtColor(frame_resized, cv2.COLOR_BGR2RGB)) transform = T.Compose([ T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) input_tensor = transform(image_pil).unsqueeze(0).to(device) # Segmentação model.eval() with torch.no_grad(): output = model(input_tensor)['out'] prediction = torch.argmax(output.squeeze(), dim=0).cpu().numpy() # Cria imagem colorida da segmentação seg_color = np.zeros_like(frame_resized) for class_id, color in enumerate(cores_bgr): seg_color[prediction == class_id] = color overlay = cv2.addWeighted(frame_resized, 0.5, seg_color, 0.5, 0) # Legenda legenda_inicio_y = 20 for i, cor in enumerate(cores_bgr): nome = id_para_nome.get(i, f"Classe {i}") pos_y = legenda_inicio_y + i * 30 cv2.rectangle(overlay, (10, pos_y - 15), (30, pos_y + 5), cor, -1) cv2.putText(overlay, nome, (40, pos_y + 2), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 1, cv2.LINE_AA) # Mostra cv2.imshow("Segmentacao em Tempo Real", overlay) if cv2.waitKey(1) & 0xFF == ord('q'): break cv2.destroyAllWindows() else: # === SEGMENTAÇÃO POR NAVEGAÇÃO ENTRE IMAGENS === if image_path and os.path.isdir(image_path): print(f"[📁] Navegando imagens em: {image_path}") # Lista de imagens extensoes = ("*.jpg", "*.png", "*.jpeg") arquivos = [] for ext in extensoes: arquivos.extend(glob.glob(os.path.join(image_path, ext))) arquivos.sort() if not arquivos: print("[!] Nenhuma imagem encontrada na pasta de teste.") exit() indice = 0 while True: caminho_img = arquivos[indice] output_predictions = segment_image(caminho_img, model) display_segmentation(caminho_img, output_predictions) print(f"[{indice+1}/{len(arquivos)}] {os.path.basename(caminho_img)}") key = cv2.waitKey(0) & 0xFF cv2.destroyAllWindows() # ← fecha imagem ao mudar if key == ord('d'): if indice < len(arquivos) - 1: indice += 1 else: indice = 0 elif key == ord('a'): if indice > 0: indice -= 1 else: indice = len(arquivos) - 1 elif key == ord('q'): print("[👋] Saindo da visualização.") break else: print(f"[!] Tecla inválida ({key}). Use A, D ou Q.")