agrobot_base/Python/OAK/datasets/_6_test_fastscnn.py

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()