194 lines
6.6 KiB
Python
194 lines
6.6 KiB
Python
|
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.")
|