160 lines
6.6 KiB
Python
160 lines
6.6 KiB
Python
import os
|
|
import time
|
|
import cv2
|
|
import glob
|
|
import argparse
|
|
import torch
|
|
import numpy as np
|
|
import depthai as dai
|
|
from PIL import Image
|
|
from fast_scnn import FastSCNN
|
|
from utils import carregar_labelmap_completo, compute_roi_indices, converter_mask_ids_para_rgb, desenhar_legenda_horizontal, desenhar_legenda_vertical, resize_keep_width
|
|
|
|
# ⚙️ Configurações
|
|
MODELO = "oak-1"
|
|
MODEL_NAME = "ervas_full"
|
|
RESOLUCAO = (384, 384)
|
|
ROI_INICIO = 0.0
|
|
ROI_TAMANHO = 1.0
|
|
dataset_path = os.path.join(MODELO, "dataset")
|
|
split_folder = "test"
|
|
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
|
|
model_path = os.path.join(MODELO, "backup", "fast_scnn", MODEL_NAME, MODEL_NAME + "_best.pth")
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--camera", action="store_true", help="Usar câmera em vez de imagens")
|
|
args = parser.parse_args()
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
_, colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
|
ignore_id = ignore_rgb[0]
|
|
|
|
model = FastSCNN(num_classes=len(classes))
|
|
model.load_state_dict(torch.load(model_path, map_location=device))
|
|
model.to(device).eval()
|
|
|
|
mean = torch.tensor([0.485, 0.456, 0.406]).reshape(3, 1, 1).to(device)
|
|
std = torch.tensor([0.229, 0.224, 0.225]).reshape(3, 1, 1).to(device)
|
|
|
|
if args.camera:
|
|
# --- Criar pipeline da OAK-1 Lite W ---
|
|
pipeline = dai.Pipeline()
|
|
cam_rgb = pipeline.createColorCamera()
|
|
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
|
|
cam_rgb.setBoardSocket(dai.CameraBoardSocket.RGB)
|
|
cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.RGB)
|
|
cam_rgb.setInterleaved(False)
|
|
cam_rgb.setFps(30)
|
|
|
|
xout_rgb = pipeline.createXLinkOut()
|
|
xout_rgb.setStreamName("rgb")
|
|
cam_rgb.video.link(xout_rgb.input)
|
|
|
|
# --- Conectar dispositivo ---
|
|
with dai.Device(pipeline) as oak_device:
|
|
rgb_queue = oak_device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
|
|
|
|
prev_time = time.time()
|
|
while True:
|
|
in_rgb = rgb_queue.get()
|
|
frame = in_rgb.getCvFrame()
|
|
|
|
H, W = frame.shape[:2]
|
|
y_fim, y_inicio = compute_roi_indices(H, ROI_INICIO, ROI_TAMANHO)
|
|
|
|
roi = frame[y_fim:y_inicio, 0:W]
|
|
roi_resized = resize_keep_width(roi, RESOLUCAO[1], RESOLUCAO[0])
|
|
roi_norm = roi_resized.astype(np.float32) / 255.0
|
|
|
|
roi_tensor = torch.from_numpy(roi_norm).permute(2, 0, 1).unsqueeze(0).to(device)
|
|
roi_tensor = (roi_tensor - mean) / std
|
|
|
|
with torch.no_grad():
|
|
pred = model(roi_tensor)
|
|
pred_ids = torch.argmax(pred.squeeze(), dim=0).cpu().numpy()
|
|
|
|
pred_rgb = converter_mask_ids_para_rgb(pred_ids, colormap_rgb, ignore_id)
|
|
pred_rgb_resized = cv2.resize(pred_rgb, (roi.shape[1], roi.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
|
|
# Overlay
|
|
overlay = frame.copy()
|
|
overlay[y_fim:y_inicio, 0:W] = cv2.addWeighted(overlay[y_fim:y_inicio, 0:W], 0.4, pred_rgb_resized, 0.6, 0)
|
|
|
|
# FPS
|
|
now = time.time()
|
|
fps = 1.0 / (now - prev_time)
|
|
prev_time = now
|
|
cv2.putText(overlay, f"FPS: {fps:.1f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
|
|
|
|
# === LEGENDA SOBRE A IMAGEM DA CÂMERA ===
|
|
legenda = desenhar_legenda_vertical(colormap_rgb, classes)
|
|
legenda_resized = cv2.resize(legenda, (150, 30 * len(colormap_rgb)), interpolation=cv2.INTER_AREA)
|
|
|
|
h, w = overlay.shape[:2]
|
|
h_leg, w_leg = legenda_resized.shape[:2]
|
|
x_offset = w - w_leg - 10
|
|
y_offset = h - h_leg - 30
|
|
|
|
overlay[y_offset:y_offset + h_leg, x_offset:x_offset + w_leg] = legenda_resized
|
|
|
|
cv2.imshow("Segmentação Fast-SCNN (OAK + PyTorch)", cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
cv2.destroyAllWindows()
|
|
else:
|
|
# Modo normal com imagens da pasta
|
|
image_paths = sorted(glob.glob(os.path.join(dataset_path, "split", split_folder, "images", "*")))
|
|
mask_paths = sorted(glob.glob(os.path.join(dataset_path, "split", split_folder, "masks", "*")))
|
|
assert len(image_paths) == len(mask_paths) and len(image_paths) > 0
|
|
|
|
idx = 0
|
|
while True:
|
|
img_path = image_paths[idx]
|
|
mask_path = mask_paths[idx]
|
|
|
|
img_rgb = np.array(Image.open(img_path).convert("RGB"))
|
|
mask_gt = np.array(Image.open(mask_path).convert("L"))
|
|
|
|
H, W = img_rgb.shape[:2]
|
|
y_fim, y_inicio = compute_roi_indices(H, ROI_INICIO, ROI_TAMANHO)
|
|
|
|
img_roi = img_rgb[y_fim:y_inicio, 0:W]
|
|
mask_roi = mask_gt[y_fim:y_inicio, 0:W]
|
|
img_resized = resize_keep_width(img_roi, RESOLUCAO[1], RESOLUCAO[0])
|
|
mask_resized = resize_keep_width(mask_roi, RESOLUCAO[1], RESOLUCAO[0])
|
|
|
|
img_norm = img_resized.astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_norm).permute(2, 0, 1).unsqueeze(0).to(device)
|
|
img_tensor = (img_tensor - mean) / std
|
|
img_tensor = img_tensor.float()
|
|
|
|
with torch.no_grad():
|
|
pred = model(img_tensor)
|
|
pred_ids = torch.argmax(pred.squeeze(), dim=0).cpu().numpy()
|
|
|
|
pred_rgb = converter_mask_ids_para_rgb(pred_ids, colormap_rgb, ignore_id)
|
|
mask_gt_rgb = converter_mask_ids_para_rgb(mask_resized, colormap_rgb, ignore_id)
|
|
|
|
resultado = np.concatenate([img_resized, mask_gt_rgb, pred_rgb], axis=1)
|
|
|
|
# Adiciona legenda abaixo
|
|
legenda = desenhar_legenda_horizontal(colormap_rgb, classes)
|
|
legenda_resized = cv2.resize(legenda, (resultado.shape[1], legenda.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
resultado_completo = np.concatenate([resultado, legenda_resized], axis=0)
|
|
cv2.imshow("Original | GroundTruth | Predito", cv2.cvtColor(resultado_completo, cv2.COLOR_RGB2BGR))
|
|
key = cv2.waitKey(0) & 0xFF
|
|
|
|
if key == ord('q'):
|
|
break
|
|
elif key == ord('d'):
|
|
idx = (idx + 1) % len(image_paths)
|
|
elif key == ord('a'):
|
|
idx = (idx - 1 + len(image_paths)) % len(image_paths)
|
|
|
|
cv2.destroyAllWindows()
|
|
|
|
if __name__ == "__main__":
|
|
main() |