agrobot_base/Python/OAK/datasets/_9_test_fastscnn.py

#!/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/<grupo>/{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()