import json 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 with open("config.json", "r") as f: config = json.load(f) MODELO = config["camera"] MODEL_NAME = config["model_name"] RESOLUCAO = config["resolucao"] ROI_INICIO = config["roi_inicio"] ROI_TAMANHO = config["roi_tamanho"] MAIN_CLASS_NAME = config["main_class_name"] use_main_class = config["use_main_class"] 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", config["modelo"], MODEL_NAME, f"{MODEL_NAME}_best{f'_f1_{MAIN_CLASS_NAME}' if use_main_class else ''}.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[0], RESOLUCAO[1], cv2.INTER_AREA) 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[0], RESOLUCAO[1], cv2.INTER_AREA) mask_resized = resize_keep_width(mask_roi, RESOLUCAO[0], RESOLUCAO[1], cv2.INTER_NEAREST) 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()