#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Teste/visualização do FastSCNN com suporte a estrutura AGRUPADA. Estrutura esperada (nova): MODELO/dataset/split/test/group//{images,masks} Fallback (legado, se não houver 'group/'): MODELO/dataset/split/test/{images,masks} Também suporta o modo câmera (--camera) igual ao original. """ 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"] model_to_use = config["model_to_use"] dataset_path = os.path.join(MODELO, "dataset") split_folder = "val" labelmap_path = os.path.join(dataset_path, "labelmap.txt") model_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME) model_name = "" if model_to_use == "geral": model_name = f"{MODEL_NAME}_best.pth" elif model_to_use == "main_class": MAIN_CLASS_NAME = config["main_class_name"] model_name = f"{MODEL_NAME}_best_f1_{MAIN_CLASS_NAME}.pth" elif model_to_use == "es": ES_CLASSES_NAME = config["es_classes"] model_name = f"{MODEL_NAME}_best_es_{ES_CLASSES_NAME}.pth" else: model_name = f"{MODEL_NAME}_best.pth" #model_name = model_name.replace(".pth", "_bkp.pth") IMG_EXTS = (".jpg", ".jpeg", ".png") MSK_EXTS = (".png", ".jpg", ".jpeg") # preferir .png quando houver def infer_ignore_id(ignore_rgb, default_id=255): """Tenta inferir ID de ignore a partir do labelmap.""" if isinstance(ignore_rgb, (list, tuple)): if len(ignore_rgb) == 1 and isinstance(ignore_rgb[0], (int, np.integer)): return int(ignore_rgb[0]) if len(ignore_rgb) == 3: return default_id if isinstance(ignore_rgb, (int, np.integer)): return int(ignore_rgb) return default_id def list_groups(group_root): if not os.path.isdir(group_root): return [] out = [] for g in sorted(os.listdir(group_root)): gdir = os.path.join(group_root, g) if os.path.isdir(os.path.join(gdir, "images")) and os.path.isdir(os.path.join(gdir, "masks")): out.append(g) return out def mask_for_base(msk_dir, base): """Encontra máscara correspondente, priorizando .png.""" best = None for ext in MSK_EXTS: cand = os.path.join(msk_dir, base + ext) if os.path.isfile(cand): if best is None: best = cand if os.path.splitext(cand)[1].lower() == ".png": return cand return best def collect_pairs_grouped(test_root, want_groups=None): """Coleta pares img/mask de test_root com estrutura 'group/'.""" group_root = os.path.join(test_root, "group") if not os.path.isdir(group_root): return [], [] groups = list_groups(group_root) if want_groups: filt = {g.strip() for g in want_groups.split(",") if g.strip()} groups = [g for g in groups if g in filt] imgs, msks, groups_idx = [], [], [] for g in groups: img_dir = os.path.join(group_root, g, "images") msk_dir = os.path.join(group_root, g, "masks") for p in sorted(glob.glob(os.path.join(img_dir, "*"))): base, ext = os.path.splitext(os.path.basename(p)) if ext.lower() not in IMG_EXTS: continue m = mask_for_base(msk_dir, base) if m: imgs.append(p) msks.append(m) groups_idx.append(g) return imgs, msks, groups_idx def collect_pairs_legacy(test_root): """Coleta pares img/mask sem 'group/'.""" img_dir = os.path.join(test_root, "images") msk_dir = os.path.join(test_root, "masks") imgs = [] msks = [] groups_idx = [] for p in sorted(glob.glob(os.path.join(img_dir, "*"))): base, ext = os.path.splitext(os.path.basename(p)) if ext.lower() not in IMG_EXTS: continue m = mask_for_base(msk_dir, base) if m: imgs.append(p) msks.append(m) groups_idx.append("legacy") return imgs, msks, groups_idx def main(): parser = argparse.ArgumentParser() parser.add_argument("--camera", action="store_true", help="Usar câmera em vez de imagens") parser.add_argument("--groups", type=str, default=None, help="Filtrar grupos (ex: chao,erva_cana)") 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 = infer_ignore_id(ignore_rgb, default_id=255) model = FastSCNN(num_classes=len(classes)) model.load_state_dict(torch.load(os.path.join(model_path, model_name), 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: # === Modo câmera (inalterado) === 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) 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 = 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) 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 = 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 imagens (agrupado + fallback) === #test_root = os.path.join(dataset_path, "split", split_folder) test_root = os.path.join(dataset_path, "512x288") image_paths, mask_paths, groups_idx = collect_pairs_grouped(test_root, want_groups=args.groups) if not image_paths: image_paths, mask_paths, groups_idx = collect_pairs_legacy(test_root) assert len(image_paths) == len(mask_paths) and len(image_paths) > 0, "Nenhuma imagem/máscara encontrada no split de teste." idx = 0 while True: img_path = image_paths[idx] mask_path = mask_paths[idx] grupo = groups_idx[idx] if groups_idx else "?" 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) 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) # Rotula o grupo na imagem cv2.putText(resultado_completo, f"grupo: {grupo}", (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,255,0), 2) 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()