276 lines
11 KiB
Python
276 lines
11 KiB
Python
#!/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)
|
|
|
|
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()
|