816 lines
28 KiB
Python
816 lines
28 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Teste/visualização do SegFormer (.pt) com suporte a estrutura AGRUPADA.
|
||
|
||
- Modo imagens (dataset/512x288 ou split/test/...)
|
||
- Modo câmera (--camera) no mesmo estilo do _9_test_fastscnn.py
|
||
- ROI + resize_keep_width + overlay + legenda
|
||
|
||
Requisitos:
|
||
pip install transformers timm
|
||
|
||
Obs:
|
||
Este script assume que seus ids de classe batem com o labelmap.txt (mask em IDs 0..K-1).
|
||
"""
|
||
|
||
from collections import deque
|
||
from enum import IntEnum
|
||
import json
|
||
import os
|
||
import time
|
||
from typing import Dict, Tuple
|
||
import cv2
|
||
import glob
|
||
import argparse
|
||
import torch
|
||
import numpy as np
|
||
import depthai as dai
|
||
from PIL import Image
|
||
|
||
from transformers import SegformerForSemanticSegmentation
|
||
|
||
# Reaproveita utilidades do seu projeto (iguais no script do FastSCNN)
|
||
from utils import (
|
||
carregar_labelmap_completo, compute_roi_indices, converter_mask_ids_para_rgb,
|
||
desenhar_legenda_horizontal, desenhar_legenda_vertical, resize_keep_width
|
||
)
|
||
|
||
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 (mesma ideia do script FastSCNN)."""
|
||
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 _extract_state_dict(ckpt):
|
||
"""
|
||
Aceita:
|
||
- state_dict puro (dict de tensores)
|
||
- checkpoint com chaves comuns: state_dict / model_state_dict / model
|
||
"""
|
||
if not isinstance(ckpt, dict):
|
||
return None
|
||
|
||
# caso já seja um state_dict puro
|
||
if any(isinstance(v, torch.Tensor) for v in ckpt.values()):
|
||
return ckpt
|
||
|
||
for k in ("state_dict", "model_state_dict", "model"):
|
||
if k in ckpt and isinstance(ckpt[k], dict):
|
||
return ckpt[k]
|
||
|
||
return None
|
||
|
||
|
||
def load_segformer_from_checkpoint(
|
||
pt_path: str,
|
||
backbone: str,
|
||
num_classes: int,
|
||
device: torch.device,
|
||
):
|
||
"""
|
||
Carrega um SegFormer (B0, B1, B2, B3...) compatível com o treino:
|
||
|
||
- Cria o modelo via from_pretrained(backbone, num_labels=num_classes)
|
||
- Carrega o state_dict salvo pelo script de treino
|
||
"""
|
||
ckpt = torch.load(pt_path, map_location="cpu", weights_only=True)
|
||
state_dict = _extract_state_dict(ckpt)
|
||
if state_dict is None:
|
||
raise RuntimeError(f"Não consegui extrair state_dict de {pt_path}. keys={list(ckpt.keys())}")
|
||
|
||
# limpar prefixos comuns
|
||
cleaned = {}
|
||
for k, v in state_dict.items():
|
||
nk = k
|
||
if nk.startswith("model."):
|
||
nk = nk[len("model."):]
|
||
if nk.startswith("module."):
|
||
nk = nk[len("module."):]
|
||
cleaned[nk] = v
|
||
|
||
# Cria o modelo igual ao treino (_8_train_segformer_b3.py)
|
||
model = SegformerForSemanticSegmentation.from_pretrained(
|
||
backbone,
|
||
num_labels=num_classes,
|
||
ignore_mismatched_sizes=True,
|
||
use_safetensors=True
|
||
)
|
||
|
||
missing, unexpected = model.load_state_dict(cleaned, strict=False)
|
||
print(f"[load] missing={len(missing)} unexpected={len(unexpected)}")
|
||
if missing:
|
||
print("[load] missing sample:", missing[:10])
|
||
if unexpected:
|
||
print("[load] unexpected sample:", unexpected[:10])
|
||
|
||
model.to(device).eval()
|
||
return model
|
||
|
||
|
||
@torch.no_grad()
|
||
def segformer_predict_ids(model, img_tensor):
|
||
"""
|
||
img_tensor: [1,3,H,W] float32 normalizado.
|
||
retorna: pred_ids [H,W] (numpy int)
|
||
"""
|
||
out = model(pixel_values=img_tensor)
|
||
logits = out.logits # [B, C, h, w] (pode ser menor que input)
|
||
# Upsample logits para o tamanho do input
|
||
logits = torch.nn.functional.interpolate(
|
||
logits,
|
||
size=img_tensor.shape[-2:],
|
||
mode="bilinear",
|
||
align_corners=False
|
||
)
|
||
pred = torch.argmax(logits, dim=1) # [B,H,W]
|
||
return pred.squeeze(0).cpu().numpy().astype(np.uint8)
|
||
|
||
|
||
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)")
|
||
parser.add_argument("--split_folder", type=str, default="val", help="split padrão (se usar split/val/test)")
|
||
parser.add_argument("--test_folder", type=str, default=None, help="pasta para teste")
|
||
args = parser.parse_args()
|
||
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
||
# Lê config do projeto (mesmo padrão do fastscnn)
|
||
with open("config.json", "r") as f:
|
||
config = json.load(f)
|
||
|
||
MODELO = config["camera"]
|
||
MODEL_NAME = config["model_name"]
|
||
RESOLUCAO = config["resolucao"] # [W,H] ex: [512,288]
|
||
ROI_INICIO = config["roi_inicio"] # fração
|
||
ROI_TAMANHO = config["roi_tamanho"] # fração
|
||
BACKBONE = config["backbone"]
|
||
|
||
model_to_use = config["model_to_use"]
|
||
dataset_path = os.path.join(MODELO, "dataset")
|
||
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
|
||
|
||
# Backbone default (se não passar por argumento)
|
||
backbone = BACKBONE
|
||
|
||
# Caminho do .pt (se não passar por argumento)
|
||
model_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)
|
||
model_name = ""
|
||
if model_to_use == "geral":
|
||
model_name = f"best_miou.pt"
|
||
elif model_to_use == "main_class":
|
||
model_name = f"best_main.pt"
|
||
else:
|
||
model_name = f"last.pt"
|
||
pt_path = os.path.join(model_path, model_name)
|
||
|
||
if not pt_path:
|
||
raise SystemExit(
|
||
"Faltou apontar o .pt do SegFormer.\n"
|
||
"Use: --pt caminho/do/best_miou.pt\n"
|
||
"ou adicione 'segformer_pt' no seu config.json."
|
||
)
|
||
|
||
# Labelmap
|
||
_, colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
||
ignore_id = infer_ignore_id(ignore_rgb, default_id=255)
|
||
num_classes = len(classes)
|
||
print(f"[cfg] classes={num_classes} | backbone={backbone}")
|
||
print(f"[cfg] pt={pt_path}")
|
||
|
||
# Modelo
|
||
model = load_segformer_from_checkpoint(pt_path, backbone=backbone, num_classes=num_classes, device=device)
|
||
|
||
# Normalização (ImageNet, padrão de muita coisa; se teu treino usou outro, troca aqui)
|
||
from _8_train_segformer import normalize_img
|
||
|
||
if args.camera:
|
||
# === Modo câmera (igual estilo do fastscnn) ===
|
||
pipeline = dai.Pipeline()
|
||
cam_rgb = pipeline.createColorCamera()
|
||
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
|
||
cam_rgb.setBoardSocket(dai.CameraBoardSocket.CAM_A)
|
||
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_bgr = in_rgb.getCvFrame() # BGR
|
||
frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||
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)
|
||
|
||
# tensor
|
||
roi_norm = roi_resized.astype(np.float32) / 255.0
|
||
img_tensor = torch.from_numpy(roi_norm).permute(2, 0, 1).unsqueeze(0).to(device)
|
||
img_tensor = normalize_img(img_tensor)
|
||
img_tensor = img_tensor.float()
|
||
|
||
pred_ids = segformer_predict_ids(model, img_tensor)
|
||
|
||
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 / max(1e-6, (now - prev_time))
|
||
prev_time = now
|
||
|
||
mask_nav = (pred_ids == ClassesSegmentacao.NAVEGAVEL.value)
|
||
status_now, status_final, debug = classificar_status_corredor(mask_nav)
|
||
|
||
cv2.putText(overlay, f"FPS: {fps:.1f} - {status_now.name}", (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 SegFormer (OAK + PyTorch)", cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
|
||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||
break
|
||
|
||
cv2.destroyAllWindows()
|
||
|
||
else:
|
||
# === Modo imagens (agrupado + fallback) ===
|
||
def collect_images_only(folder):
|
||
exts = (".jpg", ".jpeg", ".png")
|
||
paths = [
|
||
os.path.join(folder, f)
|
||
for f in sorted(os.listdir(folder))
|
||
if f.lower().endswith(exts)
|
||
]
|
||
return paths
|
||
|
||
test_root = args.test_folder if args.test_folder else ""
|
||
|
||
image_paths = []
|
||
mask_paths = []
|
||
groups_idx = []
|
||
|
||
if test_root:
|
||
# tenta modo "imagens puras"
|
||
image_paths = collect_images_only(test_root)
|
||
if image_paths:
|
||
mask_paths = None
|
||
groups_idx = None
|
||
print(f"[TEST] Modo inferência pura: {len(image_paths)} imagens")
|
||
else:
|
||
# fallback: dataset estruturado
|
||
image_paths, mask_paths, groups_idx = collect_pairs_grouped(
|
||
test_root, want_groups=args.groups
|
||
)
|
||
else:
|
||
test_root = os.path.join(dataset_path, "split", args.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) > 0, "Nenhuma imagem encontrada."
|
||
if mask_paths is not None:
|
||
assert len(image_paths) == len(mask_paths), "Mismatch imagem/máscara"
|
||
|
||
idx = 0
|
||
window_name = "Original | GroundTruth | Predito (SegFormer)"
|
||
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) # permite redimensionar/maximizar
|
||
while True:
|
||
img_path = image_paths[idx]
|
||
grupo = groups_idx[idx] if groups_idx else "?"
|
||
|
||
img_rgb = np.array(Image.open(img_path).convert("RGB"))
|
||
|
||
if mask_paths is not None:
|
||
mask_path = mask_paths[idx]
|
||
mask_gt = np.array(Image.open(mask_path).convert("L"))
|
||
else:
|
||
mask_gt = None
|
||
|
||
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]
|
||
|
||
if mask_gt is not None:
|
||
mask_roi = mask_gt[y_fim:y_inicio, 0:W]
|
||
else:
|
||
mask_roi = None
|
||
|
||
img_resized = resize_keep_width(img_roi, RESOLUCAO[0], RESOLUCAO[1], cv2.INTER_AREA)
|
||
|
||
if mask_roi is not None:
|
||
mask_resized = resize_keep_width(mask_roi, RESOLUCAO[0], RESOLUCAO[1], cv2.INTER_NEAREST)
|
||
else:
|
||
mask_resized = None
|
||
|
||
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 = normalize_img(img_tensor)
|
||
img_tensor = img_tensor.float()
|
||
|
||
pred_ids = segformer_predict_ids(model, img_tensor)
|
||
|
||
pred_rgb = converter_mask_ids_para_rgb(pred_ids, colormap_rgb, ignore_id)
|
||
|
||
overlay_pred = cv2.addWeighted(img_resized, 0.6, pred_rgb, 0.4, 0.0)
|
||
if mask_gt is not None:
|
||
mask_gt_rgb = converter_mask_ids_para_rgb(mask_resized, colormap_rgb, ignore_id)
|
||
resultado = np.concatenate([img_resized, mask_gt_rgb, overlay_pred], axis=1)
|
||
else:
|
||
resultado = np.concatenate([img_resized, overlay_pred], 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)
|
||
|
||
|
||
mask_nav = (pred_ids == ClassesSegmentacao.NAVEGAVEL.value)
|
||
classificar_status_corredor(mask_nav)
|
||
|
||
|
||
cv2.putText(resultado_completo, f"grupo: {grupo}", (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
|
||
# Converte pra BGR pra exibir
|
||
vis_bgr = cv2.cvtColor(resultado_completo, cv2.COLOR_RGB2BGR)
|
||
# Limites máximos da janela (ajusta se quiser)
|
||
MAX_WIDTH = 1600
|
||
MAX_HEIGHT = 900
|
||
h, w = vis_bgr.shape[:2]
|
||
scale = min(MAX_WIDTH / w, MAX_HEIGHT / h, 1.0)
|
||
if scale < 1.0:
|
||
new_w = int(w * scale)
|
||
new_h = int(h * scale)
|
||
vis_bgr = cv2.resize(vis_bgr, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||
cv2.imshow(window_name, vis_bgr)
|
||
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()
|
||
|
||
|
||
|
||
class ClassesSegmentacao(IntEnum):
|
||
NAONAVEGAVEL = 0
|
||
NAVEGAVEL = 1
|
||
|
||
|
||
class StatusCarroMapa(IntEnum):
|
||
Parado = 0
|
||
EntrandoRua = 1
|
||
CaminhandoRua = 2
|
||
SaindoRua = 3
|
||
Manobrando = 4
|
||
Direcionando = 5
|
||
RetornandoBase = 6 # não usamos aqui, mas mantido
|
||
Indefinido = 7
|
||
|
||
|
||
# ===== HISTER ESE TEMPORAL =====
|
||
|
||
janela_s_padrao = 1.5
|
||
# deixa espaço pra um histórico decente, a janela em segundos faz o recorte real
|
||
_status_hist = deque(maxlen=60)
|
||
|
||
|
||
def _now() -> float:
|
||
"""Tempo monotônico (não sofre com ajustes de relógio do SO)."""
|
||
return time.monotonic()
|
||
|
||
|
||
def _maioria_ultimos(janela_s: float | None = None) -> StatusCarroMapa:
|
||
"""Maioria dos statuses dentro da última janela de tempo (em segundos)."""
|
||
J = float(janela_s) if janela_s is not None else float(janela_s_padrao)
|
||
|
||
t_now = _now()
|
||
|
||
# limpa itens FORA da janela
|
||
while _status_hist and (t_now - _status_hist[0][1] > J):
|
||
_status_hist.popleft()
|
||
|
||
if not _status_hist:
|
||
# fallback razoável
|
||
return StatusCarroMapa.Direcionando
|
||
|
||
# maioria simples
|
||
cont: Dict[StatusCarroMapa, int] = {}
|
||
for st, _t in _status_hist:
|
||
cont[st] = cont.get(st, 0) + 1
|
||
|
||
top_freq = max(cont.values())
|
||
empatados = [st for st, c in cont.items() if c == top_freq]
|
||
|
||
if len(empatados) == 1:
|
||
return empatados[0]
|
||
|
||
# desempata olhando do fim pro início (mais recente primeiro)
|
||
for st, _t in reversed(_status_hist):
|
||
if st in empatados:
|
||
return st
|
||
|
||
|
||
def classificar_status_corredor(mask_nav: np.ndarray):
|
||
"""
|
||
mask_nav: (H,W) com 1 = navegável, 0 = não-navegável
|
||
|
||
Retorna:
|
||
status_now : StatusCarroMapa (instantâneo deste frame)
|
||
status_final : StatusCarroMapa (com histerese)
|
||
probs : dict[StatusCarroMapa, float] (one-hot)
|
||
debug : métricas pra log/diagnóstico
|
||
"""
|
||
if mask_nav is None or mask_nav.size == 0:
|
||
status_now = StatusCarroMapa.Manobrando
|
||
_status_hist.append((status_now, _now()))
|
||
status_final = _maioria_ultimos()
|
||
return status_now, status_final, {
|
||
"nav_near": 0.0,
|
||
"nav_mid": 0.0,
|
||
"nav_far": 0.0,
|
||
"nav_global": 0.0,
|
||
}
|
||
|
||
H, W = mask_nav.shape
|
||
nav = (mask_nav > 0).astype(np.float32)
|
||
|
||
def faixa_mean(y0: int, y1: int) -> float:
|
||
fatia = nav[y0:y1, :]
|
||
if fatia.size == 0:
|
||
return 0.0
|
||
return float(fatia.mean())
|
||
|
||
# 3 faixas verticais: far (topo), mid (meio), near (embaixo)
|
||
y_far_top = 0
|
||
y_far_bot = int(0.25 * H)
|
||
y_mid_top = y_far_bot
|
||
y_mid_bot = int(0.5 * H)
|
||
y_near_top = y_mid_bot
|
||
y_near_bot = H
|
||
|
||
nav_far = faixa_mean(y_far_top, y_far_bot)
|
||
nav_mid = faixa_mean(y_mid_top, y_mid_bot)
|
||
nav_near = faixa_mean(y_near_top, y_near_bot)
|
||
nav_global = float(nav.mean())
|
||
|
||
# diferenças entre faixas (pra medir quão "desbalanceado" está)
|
||
d_nm = abs(nav_near - nav_mid)
|
||
d_mf = abs(nav_mid - nav_far)
|
||
d_nf = abs(nav_near - nav_far)
|
||
max_delta = max(d_nm, d_mf, d_nf)
|
||
|
||
# ---- Blobs 2D na região distante (parede esquerda x direita) ----
|
||
y_blob_top = 0
|
||
y_blob_bot = int(0.35 * H) # um pouco mais profundo que o nav_far
|
||
|
||
faixa_obs = (mask_nav[y_blob_top:y_blob_bot, :] == 0).astype(np.uint8) # 1 = obstáculo
|
||
|
||
H_blob, W_blob = faixa_obs.shape
|
||
num_blobs_far = 0
|
||
corridor_nav_far = 0.0
|
||
corridor_width_frac = 0.0
|
||
|
||
if H_blob > 0 and W_blob > 0 and faixa_obs.max() > 0:
|
||
# connectedComponents espera 0/255
|
||
faixa_obs_bin = (faixa_obs * 255).astype(np.uint8)
|
||
|
||
num_labels, labels = cv2.connectedComponents(faixa_obs_bin)
|
||
|
||
# ignora blobs muito pequenos (ruído)
|
||
MIN_AREA = 0.005 * H_blob * W_blob # 0.5% da área da faixa
|
||
blobs = []
|
||
|
||
for label in range(1, num_labels): # 0 é o fundo
|
||
ys, xs = np.where(labels == label)
|
||
area = len(xs)
|
||
if area < MIN_AREA:
|
||
continue
|
||
|
||
x_min, x_max = xs.min(), xs.max()
|
||
y_min, y_max = ys.min(), ys.max()
|
||
blobs.append({
|
||
"area": area,
|
||
"bbox": (x_min, y_min, x_max, y_max),
|
||
"x_center": float(xs.mean()),
|
||
})
|
||
|
||
num_blobs_far = len(blobs)
|
||
|
||
if num_blobs_far >= 2:
|
||
# ordena da parede mais à esquerda pra mais à direita
|
||
blobs_sorted = sorted(blobs, key=lambda b: b["x_center"])
|
||
left_blob = blobs_sorted[0]
|
||
right_blob = blobs_sorted[-1]
|
||
|
||
# corredor é a região entre x_max da esquerda e x_min da direita
|
||
x_left = left_blob["bbox"][2] + 1 # x_max (inclusivo) -> +1 pra slice
|
||
x_right = right_blob["bbox"][0] # x_min
|
||
|
||
if x_right > x_left:
|
||
corridor_width = x_right - x_left
|
||
corridor_width_frac = corridor_width / float(W)
|
||
|
||
faixa_nav_corridor = (mask_nav[y_blob_top:y_blob_bot, x_left:x_right] > 0).astype(np.float32)
|
||
if faixa_nav_corridor.size > 0:
|
||
corridor_nav_far = float(faixa_nav_corridor.mean())
|
||
|
||
# ===== Limiares (podemos tunar depois) =====
|
||
THR_BLOCKED = 0.30 # global bem baixo -> quase sem caminho
|
||
THR_OPEN = 0.90 # global bem alto -> mundo aberto
|
||
THR_DELTA_COR = 0.15 # diferença "significativa" entre faixas
|
||
THR_COR_FAR_LARGO = 0.65 # far ainda relativamente alto em corredor largo
|
||
THR_OPEN_MUITO_LIMPO = 0.90 # quase tudo navegável
|
||
THR_DELTA_OBST_PEQ = 0.20 # desbalance vertical máximo para considerar "campo aberto com obstáculo pequeno"
|
||
MIN_CORRIDOR_WIDTH_FRAC = 0.10 # corredor tem que ter pelo menos ~10% da largura
|
||
MAX_CORRIDOR_WIDTH_FRAC = 0.70 # evita chamar de corredor quando é um "campo" gigante
|
||
|
||
# 1) Parado: quase não há caminho à frente (meio+frente mortos)
|
||
cond_parado = (
|
||
nav_global <= THR_BLOCKED
|
||
and nav_mid < 0.20
|
||
and nav_far < 0.10
|
||
)
|
||
|
||
# 2) Direcionando: mundo aberto, sem corredor marcado
|
||
# 2.1 base: tudo alto e muito homogêneo
|
||
cond_direcionando_base = (
|
||
nav_global >= THR_OPEN
|
||
and max_delta < 0.10
|
||
)
|
||
|
||
# 2.2 modo "campo aberto com obstáculo pequeno":
|
||
# quase tudo navegável, e até o far é bem alto
|
||
cond_direcionando_obst_peq = (
|
||
nav_global >= THR_OPEN_MUITO_LIMPO and # >= 0.90
|
||
nav_near >= 0.95 and
|
||
nav_mid >= 0.80 and # um pouco mais permissivo
|
||
nav_far >= 0.50 and # aceita far um pouco mais fechado
|
||
num_blobs_far <= 1 # no máximo UMA parede grande
|
||
)
|
||
|
||
# 2.3 modo "campo aberto com parede na frente":
|
||
# chão bem navegável perto, sem corredor definido, e FAR quase todo bloqueado
|
||
cond_direcionando_frente_fe_chada = (
|
||
nav_near >= 0.80 and # perto bem aberto
|
||
nav_mid >= 0.30 and # meio ainda razoável
|
||
nav_far <= 0.10 and # topo praticamente bloqueado (parede)
|
||
nav_global >= 0.50 and # ainda tem bastante área navegável no frame
|
||
num_blobs_far <= 1 # no máximo uma "parede", nada de corredor
|
||
)
|
||
|
||
# 2.3 modo "campo aberto com borda lateral":
|
||
# cena razoavelmente aberta, uma parede forte de um lado, mas sem corredor fechado
|
||
cond_direcionando_borda_lateral = (
|
||
nav_global >= 0.60 and # já tem boa área navegável
|
||
nav_near >= 0.70 and
|
||
nav_mid >= 0.50 and
|
||
nav_far >= 0.40 and # far não está "morrendo", só mais sujo
|
||
nav_far <= 0.80 and # não é mundão 100% limpo
|
||
num_blobs_far == 1 # exatamente UMA parede grande
|
||
)
|
||
|
||
cond_direcionando_aberto = (
|
||
nav_global >= 0.75 and
|
||
nav_near >= 0.70 and
|
||
nav_mid >= 0.70 and
|
||
nav_far >= 0.70 and
|
||
max_delta <= 0.12 and
|
||
num_blobs_far <= 1
|
||
)
|
||
|
||
cond_direcionando = (
|
||
cond_direcionando_base
|
||
or cond_direcionando_obst_peq
|
||
or cond_direcionando_frente_fe_chada
|
||
or cond_direcionando_borda_lateral
|
||
or cond_direcionando_aberto
|
||
)
|
||
|
||
# 3) CaminhandoRua: dentro do corredor "clássico"
|
||
cond_caminhando_base = (
|
||
nav_near >= 0.55 and
|
||
nav_mid >= 0.25 and
|
||
nav_far <= 0.50 and
|
||
(nav_near - nav_far) >= 0.20 and
|
||
num_blobs_far >= 2 # precisa de DUAS paredes
|
||
)
|
||
|
||
# 3.1 CaminhandoRua em corredor mais largo, com parede só de um lado
|
||
# enquadra bem os casos:
|
||
# nav_global ~0.66–0.77, near ~0.75–0.88, mid ~0.57–0.70, far ~0.56–0.63
|
||
cond_caminhando_largo = (
|
||
nav_global >= 0.60 and
|
||
nav_near >= 0.75 and
|
||
nav_mid >= 0.50 and
|
||
nav_far >= 0.50 and
|
||
nav_far <= THR_COR_FAR_LARGO and
|
||
(nav_near - nav_far) >= THR_DELTA_COR and
|
||
num_blobs_far >= 2 # corredor largo, mas ainda corredor
|
||
)
|
||
|
||
cond_caminhando_multi_corredores = (
|
||
nav_global >= 0.50 and # tem chão suficiente
|
||
nav_mid >= 0.55 and # meio bem limpo
|
||
nav_near >= 0.50 and # perto também ok
|
||
num_blobs_far >= 2 and # pelo menos duas "paredes"
|
||
corridor_width_frac >= MIN_CORRIDOR_WIDTH_FRAC and
|
||
corridor_width_frac <= MAX_CORRIDOR_WIDTH_FRAC and
|
||
corridor_nav_far >= 0.55 # corredor entre paredes bem navegável
|
||
)
|
||
|
||
cond_caminhando = (
|
||
cond_caminhando_base
|
||
or cond_caminhando_largo
|
||
or cond_caminhando_multi_corredores
|
||
)
|
||
|
||
cond_entrando = (
|
||
nav_global >= 0.75 and
|
||
nav_near >= 0.90 and
|
||
nav_mid >= 0.60 and
|
||
nav_far >= 0.40 and
|
||
nav_far <= 0.85 and
|
||
(nav_near - nav_far) >= 0.10 and # far mais fechado que near
|
||
num_blobs_far >= 2 and # duas paredes detectadas
|
||
corridor_width_frac >= MIN_CORRIDOR_WIDTH_FRAC and
|
||
corridor_nav_far >= 0.60 # corredor entre as paredes bem navegável
|
||
)
|
||
|
||
# 5) SaindoRua
|
||
cond_saindo_1 = (
|
||
nav_near >= 0.50 and
|
||
nav_mid >= 0.25 and
|
||
nav_far >= 0.55 and
|
||
(nav_far - nav_mid) >= 0.10 and
|
||
nav_far >= nav_near - 0.15
|
||
)
|
||
|
||
cond_saindo_2 = (
|
||
nav_global >= 0.75 and
|
||
nav_near >= 0.70 and
|
||
nav_mid >= 0.60 and
|
||
nav_far >= 0.75 and
|
||
nav_far >= nav_mid
|
||
)
|
||
|
||
cond_saindo = cond_saindo_1 or cond_saindo_2
|
||
|
||
# ===== Decisão (ordem importa!) =====
|
||
if cond_parado:
|
||
status = StatusCarroMapa.Parado
|
||
elif cond_direcionando:
|
||
status = StatusCarroMapa.Direcionando
|
||
elif cond_saindo:
|
||
status = StatusCarroMapa.SaindoRua
|
||
elif cond_entrando:
|
||
status = StatusCarroMapa.EntrandoRua
|
||
elif cond_caminhando:
|
||
status = StatusCarroMapa.CaminhandoRua
|
||
else:
|
||
status = StatusCarroMapa.Indefinido
|
||
|
||
debug = {
|
||
"nav_global": nav_global,
|
||
"nav_near": nav_near,
|
||
"nav_mid": nav_mid,
|
||
"nav_far": nav_far,
|
||
"d_nm": d_nm,
|
||
"d_mf": d_mf,
|
||
"d_nf": d_nf,
|
||
"max_delta": max_delta,
|
||
"num_blobs_far": num_blobs_far,
|
||
"corridor_nav_far": corridor_nav_far,
|
||
"corridor_width_frac": corridor_width_frac,
|
||
"cond_parado": cond_parado,
|
||
"cond_direcionando_base": cond_direcionando_base,
|
||
"cond_direcionando_obst_peq": cond_direcionando_obst_peq,
|
||
"cond_caminhando_base": cond_caminhando_base,
|
||
"cond_caminhando_largo": cond_caminhando_largo,
|
||
"cond_entrando": cond_entrando,
|
||
"cond_saindo": cond_saindo,
|
||
"THR_BLOCKED": THR_BLOCKED,
|
||
"THR_OPEN": THR_OPEN,
|
||
"THR_DELTA_COR": THR_DELTA_COR,
|
||
"THR_COR_FAR_LARGO": THR_COR_FAR_LARGO,
|
||
"THR_OPEN_MUITO_LIMPO": THR_OPEN_MUITO_LIMPO,
|
||
}
|
||
|
||
status_now = status
|
||
print(status_now.name, debug)
|
||
|
||
_status_hist.append((status_now, _now()))
|
||
status_final = _maioria_ultimos()
|
||
|
||
return status_now, status_final, debug
|
||
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|