diff --git a/AgroBase/AgroBase/Models/TrajetoriaMapaOperacaoModel.cs b/AgroBase/AgroBase/Models/TrajetoriaMapaOperacaoModel.cs index fc9e8f371..f9a90ef4a 100644 --- a/AgroBase/AgroBase/Models/TrajetoriaMapaOperacaoModel.cs +++ b/AgroBase/AgroBase/Models/TrajetoriaMapaOperacaoModel.cs @@ -2496,40 +2496,64 @@ namespace AgroBase.Models public enum AutonomiaCorredorStatus { Desconhecido = 0, - OkCompleto, // Bateria + herbicida suficientes pra entrar e pulverizar - OkSomenteTransito, // Bateria ok, herbicida insuficiente pra pulverizar tudo - OkSupervisionado, // Liberado manualmente - Critico // Nem a bateria garante atravessar com segurança + OkCompleto, // Bateria + herbicida suficientes pra entrar e pulverizar + OkSomenteTransito, // Bateria ok, herbicida insuficiente pra pulverizar tudo (mas pode atravessar) + OkSupervisionado, // Liberado manualmente MESMO com parâmetro não OK (override humano) + AguardandoLiberacaoManual, // Travado após ficar crítico por X segundos, só sai com liberação humana + Critico // Nem a bateria garante atravessar com segurança (ou herbicida mandatório insuficiente) } public class AutonomiaCorredorModel { + // ========================== + // Configurações + // ========================== + private readonly TimeSpan _tempoParaTravar = TimeSpan.FromSeconds(10); // X segundos em crítico -> trava + + // ========================== + // Estado público + // ========================== [JsonProperty] public bool Iniciado { get; private set; } + [JsonProperty] - public AutonomiaCorredorStatus Status { get; private set; } - private bool _liberado => - Status == AutonomiaCorredorStatus.OkCompleto || - Status == AutonomiaCorredorStatus.OkSomenteTransito || - Status == AutonomiaCorredorStatus.OkSupervisionado; + public AutonomiaCorredorStatus Status { get; private set; } = AutonomiaCorredorStatus.Desconhecido; + [JsonProperty] - public bool Liberado { get; private set; } + public bool Liberado { get; private set; } = true; + [JsonProperty] public double DistanciaCorredor_m { get; private set; } + [JsonProperty] public double DistanciaSeguraBateria_m { get; private set; } + [JsonProperty] public double DistanciaSeguraHerbicida_m { get; private set; } - private Dictionary CorredoresLiberados { get; set; } = new Dictionary(); - // Flags de conveniência - public bool BateriaSuficiente => DistanciaSeguraBateria_m >= DistanciaCorredor_m; - public bool HerbicidaSuficiente => DistanciaSeguraHerbicida_m >= DistanciaCorredor_m; - - // Texto pra UI/log [JsonProperty] public string Motivo { get; private set; } = ""; + // ========================== + // Flags de conveniência + // ========================== + public bool BateriaSuficiente => DistanciaSeguraBateria_m >= DistanciaCorredor_m; + public bool HerbicidaSuficiente => DistanciaSeguraHerbicida_m >= DistanciaCorredor_m; + + // ========================== + // Controle de histórico por corredor + // ========================== + private Dictionary CorredoresLiberados { get; set; } = new Dictionary(); + + // ========================== + // Trava (latch) por corredor + // ========================== + private readonly HashSet _corredoresTravados = new HashSet(); + private readonly Dictionary _inicioCriticoUtc = new Dictionary(); + + // ========================== + // Ciclo de vida + // ========================== public void Iniciar() { Iniciado = true; @@ -2538,10 +2562,18 @@ namespace AgroBase.Models public void Parar() { Iniciado = false; + // opcional: não limpar travas aqui, pois pode querer manter estado até o operador decidir. + // se preferir limpar, descomente: + // _corredoresTravados.Clear(); + // _inicioCriticoUtc.Clear(); } + // ========================== + // Atualização principal + // ========================== public void AtualizarDados(int idxCorredor, bool bateria_liberada, bool reservatorio_liberado) { + // Se não iniciado, não bloqueia nada if (!Iniciado) { Liberado = true; @@ -2550,12 +2582,18 @@ namespace AgroBase.Models return; } + // Valida índice do corredor if (Variaveis.OperacaoEmAndamento.Trajetoria?._Corredores?.Count <= idxCorredor) return; + // Se já registrou esse corredor como liberado e travado/feito, não recalcula if (CorredoresLiberados.ContainsKey(idxCorredor) && CorredoresLiberados[idxCorredor].Liberado) return; + // Atalho: interpretar que o operador “clicou em liberar” (pelo menos um sinal veio) + bool solicitouLiberacaoHumana = bateria_liberada || reservatorio_liberado; + + // Carrega distâncias e contadores double distanciaCorredor = Variaveis.OperacaoEmAndamento.Trajetoria._Corredores[idxCorredor]?.DistanciaTotal ?? 0; var _Bateria = Variaveis.OperacaoEmAndamento.DispSen?.Dados?.ContadorCarga; @@ -2567,41 +2605,157 @@ namespace AgroBase.Models bool bateriaOk = BateriaSuficiente && (_Bateria?.Iniciado ?? false); bool herbOk = HerbicidaSuficiente && (_Reservatorio?.Iniciado ?? false); - bool exigirPulverizacaoTotal = Variaveis.OperacaoEmAndamento.Parametros?.ModulosMandatorios?.Any(x => x.Dispositivo == T_Code.Atu && x.Mandatorio && x.Utilizar) ?? false; - List status_mods = new List(); - string _motivo = ""; - string _motivo_bat = "Bateria " + (bateriaOk ? "suficiente" : "insuficiente") + $". Distância segura: {DistanciaSeguraBateria_m:F1} m, corredor: {distanciaCorredor:F1} m"; - string _motivo_herb = "Herbicida " + (herbOk ? "suficiente" : "insuficiente") + $" para pulverizar o corredor. Distância segura pulverizando: {DistanciaSeguraHerbicida_m:F1} m, corredor: {distanciaCorredor:F1} m"; + bool exigirPulverizacaoTotal = + Variaveis.OperacaoEmAndamento.Parametros?.ModulosMandatorios? + .Any(x => x.Dispositivo == T_Code.Atu && x.Mandatorio && x.Utilizar) ?? false; - if (!bateriaOk) + // ========================== + // 1) Calcula status "base" (sem considerar trava) + // ========================== + AutonomiaCorredorStatus statusBase; + string motivoBase = ""; + string motivoBat = "Bateria " + (bateriaOk ? "suficiente" : "insuficiente") + + $". Distância segura: {DistanciaSeguraBateria_m:F1} m, corredor: {distanciaCorredor:F1} m"; + string motivoHerb = "Herbicida " + (herbOk ? "suficiente" : "insuficiente") + + $" para pulverizar o corredor. Distância segura pulverizando: {DistanciaSeguraHerbicida_m:F1} m, corredor: {distanciaCorredor:F1} m"; + + // Regra: bateria insuficiente é sempre crítica (não atravessa com segurança) + bool criticoBateria = !bateriaOk; + + // Regra: herbicida insuficiente só vira crítico se pulverização total é mandatória + bool criticoHerbicida = !herbOk && exigirPulverizacaoTotal; + + if (criticoBateria || criticoHerbicida) { - status_mods.Add(!bateria_liberada ? AutonomiaCorredorStatus.Critico : AutonomiaCorredorStatus.OkSupervisionado); - if (bateria_liberada) _motivo_bat = $"Liberação humana: {_motivo_bat}"; - _motivo += _motivo_bat; + statusBase = AutonomiaCorredorStatus.Critico; + + if (criticoBateria) + motivoBase += motivoBat; + + if (!herbOk) + motivoBase += (string.IsNullOrEmpty(motivoBase) ? "" : "; ") + motivoHerb; } - - if (!herbOk) + else { - if (exigirPulverizacaoTotal) - status_mods.Add(!reservatorio_liberado ? AutonomiaCorredorStatus.Critico : AutonomiaCorredorStatus.OkSupervisionado); + // Não crítico + if (!herbOk && !exigirPulverizacaoTotal) + { + // Pode atravessar, mas não pulveriza tudo + statusBase = AutonomiaCorredorStatus.OkSomenteTransito; + motivoBase = motivoHerb; + } else - status_mods.Add(AutonomiaCorredorStatus.OkSomenteTransito); - if (reservatorio_liberado) _motivo_herb = $"Liberação humana: {_motivo_herb}"; - _motivo += (string.IsNullOrEmpty(_motivo) ? "" : "; ") + _motivo_herb; + { + // Tudo ok + statusBase = AutonomiaCorredorStatus.OkCompleto; + } } - Status = - status_mods.Contains(AutonomiaCorredorStatus.Critico) ? AutonomiaCorredorStatus.Critico : - status_mods.Contains(AutonomiaCorredorStatus.OkSupervisionado) ? AutonomiaCorredorStatus.OkSupervisionado : - status_mods.Contains(AutonomiaCorredorStatus.OkSomenteTransito) ? AutonomiaCorredorStatus.OkSomenteTransito : - AutonomiaCorredorStatus.OkCompleto; + // ========================== + // 2) Atualiza timers de crítico e ativa trava se necessário + // ========================== + bool estaCriticoAgora = (statusBase == AutonomiaCorredorStatus.Critico); + + if (estaCriticoAgora) + { + if (!_inicioCriticoUtc.ContainsKey(idxCorredor)) + _inicioCriticoUtc[idxCorredor] = DateTime.UtcNow; + + var tempoCritico = DateTime.UtcNow - _inicioCriticoUtc[idxCorredor]; + + if (tempoCritico >= _tempoParaTravar) + _corredoresTravados.Add(idxCorredor); + } + else + { + // Saiu do crítico -> reseta timer (mas não destrava automaticamente se já travou) + if (_inicioCriticoUtc.ContainsKey(idxCorredor)) + _inicioCriticoUtc.Remove(idxCorredor); + } + + // ========================== + // 3) Aplicação da trava + liberação humana + // ========================== + bool corredorTravado = _corredoresTravados.Contains(idxCorredor); + + AutonomiaCorredorStatus statusFinal = statusBase; + bool liberadoFinal = false; + string motivoFinal = ""; + + if (corredorTravado) + { + // Se travou, não libera sozinho nunca + if (!solicitouLiberacaoHumana) + { + statusFinal = AutonomiaCorredorStatus.AguardandoLiberacaoManual; + liberadoFinal = false; + + // Motivo ajuda o operador a entender o porquê + string detalhes = string.IsNullOrEmpty(motivoBase) + ? "Autonomia insuficiente detectada anteriormente." + : motivoBase; + + motivoFinal = $"Corredor {idxCorredor + 1}: BLOQUEADO. {detalhes} " + + $"Aguardando liberação do operador pela base."; + } + else + { + // Operador solicitou liberação + if (statusBase != AutonomiaCorredorStatus.Critico) + { + // ✅ Se agora está OK, libera como OK normal (não supervisionado) + statusFinal = statusBase; // OkCompleto ou OkSomenteTransito + liberadoFinal = true; + + motivoFinal = $"Corredor {idxCorredor + 1}: LIBERADO pelo operador (parâmetros OK). " + + (string.IsNullOrEmpty(motivoBase) ? "Liberado" : motivoBase); + + // Destrava o corredor + _corredoresTravados.Remove(idxCorredor); + _inicioCriticoUtc.Remove(idxCorredor); + } + else + { + // ❌ Ainda está crítico, mas operador insistiu -> OkSupervisionado (override) + statusFinal = AutonomiaCorredorStatus.OkSupervisionado; + liberadoFinal = true; + + motivoFinal = $"Corredor {idxCorredor + 1}: LIBERAÇÃO MANUAL (override). " + + (string.IsNullOrEmpty(motivoBase) ? "Autonomia insuficiente, mas liberado." : motivoBase); + + // Destrava o corredor (o operador assumiu) + _corredoresTravados.Remove(idxCorredor); + _inicioCriticoUtc.Remove(idxCorredor); + } + } + } + else + { + // Sem trava: comportamento raiz + statusFinal = statusBase; + + liberadoFinal = + statusFinal == AutonomiaCorredorStatus.OkCompleto || + statusFinal == AutonomiaCorredorStatus.OkSomenteTransito; + + motivoFinal = $"Corredor {idxCorredor + 1}: " + + (liberadoFinal + ? (string.IsNullOrEmpty(motivoBase) ? "Liberado" : motivoBase) + : (string.IsNullOrEmpty(motivoBase) ? "Bloqueado" : motivoBase)); + } + + // ========================== + // 4) Persiste/Loga resultado como você já fazia (mantendo raiz) + // ========================== + Status = statusFinal; var res = new AutonomiaCorredorModel() { - Liberado = _liberado, - Motivo = $"Corredor {idxCorredor + 1}: " + (string.IsNullOrEmpty(_motivo) ? "Liberado" : $"{_motivo}"), - Status = Status, + Iniciado = this.Iniciado, + Liberado = liberadoFinal, + Motivo = motivoFinal, + Status = statusFinal, DistanciaCorredor_m = DistanciaCorredor_m, DistanciaSeguraBateria_m = DistanciaSeguraBateria_m, DistanciaSeguraHerbicida_m = DistanciaSeguraHerbicida_m, @@ -2611,11 +2765,17 @@ namespace AgroBase.Models if (CorredoresLiberados.ContainsKey(idxCorredor)) { - if (CorredoresLiberados[idxCorredor].Liberado == false && _liberado) + // Se estava bloqueado e agora liberou, registra mudança + if (CorredoresLiberados[idxCorredor].Liberado == false && liberadoFinal) { CorredoresLiberados[idxCorredor] = res; addLog = true; } + else + { + // Mantém o mais recente motivo/status para UI (opcional) + CorredoresLiberados[idxCorredor] = res; + } } else { @@ -2625,18 +2785,28 @@ namespace AgroBase.Models if (addLog) { - Variaveis.OperacaoEmAndamento.Sensoriamento.InserirLog(T_Code.Trj, _liberado ? StatusModulo.Operante : StatusModulo.Alerta, _liberado ? 100 : 0, res.Motivo); + // Mantive tua ideia: loga operante quando liberado, alerta quando não + Variaveis.OperacaoEmAndamento.Sensoriamento.InserirLog( + T_Code.Trj, + liberadoFinal ? StatusModulo.Operante : StatusModulo.Alerta, + liberadoFinal ? 100 : 0, + res.Motivo + ); } - Liberado = CorredoresLiberados[idxCorredor].Liberado; - Motivo = CorredoresLiberados[idxCorredor].Motivo; + // Atualiza estado exposto + Liberado = res.Liberado; + Motivo = res.Motivo; } - + // ========================== + // Clone (mantido) + // ========================== public AutonomiaCorredorModel Clone() { return new AutonomiaCorredorModel() { + Iniciado = Iniciado, Liberado = Liberado, Motivo = Motivo, Status = Status, diff --git a/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/health_worker/main.py b/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/health_worker/main.py index ef65624f3..5f820c63f 100644 --- a/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/health_worker/main.py +++ b/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/health_worker/main.py @@ -49,7 +49,7 @@ def main(): try: # Atualizar SAÚDE DOS MÓDULOS a cada 2.5s - if (agora - ultimo_saude_modulos) >= 2.5: + if (agora - ultimo_saude_modulos) >= 0.5: try: for t_code, modulo in modulos.items(): modulo.atualizar_saude() diff --git a/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/health_worker/modulos/ponte_ip.py b/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/health_worker/modulos/ponte_ip.py index 4e3b7380e..4284fb13d 100644 --- a/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/health_worker/modulos/ponte_ip.py +++ b/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/health_worker/modulos/ponte_ip.py @@ -1,9 +1,11 @@ +import re import time import threading import statistics import socket import psutil import subprocess +from concurrent.futures import ThreadPoolExecutor from collections import deque import paho.mqtt.client as mqtt from shared.contexto_global_redis import ContextoGlobalRedis @@ -19,7 +21,7 @@ class ModuloIPBribge(ModuloDiagnosticoBase): self._thread_saude = None self._running = False - self.nic_name = None + self.nic_name = self.descobrir_nic_para_base() self.window_size = window_size self.rtts = deque(maxlen=window_size) # ping válidos (ms) @@ -37,9 +39,88 @@ class ModuloIPBribge(ModuloDiagnosticoBase): self.last_heartbeat_ts = time.time() # atualize isso de fora self.rover_id = None - self._mqtt_conectado = False self.sub = False - self._start_mqtt_heartbeat() + + self._mqtt = None + self._mqtt_conectado = False + self._mqtt_executor = ThreadPoolExecutor(max_workers=1) + self._mqtt_future = None + self._mqtt_lock = threading.Lock() + self._mqtt_last_attempt_ts = 0.0 + self._mqtt_backoff_s = 1.0 # começa rápido + self._mqtt_backoff_max_s = 15.0 # limite + self._mqtt_connecting = False + + self._ping_executor = ThreadPoolExecutor(max_workers=1) + self._ping_future = None + self._ping_lock = threading.Lock() + + # snapshot do último ping concluído + self._last_ping_ts = 0.0 + self._last_ping_ok = False + self._last_ping_rtt_ms = None + self._last_ping_loss_pct = 100.0 + + # parâmetros do ping + self._ping_timeout_ms = 400 # recomendado: 300–500 + self._ping_count = 2 # recomendado: 1 (janela já suaviza) + self._ping_min_interval = 0.6 # não precisa pingar a cada 100ms + + + + def _start_ping_async_if_needed(self): + """Dispara ping em background se não houver um em andamento.""" + now = time.time() + + # respeita intervalo mínimo (evita spam de ping) + if (now - self._last_ping_ts) < self._ping_min_interval: + return + + # já tem um ping rodando? + if self._ping_future is not None and not self._ping_future.done(): + return + + base_ip = self.get_base_ip() + if not base_ip: + # sem rota ainda -> considera "sem dados" (não bloqueia) + with self._ping_lock: + self._last_ping_ts = now + self._last_ping_ok = False + self._last_ping_rtt_ms = None + self._last_ping_loss_pct = 100.0 + return + + # dispara o ping em background + self._ping_future = self._ping_executor.submit( + self._ping_once, base_ip, self._ping_count, self._ping_timeout_ms + ) + + def _consume_ping_result_if_ready(self): + """Se o ping terminou, atualiza o snapshot sem bloquear o loop.""" + if self._ping_future is None or not self._ping_future.done(): + return + + try: + # _ping_once deve retornar: (ok: bool, avg_rtt_ms: Optional[float], loss_pct: float) + ok, avg_rtt_ms, loss_pct = self._ping_future.result(timeout=0) + except Exception: + ok, avg_rtt_ms, loss_pct = False, None, 100.0 + + with self._ping_lock: + self._last_ping_ts = time.time() + self._last_ping_ok = ok + self._last_ping_rtt_ms = avg_rtt_ms + self._last_ping_loss_pct = loss_pct + + # limpa + self._ping_future = None + + def _get_ping_snapshot(self): + """Lê o último resultado do ping (snapshot thread-safe).""" + with self._ping_lock: + return (self._last_ping_ok, self._last_ping_rtt_ms, self._last_ping_loss_pct, self._last_ping_ts) + + def atualizar_saude(self): @@ -72,38 +153,107 @@ class ModuloIPBribge(ModuloDiagnosticoBase): self._mqtt.connect(base_ip, 1883, 60) self._mqtt.loop_start() except: - pass - #print("Erro ao se conectar no broker mqtt") + #pass + print("Erro ao se conectar no broker mqtt") + + def _start_mqtt_heartbeat_async(self): + """Dispara tentativa de conexão MQTT sem bloquear o loop.""" + now = time.time() + + with self._mqtt_lock: + if self._mqtt_conectado: + return + if self._mqtt_connecting: + return + + # respeita backoff + if (now - self._mqtt_last_attempt_ts) < self._mqtt_backoff_s: + return + + self._mqtt_last_attempt_ts = now + self._mqtt_connecting = True + + base_ip = self.get_base_ip() + if base_ip is None: + with self._mqtt_lock: + self._mqtt_connecting = False + return + + # se já tem uma future rodando, não dispara outra + if self._mqtt_future is not None and not self._mqtt_future.done(): + return + + self._mqtt_future = self._mqtt_executor.submit(self._mqtt_connect_worker, base_ip) + + def _mqtt_connect_worker(self, base_ip: str): + """Roda em background. Pode bloquear aqui sem travar o health.""" + try: + client = mqtt.Client() + client.on_connect = self._on_connect + client.on_message = self._on_message + client.on_disconnect = self._on_disconnect + + # dica V1: keepalive menor ajuda a detectar quedas mais rápido + keepalive = 15 + + # conecta (bloqueante) + client.connect(base_ip, 1883, keepalive) + + # se chegou aqui, iniciou. loop_start é não-bloqueante + client.loop_start() + + # troca o client com lock (evita race) + with self._mqtt_lock: + # se tinha um client antigo, tenta parar + if self._mqtt is not None: + try: + self._mqtt.loop_stop() + self._mqtt.disconnect() + except: + pass + + self._mqtt = client + # on_connect vai setar _mqtt_conectado quando confirmar + + except Exception: + # falhou: libera pra tentar de novo com backoff maior + with self._mqtt_lock: + self._mqtt_conectado = False + self._mqtt_connecting = False + self._mqtt_backoff_s = min(self._mqtt_backoff_s * 2.0, self._mqtt_backoff_max_s) + def _on_connect(self, client, userdata, flags, rc): - if rc == 0: - print("[HEARTBEAT] MQTT conectado") - self._mqtt_conectado = True - - self._reset_janelas() - - self.rover_id = ContextoGlobalRedis.get_equipamento().get("serial_number") - - if self.sub == False: - topic = f"agrobot/v1/rover/{self.rover_id}/heartbeat" - client.subscribe(topic) - print("[HEARTBEAT] Subscribado em:", topic) - self.sub = True - else: - print("Erro ao conectar MQTT:", rc) + with self._mqtt_lock: + self._mqtt_conectado = (rc == 0) + self._mqtt_connecting = False + # reset backoff quando conecta + if self._mqtt_conectado: + self._mqtt_backoff_s = 1.0 + self._reset_janelas() + self.rover_id = ContextoGlobalRedis.get_equipamento().get("serial_number") + if self.sub == False: + topic = f"agrobot/v1/rover/{self.rover_id}/heartbeat" + client.subscribe(topic) + print("[HEARTBEAT] Subscribado em:", topic) + self.sub = True + else: + print("Erro ao conectar MQTT:", rc) def _on_disconnect(self, client, userdata, rc): - print("[HEARTBEAT] MQTT desconectado!", rc) - self._mqtt_conectado = False - self.sub = False + with self._mqtt_lock: + self._mqtt_conectado = False + self._mqtt_connecting = False + self.sub = False + print("[HEARTBEAT] MQTT desconectado!", rc) - # se rc != 0 significa desconexão inesperada - if rc != 0: - print("[HEARTBEAT] Desconexão inesperada — pode ter perdido o link 900 MHz") - - # Marca como "sem heartbeat" instantâneo (opcional) - # Aqui podemos forçar atraso grande - self.last_heartbeat_ts = 0 + # se rc != 0 significa desconexão inesperada + if rc != 0: + print("[HEARTBEAT] Desconexão inesperada — pode ter perdido o link 900 MHz") + + # Marca como "sem heartbeat" instantâneo (opcional) + # Aqui podemos forçar atraso grande + self.last_heartbeat_ts = 0 def _on_message(self, client, userdata, msg): topic = msg.topic @@ -185,7 +335,7 @@ class ModuloIPBribge(ModuloDiagnosticoBase): self.last_bw_counters = counters self.last_bw_ts = now - def _ping_once(self): + def _ping_once_old(self, pings=5, timeout=500): """ Retorna RTT em ms se ok, ou None se timeout. Aqui exemplo pra Windows usando 'ping -n 1'. @@ -198,14 +348,14 @@ class ModuloIPBribge(ModuloDiagnosticoBase): # timeout de 1000 ms proc = subprocess.run( - ["ping", "-n", "5", "-w", "500", base_ip], + ["ping", "-n", f"{pings}", "-w", f"{timeout}", base_ip], capture_output=True, text=True ) if proc.returncode != 0: - #print("timeout") + print("timeout") return None, 100 - #print(proc.stdout) + print(proc.stdout) perda = None media_tempo = None @@ -224,6 +374,51 @@ class ModuloIPBribge(ModuloDiagnosticoBase): print("erro timeout") return None, 100 + def _ping_once(self, base_ip: str, pings: int, timeout_ms: int): + """ + Retorna (ok: bool, avg_rtt_ms: float|None, loss_pct: float) + Compatível com o ping async. + """ + try: + cmd = ["ping", "-n", str(pings), "-w", str(timeout_ms), base_ip] + proc = subprocess.run(cmd, capture_output=True, text=True) + out = (proc.stdout or "") + "\n" + (proc.stderr or "") + + # Reply confiável: TTL= + replies = len(re.findall(r"TTL=", out, flags=re.IGNORECASE)) + ok = replies > 0 + + # Tempo: time=12ms / time<1ms / tempo=12ms + times = [] + for m in re.findall(r"(?:time|tempo)[=<]\s*(\d+)\s*ms", out, flags=re.IGNORECASE): + try: + times.append(float(m)) + except: + pass + avg_rtt_ms = (sum(times) / len(times)) if times else None + + # Loss: "Lost = X" ou "Perdidos = X" + lost = None + m = re.search(r"(?:Lost|Perdidos)\s*=\s*(\d+)", out, flags=re.IGNORECASE) + if m: + lost = int(m.group(1)) + else: + # fallback pelo número de replies + lost = max(0, int(pings) - replies) + + loss_pct = (lost / max(1, int(pings))) * 100.0 + + # (opcional) debug leve em caso de falha + # if not ok: + # print("PING FAIL:", base_ip, "rc:", proc.returncode) + # print(out) + + return ok, avg_rtt_ms, loss_pct + + except Exception as e: + # (opcional) print(e) + return False, None, 100.0 + def _update_nic_errors(self): if self.nic_name is None: self.nic_name = self.descobrir_nic_para_base() @@ -306,7 +501,7 @@ class ModuloIPBribge(ModuloDiagnosticoBase): loss_score = 0 # JITTER - if len(self.rtts) > 1: + if len(self.rtts) > 2: jitter = statistics.pstdev(self.rtts) if jitter <= 10: jitter_score = 100 @@ -319,7 +514,7 @@ class ModuloIPBribge(ModuloDiagnosticoBase): else: jitter = 0.0 # Sem dados suficientes, melhor considerar "desconhecido" => score neutro/baixo - jitter_score = 0 + jitter_score = 1 # NIC ERRORS if self.nic_error_rates: @@ -337,14 +532,11 @@ class ModuloIPBribge(ModuloDiagnosticoBase): nic_score = 0 # HEARTBEAT (score usando "atraso") - if atraso <= 1.5: - hb_score = 100 - elif atraso <= 2: - hb_score = 80 - elif atraso <= 5: - hb_score = 30 - else: - hb_score = 0 + t0, t1 = 1.0, 3.0 # 1s perfeito, 3s zerou + x = (atraso - t0) / (t1 - t0) # 0..1 + x = max(0.0, min(1.0, x)) + gamma = 1.6 # >1 deixa cair mais rápido perto do final + hb_score = int(round(100 * (1.0 - (x ** gamma)))) # retorna scores + brutos return ( @@ -353,8 +545,9 @@ class ModuloIPBribge(ModuloDiagnosticoBase): ) def atualizar_saude_interno(self): + #print("atualizando saude IPB...") try: - self._start_mqtt_heartbeat() + self._start_mqtt_heartbeat_async() SAUDE_MIN_ALERTA = 80 @@ -363,14 +556,27 @@ class ModuloIPBribge(ModuloDiagnosticoBase): saude_individual = [] conectado = self.get_base_ip() is not None and self.nic_name is not None and self._mqtt_conectado - (rtt, loss) = self._ping_once() - #print(f"rtt: {rtt}, loss: {loss}") + # 1) consome resultado pronto (não bloqueia) + self._consume_ping_result_if_ready() + # 2) dispara novo ping se precisar (não bloqueia) + self._start_ping_async_if_needed() + # 3) usa snapshot pra alimentar suas janelas rtts/loses/timeouts + ok, rtt, loss, ping_ts = self._get_ping_snapshot() + #print(f"ok: {ok}, rtt: {rtt}, loss: {loss}, ping_ts: {ping_ts}") self.timeouts.append(loss == 100) - if rtt is not None: + if ok and rtt is not None: self.rtts.append(rtt) if loss is not None: self.loses.append(loss) + #(rtt, loss) = self._ping_once_old(pings=3, timeout=200) + #print(f"rtt: {rtt}, loss: {loss}") + #self.timeouts.append(loss == 100) + #if rtt is not None: + # self.rtts.append(rtt) + #if loss is not None: + # self.loses.append(loss) + self._update_nic_errors() self._update_nic_bandwidth() diff --git a/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/manager_worker/modulos/movimentacao.py b/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/manager_worker/modulos/movimentacao.py index 39292a183..8b1473466 100644 --- a/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/manager_worker/modulos/movimentacao.py +++ b/AgroBase/AgroBase/bin/x64/Debug/Python/Scripts/workers/manager_worker/modulos/movimentacao.py @@ -50,10 +50,15 @@ def definir_comando(parada_necessaria: bool, erro: bool, filtro_vel: FiltroVeloc pontos_min = 3 vel_min = 15.0 - ervas_min_percenet = 0.005 # 5% + ervas_min_percenet = 0.005 # 0.5% + #pct_erva_bico_on = _contexto.get("Atu", {}).get("pct_erva_bico_on", 0.03) + #qtd_bicos = _contexto.get("Atu", {}).get("qtd_bicos", 3.0) + #ervas_min_percenet = round(pct_erva_bico_on / qtd_bicos, 5) vel_com_ervas = _operacao.get("Mov", {}).get("percent_vel_min", 0) vel_sem_ervas = _operacao.get("Mov", {}).get("percent_vel_max", 100) ang_max = _operacao.get("Dir", {}).get("angulo_max", 30) + + status_ipb = StatusModulo(ContextoGlobalRedis.get_modulo(T_Code.Ipb).get("saude", {}).get("status", StatusModulo.DESCONECTADO.value)) status_carro = StatusCarroMapa.Parado reduzir_para_pulverizar = False @@ -70,7 +75,9 @@ def definir_comando(parada_necessaria: bool, erro: bool, filtro_vel: FiltroVeloc velocidade_sp = vel_com_ervas - if status_carro == StatusCarroMapa.Parado: + if status_ipb != StatusModulo.OPERANTE: + velocidade_sp = vel_min + elif status_carro == StatusCarroMapa.Parado: velocidade_sp = 0 elif status_carro in [StatusCarroMapa.Direcionando, StatusCarroMapa.RetornandoBase]: velocidade_sp = calcular_velocidade_relativa(vel_com_ervas, vel_sem_ervas, ang_max, erro_orientacao, k=0.85, curva=0.8, dead=10.0) diff --git a/AgroBase/OperationControl/Windows/MainWindow.xaml.cs b/AgroBase/OperationControl/Windows/MainWindow.xaml.cs index 74938fa3e..8c18c2424 100644 --- a/AgroBase/OperationControl/Windows/MainWindow.xaml.cs +++ b/AgroBase/OperationControl/Windows/MainWindow.xaml.cs @@ -1262,6 +1262,12 @@ namespace OperationControl.Windows txtTrajetoria_HerbOk.Text = (dados?.HerbicidaSuficiente ?? false) ? "Sim" : "Não"; btnTrajetoria_Confirmar.IsEnabled = !(dados?.Liberado ?? false); + + if (inicial) + { + chbTrajetoria_BatLiberada.IsChecked = dados?.BateriaSuficiente ?? false; + chbTrajetoria_HerbLiberado.IsChecked = dados?.HerbicidaSuficiente ?? false; + } break; } case AgroBase.Models.Enums.T_Code.Npc: diff --git a/Python/OAK/datasets/_12_check_percent_class.py b/Python/OAK/datasets/_12_check_percent_class.py index 50f1a8ca0..8e75af963 100644 --- a/Python/OAK/datasets/_12_check_percent_class.py +++ b/Python/OAK/datasets/_12_check_percent_class.py @@ -13,6 +13,7 @@ Uso: python _12_check_percent_class_labelmap.py """ +import argparse import os, json import numpy as np from PIL import Image @@ -28,8 +29,6 @@ ROI_TAMANHO = config["roi_tamanho"] pasta_base = os.path.join(MODELO, "dataset") labelmap_path = os.path.join(pasta_base, "labelmap.txt") -root = os.path.join(pasta_base, "split", "train", "group") -#root = os.path.join(pasta_base, "576x320", "group") # Limite de amostras por grupo (para rodar rápido). Ajuste se quiser. MAX_SAMPLES_PER_GROUP = 1000 @@ -58,52 +57,62 @@ def roi_slice(h): y_fim, y_ini = max(0, h - int(ROI_TAMANHO * h)), h return slice(y_fim, y_ini) -# -------------- Labelmap -------------- -_, colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path) -ignore_id = infer_ignore_id(ignore_rgb, default_id=255) +def main(args): + # -------------- Labelmap -------------- + _, colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path) + ignore_id = infer_ignore_id(ignore_rgb, default_id=255) -# 'classes' esperado como dict: id -> nome -# Ordena por id para imprimir de forma estável -class_ids_sorted = sorted(classes.keys()) -class_names_sorted = [classes[cid] for cid in class_ids_sorted] + # 'classes' esperado como dict: id -> nome + # Ordena por id para imprimir de forma estável + class_ids_sorted = sorted(classes.keys()) + class_names_sorted = [classes[cid] for cid in class_ids_sorted] -# -------------- Coleta -------------- -if not os.path.isdir(root): - raise SystemExit(f"Nenhum diretório encontrado em {root}") + root = os.path.join(pasta_base, "split", args.split, "group") + #root = os.path.join(pasta_base, "576x320", "group") -grupos = [g for g in os.listdir(root) if os.path.isdir(os.path.join(root, g))] + # -------------- Coleta -------------- + if not os.path.isdir(root): + raise SystemExit(f"Nenhum diretório encontrado em {root}") -for g in sorted(grupos): - mdir = os.path.join(root, g, "masks") - if not os.path.isdir(mdir): - continue + grupos = [g for g in os.listdir(root) if os.path.isdir(os.path.join(root, g))] - totals = {cid: 0 for cid in class_ids_sorted} - n = 0 - - for fname in os.listdir(mdir): - if not fname.lower().endswith(".png"): + for g in sorted(grupos): + mdir = os.path.join(root, g, "masks") + if not os.path.isdir(mdir): continue - m = np.array(Image.open(os.path.join(mdir, fname)).convert("L")) - rs = roi_slice(m.shape[0]) - roi = m[rs, :] - # Acumula só das classes válidas do labelmap (ignorando 'ignore' e outros valores) + totals = {cid: 0 for cid in class_ids_sorted} + n = 0 + + for fname in os.listdir(mdir): + if not fname.lower().endswith(".png"): + continue + m = np.array(Image.open(os.path.join(mdir, fname)).convert("L")) + rs = roi_slice(m.shape[0]) + roi = m[rs, :] + + # Acumula só das classes válidas do labelmap (ignorando 'ignore' e outros valores) + for cid in class_ids_sorted: + totals[cid] += int((roi == cid).sum()) + + n += 1 + if n >= MAX_SAMPLES_PER_GROUP: + break + + s = sum(totals.values()) + s = s if s > 0 else 1 # evita div/0 + + # Monta string dinâmica "nome=xx.xx%" + parts = [] for cid in class_ids_sorted: - totals[cid] += int((roi == cid).sum()) + name = classes[cid] + perc = totals[cid] / s + parts.append(f"{name}={perc:6.2%}") - n += 1 - if n >= MAX_SAMPLES_PER_GROUP: - break + print(f"{g:16s} " + " ".join(parts) + f" (amostras={n})") - s = sum(totals.values()) - s = s if s > 0 else 1 # evita div/0 - - # Monta string dinâmica "nome=xx.xx%" - parts = [] - for cid in class_ids_sorted: - name = classes[cid] - perc = totals[cid] / s - parts.append(f"{name}={perc:6.2%}") - - print(f"{g:16s} " + " ".join(parts) + f" (amostras={n})") +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="Validacao de classes no dataset") + ap.add_argument("--split", type=str, default="train", help="Pasta para verificar (ex: train,val,test).") + args = ap.parse_args() + main(args) \ No newline at end of file diff --git a/Python/OAK/datasets/_2_create_full_mask_raw.py b/Python/OAK/datasets/_2_create_full_mask_raw.py index 7728da6b5..1149be10c 100644 --- a/Python/OAK/datasets/_2_create_full_mask_raw.py +++ b/Python/OAK/datasets/_2_create_full_mask_raw.py @@ -199,10 +199,11 @@ def localizar_raw_correspondente(pasta_new_raws, nome_img: str) -> str | None: return candidate return None -def processar_novas_imagens(classe, cor_classe_rgb, fazer_copia_final=True, manifesto_csv=None): - pasta_new_previews = os.path.join(MODELO, "dataset", "brutas", "group", classe, "previews") - pasta_new_masks = os.path.join(MODELO, "dataset", "brutas", "group", classe, "masks") - pasta_new_raws = os.path.join(MODELO, "dataset", "brutas", "group", classe, "raws") +def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_final=True, manifesto_csv=None, orignais=False): + source = os.path.join(MODELO, "dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join(MODELO, "dataset", "original") + pasta_new_previews = os.path.join(source, "group", grupo, "previews") + pasta_new_masks = os.path.join(source, "group", grupo, "masks") + pasta_new_raws = os.path.join(source, "group", grupo, "raws") garantir_pasta(pasta_new_previews) garantir_pasta(pasta_new_masks) @@ -281,8 +282,12 @@ def build_cli(): ap = argparse.ArgumentParser( description="Gera máscaras sólidas para novas imagens de UMA classe (via labelmap) e copia para dataset final com dedup." ) + ap.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.") + ap.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.") ap.add_argument("--classe", required=True, help="Nome da classe (como está no labelmap.txt). Ex: chao, cana, erva") + ap.add_argument("--grupo", required=True, help="Nome da pasta grupo. Ex: chao, cana, erva, chao_cana, chao_cana_erva") ap.add_argument("--no-copy", action="store_true", help="Não copia para as pastas finais (só cria masks em new_masks).") + ap.add_argument("--from-originals", action="store_true", help="Faz o procedimento na pasta em originais") ap.add_argument("--manifest", default="", help="Caminho do CSV de manifesto a gerar (ou vazio para não gerar).") return ap @@ -303,8 +308,11 @@ if __name__ == "__main__": manifesto_csv = args.manifest if args.manifest else None processar_novas_imagens( - classe=args.classe, + cana=args.cana, + horario=args.horario, + grupo=args.grupo, cor_classe_rgb=cor_rgb, fazer_copia_final=fazer_copia, manifesto_csv=manifesto_csv, + orignais=args.from_originals ) diff --git a/Python/OAK/datasets/_2_ingest_new_masks.py b/Python/OAK/datasets/_2_ingest_new_masks.py index 482f8978c..a130fb559 100644 --- a/Python/OAK/datasets/_2_ingest_new_masks.py +++ b/Python/OAK/datasets/_2_ingest_new_masks.py @@ -12,7 +12,7 @@ with open("config.json", "r") as f: MODELO = config["camera"] # Raiz das brutas agrupadas -PASTA_BRUTAS_GROUP_ROOT = os.path.join(MODELO, "dataset", "brutas", "group") +PASTA_BRUTAS_GROUP_ROOT = os.path.join(MODELO, "dataset", "brutas") # Onde as máscaras novas (rotuladas externamente) são colocadas PASTA_NEW_MASKS = os.path.join(MODELO, "dataset", "new_masks") @@ -180,6 +180,8 @@ def build_cli(): "nas subpastas de dataset/brutas/group/** e copiando tudo para dataset/original." ) ) + ap.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.") + ap.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.") ap.add_argument("--copy", action="store_true", help="Copia as máscaras em vez de mover (por padrão, move e esvazia new_masks conforme ingere).") ap.add_argument("--manifest", default="", help="Caminho para CSV de manifesto (opcional).") return ap @@ -191,7 +193,9 @@ if __name__ == "__main__": manifesto = args.manifest if args.manifest.strip() else None mover = not args.copy - print(f"[INFO] MODELO : {MODELO}") + PASTA_BRUTAS_GROUP_ROOT = os.path.join(PASTA_BRUTAS_GROUP_ROOT, f"cana_{args.cana}", args.horario, "group") + + print(f"[INFO] MODELO : {MODELO}") print(f"[INFO] Brutas (group root): {PASTA_BRUTAS_GROUP_ROOT}") print(f"[INFO] New masks : {PASTA_NEW_MASKS}") print(f"[INFO] Mover máscaras? : {mover}") diff --git a/Python/OAK/datasets/_3_adjust_ignore_class.py b/Python/OAK/datasets/_3_adjust_ignore_class.py deleted file mode 100644 index 1d45b573a..000000000 --- a/Python/OAK/datasets/_3_adjust_ignore_class.py +++ /dev/null @@ -1,68 +0,0 @@ -import json -import os -import cv2 -import numpy as np - -# ⚙️ Configurações -with open("config.json", "r") as f: - config = json.load(f) -MODELO = config["camera"] - -pasta_mascaras = os.path.join(MODELO, "dataset", "original", "masks") - -# Regras de substituição (cores em RGB) -# Exemplo: trocar (255, 0, 0) por branco (255,255,255) com tolerância 10 -SUBSTITUICOES = [ - #{"target_rgb": (255, 255, 255), "tolerancia": 10, "replace_rgb": (128, 0, 0)}, # branco -> vermelho - - #{"target_rgb": (128, 0, 0), "tolerancia": 50, "replace_rgb": (128, 0, 0)}, # chao - #{"target_rgb": (0, 128, 0), "tolerancia": 50, "replace_rgb": (0, 128, 0)}, # erva - #{"target_rgb": (0, 0, 128), "tolerancia": 50, "replace_rgb": (0, 0, 128)}, # cana - - {"target_rgb": (128, 0, 0), "tolerancia": 50, "replace_rgb": (128, 0, 0)}, # chao - {"target_rgb": (0, 128, 0), "tolerancia": 50, "replace_rgb": (0, 128, 0)}, # cana - {"target_rgb": (0, 0, 128), "tolerancia": 50, "replace_rgb": (0, 0, 128)}, # obstaculo -] - -def aplicar_substituicoes(img_bgr): - """Recebe imagem BGR (OpenCV) e aplica regras RGB com tolerância, de forma vetorizada.""" - # Converte uma vez pra RGB só para fazer o match nas cores “humanas” - img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) - - for rule in SUBSTITUICOES: - (tr, tg, tb) = rule["target_rgb"] - tol = int(rule.get("tolerancia", 0)) - (rr, rg, rb) = rule["replace_rgb"] - - # Faixa inferior/superior com tolerância (clamp 0..255) - lower = np.array([max(tr - tol, 0), max(tg - tol, 0), max(tb - tol, 0)], dtype=np.uint8) - upper = np.array([min(tr + tol, 255), min(tg + tol, 255), min(tb + tol, 255)], dtype=np.uint8) - - # Máscara booleana dos pixels a substituir - mask = cv2.inRange(img_rgb, lower, upper) # 255 onde bate - - if np.any(mask): - # Cria uma imagem de destino RGB com a cor de replace - replace_rgb = np.zeros_like(img_rgb) - replace_rgb[:] = (rr, rg, rb) - # Faz o blend: onde mask==255, põe replace; onde não, mantém original - img_rgb = np.where(mask[..., None] == 255, replace_rgb, img_rgb) - - # Volta pra BGR pro OpenCV salvar - return cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) - -def corrigir_mascara(caminho_img): - img_bgr = cv2.imread(caminho_img, cv2.IMREAD_COLOR) - if img_bgr is None: - print(f"[ERRO] Não abriu: {caminho_img}") - return - - out_bgr = aplicar_substituicoes(img_bgr) - if not np.array_equal(out_bgr, img_bgr): - cv2.imwrite(caminho_img, out_bgr) - print(f"Ajustado: {os.path.basename(caminho_img)}") - -if __name__ == "__main__": - for nome in os.listdir(pasta_mascaras): - if nome.lower().endswith((".png", ".jpg", ".jpeg")): - corrigir_mascara(os.path.join(pasta_mascaras, nome)) diff --git a/Python/OAK/datasets/_3_change_mask_color.py b/Python/OAK/datasets/_3_change_mask_color.py new file mode 100644 index 000000000..2ea0659db --- /dev/null +++ b/Python/OAK/datasets/_3_change_mask_color.py @@ -0,0 +1,147 @@ +import argparse +from pathlib import Path + +import numpy as np +from PIL import Image + + +def process_mask_rgb( + in_path: Path, + out_path: Path, + src_rgb: tuple[int, int, int], + dst_rgb: tuple[int, int, int], + keep_other_colors: bool, +): + """ + Troca pixels de uma cor RGB exata por outra cor RGB. + + - src_rgb: (R, G, B) que será substituída + - dst_rgb: (R, G, B) nova cor + - keep_other_colors: + True -> mantém outras cores como estão + False -> tudo que não for src_rgb vira preto (0,0,0) + """ + img = Image.open(in_path).convert("RGB") + arr = np.array(img) # (H, W, 3) uint8 + + src_r, src_g, src_b = src_rgb + dst_r, dst_g, dst_b = dst_rgb + + # Máscara de pixels que são exatamente a cor origem + mask = ( + (arr[:, :, 0] == src_r) + & (arr[:, :, 1] == src_g) + & (arr[:, :, 2] == src_b) + ) + + if keep_other_colors: + # Só troca onde é src_rgb + arr[mask] = np.array([dst_r, dst_g, dst_b], dtype=np.uint8) + else: + # Primeiro zera tudo que não é src_rgb + arr[~mask] = np.array([0, 0, 0], dtype=np.uint8) + # Depois coloca dst_rgb onde era src_rgb + arr[mask] = np.array([dst_r, dst_g, dst_b], dtype=np.uint8) + + out_path.parent.mkdir(parents=True, exist_ok=True) + Image.fromarray(arr, mode="RGB").save(out_path) + + +def parse_rgb(s: str) -> tuple[int, int, int]: + """ + Converte string no formato "R,G,B" para tupla (R,G,B). + Ex: "128,0,0" -> (128, 0, 0) + """ + parts = s.split(",") + if len(parts) != 3: + raise argparse.ArgumentTypeError( + f"RGB inválido: '{s}'. Use o formato R,G,B (ex: 128,0,0)." + ) + try: + r, g, b = [int(p.strip()) for p in parts] + except ValueError: + raise argparse.ArgumentTypeError( + f"RGB inválido: '{s}'. Valores devem ser inteiros 0..255." + ) + for v in (r, g, b): + if not (0 <= v <= 255): + raise argparse.ArgumentTypeError( + f"RGB inválido: '{s}'. Cada valor deve estar entre 0 e 255." + ) + return (r, g, b) + + +def main(): + parser = argparse.ArgumentParser( + description="Troca uma cor RGB exata por outra em máscaras (png/jpg/etc)." + ) + parser.add_argument( + "--masks_dir", + type=str, + required=True, + help="Pasta com as máscaras de entrada (procura recursivamente por png/jpg/etc).", + ) + parser.add_argument( + "--out_dir", + type=str, + default=None, + help="Pasta de saída. Se não for informada, sobrescreve as máscaras originais.", + ) + parser.add_argument( + "--src_rgb", + type=parse_rgb, + required=True, + help='Cor origem no formato "R,G,B" (ex: "128,0,0").', + ) + parser.add_argument( + "--dst_rgb", + type=parse_rgb, + required=True, + help='Cor destino no formato "R,G,B" (ex: "0,0,0").', + ) + parser.add_argument( + "--keep_other_colors", + action="store_true", + help="Se setado, mantém outras cores como estão. " + "Se omitido, tudo que não for src_rgb vira preto (0,0,0).", + ) + + args = parser.parse_args() + + masks_dir = Path(args.masks_dir) + if not masks_dir.is_dir(): + raise SystemExit(f"Pasta de máscaras não encontrada: {masks_dir}") + + if args.out_dir is None: + out_dir = masks_dir + print(f"[INFO] Sem out_dir, sobrescrevendo arquivos em {masks_dir}") + else: + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + print(f"[INFO] Salvando máscaras convertidas em {out_dir}") + + exts = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"} + files = [p for p in masks_dir.rglob("*") if p.suffix.lower() in exts] + + print(f"[INFO] Encontrados {len(files)} arquivos de máscara.") + + for i, in_path in enumerate(files, 1): + rel = in_path.relative_to(masks_dir) + out_path = out_dir / rel + + process_mask_rgb( + in_path=in_path, + out_path=out_path, + src_rgb=args.src_rgb, + dst_rgb=args.dst_rgb, + keep_other_colors=args.keep_other_colors, + ) + + if i % 50 == 0 or i == len(files): + print(f"[INFO] Processados {i}/{len(files)} arquivos") + + print("[OK] Conversão concluída.") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/Python/OAK/datasets/_3_copy_selected_images_to_mask.py b/Python/OAK/datasets/_3_copy_selected_images_to_mask.py new file mode 100644 index 000000000..c1a05bf0f --- /dev/null +++ b/Python/OAK/datasets/_3_copy_selected_images_to_mask.py @@ -0,0 +1,328 @@ +import json +import os +import shutil +import argparse +import csv + +# ================= CONFIG ================= + +with open("config.json", "r") as f: + config = json.load(f) + +MODELO = config["camera"] + +# Raiz das brutas (todas as canas/horários/grupos) +PASTA_BRUTAS_ROOT = os.path.join(MODELO, "dataset", "brutas") + +# Onde você coloca os previews selecionados (tudo misturado) +PASTA_SELECTED_PREVIEWS = os.path.join(MODELO, "dataset", "selected_previews") + +# Destino final, organizado por grupo: +# dataset/original/group/{GRUPO}/{previews,raws,metas,masks} +PASTA_ORIGINAL_GROUP_ROOT = os.path.join(MODELO, "dataset", "original", "group") + +EXT_PREVIEWS = (".png", ".jpg", ".jpeg") +EXT_RAWS = (".raw",) +EXT_MASKS = (".png",) +# Ajusta se suas metas tiverem outra extensão +EXT_METAS = (".json", ".yml", ".yaml", ".txt", ".csv") + + +# ================= HELPERS ================= + +def garantir_pasta(p: str): + os.makedirs(p, exist_ok=True) + + +def nome_disponivel(dest_dir: str, base: str, ext: str) -> str: + """ + Devolve um caminho disponível em dest_dir com base + ext. + Se já existir, adiciona _001, _002, ... + """ + p = os.path.join(dest_dir, base + ext) + if not os.path.exists(p): + return p + + i = 1 + while True: + p = os.path.join(dest_dir, f"{base}_{i:03d}{ext}") + if not os.path.exists(p): + return p + i += 1 + + +def buscar_por_base(pasta: str, base: str, exts: tuple[str, ...]) -> str | None: + """ + Procura um arquivo em 'pasta' com o mesmo 'base' e qualquer extensão em 'exts'. + Retorna o caminho completo ou None. + """ + if not os.path.isdir(pasta): + return None + + for nome in os.listdir(pasta): + nome_lower = nome.lower() + root, ext = os.path.splitext(nome_lower) + if root == base.lower() and ext in exts: + return os.path.join(pasta, nome) + + return None + + +def indexar_brutas_por_raw(root_brutas: str): + """ + Varre dataset/brutas recursivamente, olhando pastas 'raws' e montando índice: + + base -> { + "raw": caminho_raw, + "cana": "cana_alta" | "cana_baixa" | ... (se conseguir inferir), + "horario": "meio_dia" | "cedo" | ... (se conseguir inferir), + "grupo": "", + "group_dir": caminho_da_pasta_do_grupo + } + + Estrutura esperada (relativa a root_brutas): + cana_alta/meio_dia/group//raws/*.raw + """ + index = {} + + if not os.path.isdir(root_brutas): + print(f"[AVISO] Pasta de brutas não existe: {root_brutas}") + return index + + for dirpath, dirnames, filenames in os.walk(root_brutas): + base_dir = os.path.basename(dirpath).lower() + + if base_dir != "raws": + continue + + # dirpath = .../cana_x/horario/group//raws + group_dir = os.path.dirname(dirpath) # .../cana_x/horario/group/ + + rel = os.path.relpath(dirpath, root_brutas) + parts = rel.split(os.sep) + + # Defaults + cana = None + horario = None + grupo = None + + if len(parts) >= 5: + # [0] = cana_<...> + # [1] = horario + # [2] = "group" + # [3] = + cana = parts[0] + horario = parts[1] + # parts[2] deve ser "group" + grupo = parts[3] + else: + # fallback bem genérico + if "group" in parts: + i = parts.index("group") + if i + 1 < len(parts): + grupo = parts[i + 1] + + for nome in filenames: + if not nome.lower().endswith(EXT_RAWS): + continue + + base = os.path.splitext(nome)[0] + + if base in index: + # Conflito (mesmo base em dois lugares) -> loga e mantém o primeiro + print(f"[CONFLITO] base repetida em raws: {base}") + continue + + raw_path = os.path.join(dirpath, nome) + + index[base] = { + "raw": raw_path, + "cana": cana, + "horario": horario, + "grupo": grupo, + "group_dir": group_dir, + } + + print(f"[INDEX] Entradas indexadas por RAW: {len(index)}") + return index + + +def organizar_selected_previews(copy_only: bool = False, manifesto_csv: str | None = None): + if not os.path.isdir(PASTA_SELECTED_PREVIEWS): + raise SystemExit(f"[ERRO] Pasta selected_previews não existe: {PASTA_SELECTED_PREVIEWS}") + + # Indexa todas as brutas a partir dos RAWs + index_raws = indexar_brutas_por_raw(PASTA_BRUTAS_ROOT) + + registros = [] + total, movidos, ignorados = 0, 0, 0 + + for nome in os.listdir(PASTA_SELECTED_PREVIEWS): + caminho_preview_sel = os.path.join(PASTA_SELECTED_PREVIEWS, nome) + + if not os.path.isfile(caminho_preview_sel): + continue + + if not nome.lower().endswith(EXT_PREVIEWS): + continue + + total += 1 + + base = os.path.splitext(nome)[0] + + info = index_raws.get(base) + if info is None: + ignorados += 1 + print(f"[SKIP] {base} -> não encontrado em dataset/brutas (via RAW)") + continue + + grupo = info.get("grupo") or "unknown" + cana = info.get("cana") or "unknown" + horario = info.get("horario") or "unknown" + group_dir = info["group_dir"] + raw_src = info["raw"] + + # Pastas irmãs em brutas + metas_src_dir = os.path.join(group_dir, "metas") + masks_src_dir = os.path.join(group_dir, "masks") + + meta_src = buscar_por_base(metas_src_dir, base, EXT_METAS) + mask_src = buscar_por_base(masks_src_dir, base, EXT_MASKS) + + # Destino: dataset/original/group/{GRUPO}/{previews,raws,metas,masks} + dest_group_root = os.path.join(PASTA_ORIGINAL_GROUP_ROOT, grupo) + dest_prev_dir = os.path.join(dest_group_root, "previews") + dest_raw_dir = os.path.join(dest_group_root, "raws") + dest_meta_dir = os.path.join(dest_group_root, "metas") + dest_mask_dir = os.path.join(dest_group_root, "masks") + + garantir_pasta(dest_prev_dir) + garantir_pasta(dest_raw_dir) + garantir_pasta(dest_meta_dir) + garantir_pasta(dest_mask_dir) + + # Define extensões + prev_ext_sel = os.path.splitext(nome)[1].lower() + raw_ext = os.path.splitext(raw_src)[1].lower() + + # Gera nome final único com base no preview + dst_preview = nome_disponivel(dest_prev_dir, base, prev_ext_sel) + new_base = os.path.splitext(os.path.basename(dst_preview))[0] + + dst_raw = os.path.join(dest_raw_dir, new_base + raw_ext) + dst_meta = None + dst_mask = None + + if meta_src is not None: + meta_ext = os.path.splitext(meta_src)[1].lower() + dst_meta = os.path.join(dest_meta_dir, new_base + meta_ext) + + if mask_src is not None: + mask_ext = os.path.splitext(mask_src)[1].lower() + dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext) + + # Copia/move preview selecionado + if copy_only: + shutil.copy2(caminho_preview_sel, dst_preview) + else: + shutil.move(caminho_preview_sel, dst_preview) + + # Copia RAW e metas/masks (se existirem) + if copy_only: + shutil.copy2(raw_src, dst_raw) + else: + shutil.move(raw_src, dst_raw) + + if meta_src is not None: + if copy_only: + shutil.copy2(meta_src, dst_meta) + else: + shutil.move(meta_src, dst_meta) + + if mask_src is not None: + if copy_only: + shutil.copy2(mask_src, dst_mask) + else: + shutil.move(mask_src, dst_mask) + + movidos += 1 + + registros.append([ + base, + grupo, + cana, + horario, + caminho_preview_sel, + raw_src, + meta_src or "", + mask_src or "", + dst_preview, + dst_raw, + dst_meta or "", + dst_mask or "", + ]) + + print(f"[OK] {new_base} -> grupo={grupo} | cana={cana} | horario={horario}") + + if manifesto_csv and registros: + with open(manifesto_csv, "w", newline="", encoding="utf-8") as f: + w = csv.writer(f) + w.writerow([ + "base", + "grupo", + "cana", + "horario", + "src_preview_selected", + "src_raw", + "src_meta", + "src_mask", + "dst_preview", + "dst_raw", + "dst_meta", + "dst_mask", + ]) + w.writerows(registros) + print(f"[MANIFESTO] {manifesto_csv} salvo ({len(registros)} entradas).") + + print(f"\nResumo: total_selected={total} | organizadas={movidos} | ignoradas={ignorados}") + + +# ================= CLI ================= + +def build_cli(): + ap = argparse.ArgumentParser( + description=( + "Organiza previews selecionadas (dataset/selected_previews) " + "descobrindo cana/horário/grupo em dataset/brutas e copiando/movendo " + "preview + raw + meta (+ mask se existir) para dataset/original/group/{GRUPO}." + ) + ) + ap.add_argument( + "--copy", + action="store_true", + help="Copia os previews em vez de mover (por padrão, move e esvazia selected_previews conforme organiza)." + ) + ap.add_argument( + "--manifest", + default="", + help="Caminho para CSV de manifesto (opcional)." + ) + return ap + + +if __name__ == "__main__": + ap = build_cli() + args = ap.parse_args() + + manifesto = args.manifest if args.manifest.strip() else None + copy_only = args.copy + + print(f"[INFO] MODELO : {MODELO}") + print(f"[INFO] Brutas root : {PASTA_BRUTAS_ROOT}") + print(f"[INFO] Selected previews : {PASTA_SELECTED_PREVIEWS}") + print(f"[INFO] Original group root : {PASTA_ORIGINAL_GROUP_ROOT}") + print(f"[INFO] Copy only? : {copy_only}") + print(f"[INFO] Manifesto : {manifesto or '(nenhum)'}") + print() + + organizar_selected_previews(copy_only=copy_only, manifesto_csv=manifesto) \ No newline at end of file diff --git a/Python/OAK/datasets/_4_group_images_by_class.py b/Python/OAK/datasets/_4_group_images_by_class.py index 07d14e6a6..e1382d3a0 100644 --- a/Python/OAK/datasets/_4_group_images_by_class.py +++ b/Python/OAK/datasets/_4_group_images_by_class.py @@ -25,7 +25,7 @@ import shutil import argparse import numpy as np -from utils import carregar_labelmap_completo +from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids # ====================== Configurações base ====================== @@ -33,19 +33,22 @@ def carregar_config_e_paths(): with open("config.json", "r", encoding="utf-8") as f: config = json.load(f) MODELO = config.get("camera") + USE_MASKS2 = config.get("dual_head", False) pasta_base = os.path.join(MODELO, "dataset") labelmap_path = os.path.join(pasta_base, "labelmap.txt") # Pastas origem/destino PASTA_NEW_IMAGES = os.path.join(pasta_base, "original", "images") PASTA_NEW_MASKS = os.path.join(pasta_base, "original", "masks") + PASTA_NEW_MASKS2 = os.path.join(pasta_base, "original", "masks2") PASTA_FINAL = os.path.join(pasta_base, "original", "group") - return MODELO, pasta_base, labelmap_path, PASTA_NEW_IMAGES, PASTA_NEW_MASKS, PASTA_FINAL + return MODELO, pasta_base, labelmap_path, PASTA_NEW_IMAGES, PASTA_NEW_MASKS, PASTA_NEW_MASKS2, USE_MASKS2, PASTA_FINAL # Extensões aceitas EXT_IMAGENS = (".jpg", ".jpeg", ".png") EXT_MASKS = (".png", ".jpg", ".jpeg") # prioridade será .png quando houver +EXT_MASKS2 = (".png", ".jpg", ".jpeg") # idem # Manifesto padrão MANIFESTO_DEFAULT = "manifest.csv" @@ -84,6 +87,25 @@ def mapear_masks_por_base(pasta_masks): mapa[base] = cam return mapa +def mapear_masks2_por_base(pasta_masks2): + """Retorna {base: caminho_mask2}, priorizando .png se houver múltiplas por base.""" + if not pasta_masks2 or not os.path.isdir(pasta_masks2): + return {} + mapa = {} + for nome in os.listdir(pasta_masks2): + lower = nome.lower() + if not lower.endswith(EXT_MASKS2): + continue + base, ext = os.path.splitext(nome) + cam = os.path.join(pasta_masks2, nome) + if base not in mapa: + mapa[base] = cam + else: + atual_ext = os.path.splitext(mapa[base])[1].lower() + if atual_ext != ".png" and ext.lower() == ".png": + mapa[base] = cam + return mapa + def localizar_imagem_por_base(pasta_imgs, base): """Retorna caminho da imagem correspondente ao base se existir.""" for ext in EXT_IMAGENS: @@ -118,55 +140,70 @@ def inferir_ignore_id(ignore_rgb, cor_para_id): def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True): """ - Extrai IDs de classes presentes na máscara. - - Se a máscara for 1 canal: retorna valores únicos como IDs diretamente. - - Se for 3 canais: pega cores únicas (BGR), converte para RGB (se assume_rgb=True), - e mapeia usando cor_para_id[(R,G,B)] -> id. + Versão corrigida e otimizada: + + - Se máscara for 1 canal: np.unique direto -> IDs. + - Se for 3 canais: + * se assume_rgb=True: labelmap está em RGB, + mas OpenCV lê BGR -> convertemos BGR -> RGB. + * se assume_rgb=False: labelmap está em BGR, + mantemos BGR como está. + - Usa converter_mask_rgb_para_ids em amostragem + fallback full-scan. + Retorna: set(ids_presentes) """ m = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED) if m is None: raise RuntimeError(f"Falha ao abrir máscara: {mask_path}") - # grayscale / paleta indexada + # --------------------------- + # CASO 1: máscara indexada + # --------------------------- if len(m.shape) == 2 or (len(m.shape) == 3 and m.shape[2] == 1): - vals = np.unique(m).tolist() + vals = np.unique(m) return set(int(v) for v in vals) - # 3 canais (OpenCV lê BGR) - h, w, c = m.shape - flat = m.reshape(-1, 3) - uniq_bgr = np.unique(flat, axis=0) - ids = set() - for b, g, r in uniq_bgr: - if assume_rgb: - key = (int(r), int(g), int(b)) # converte para RGB - else: - key = (int(b), int(g), int(r)) # já em BGR no labelmap - id_ = cor_para_id.get(key) - if id_ is not None: - try: - ids.add(int(id_)) - except Exception: - pass + # --------------------------- + # CASO 2: máscara RGB + # --------------------------- + # Se o labelmap está em RGB (assume_rgb=True), + # convertemos a imagem BGR->RGB para casar com as chaves. + if assume_rgb: + img = cv2.cvtColor(m, cv2.COLOR_BGR2RGB) + else: + # labelmap já está em BGR; OpenCV entrega BGR; deixa como está + img = m + + # Aqui as chaves de cor_para_id estão no MESMO espaço de cor da imagem. + mapa_rgb = cor_para_id + max_classes = len(cor_para_id) + + # ---------- AMOSTRAGEM RÁPIDA ---------- + step = 8 # pode ajustar para 4 se quiser mais precisão + amostra = img[::step, ::step] + amostra_ids = converter_mask_rgb_para_ids(amostra, mapa_rgb, ignore_id=255) + ids = set(int(x) for x in np.unique(amostra_ids) if x != 255) + + if len(ids) >= max_classes: + return ids + + # ---------- FULL-SCAN (fallback) ---------- + full_ids = converter_mask_rgb_para_ids(img, mapa_rgb, ignore_id=255) + ids = set(int(x) for x in np.unique(full_ids) if x != 255) return ids def montar_nome_grupo(ids_presentes, id_para_nome): """ - Constrói o nome do grupo a partir dos nomes das classes dos IDs presentes. - Preferência de ordenação: chao < erva < cana; demais nomes em ordem alfabética. + Constrói o nome do grupo respeitando a ordem natural dos IDs do labelmap. + Ex: {0,1,2} -> chao_cana_obstaculo """ - nomes = [] - for cid in sorted(ids_presentes): - nome = id_para_nome.get(cid, str(cid)) - nomes.append(nome) + #print(f"ids_presentes: {ids_presentes}") + if not ids_presentes: + return "sem_classe" + nomes = [id_para_nome.get(cid, str(cid)) for cid in sorted(ids_presentes)] + return "_".join(nomes) - # aplicar ordenação preferida quando disponíveis - prefer = {"chao": 0, "erva": 1, "cana": 2} - nomes = sorted(nomes, key=lambda n: (prefer.get(n, 99), n)) - return "_".join(nomes) if nomes else "sem_classe" - -def copiar_ou_mover(img_src, mask_src, dest_img_dir, dest_mask_dir, mover=False): +def copiar_ou_mover(img_src, mask_src, dest_img_dir, dest_mask_dir, mover=False, mask2_src=None, dest_mask2_dir=None): garantir_pasta(dest_img_dir) garantir_pasta(dest_mask_dir) @@ -177,6 +214,7 @@ def copiar_ou_mover(img_src, mask_src, dest_img_dir, dest_mask_dir, mover=False) dst_img = nome_disponivel(dest_img_dir, base_img, img_ext) new_base = os.path.splitext(os.path.basename(dst_img))[0] dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext) + dst_mask2 = None if os.path.exists(dst_mask): # evita colisão invertendo a ordem do "único" para a máscara @@ -186,37 +224,54 @@ def copiar_ou_mover(img_src, mask_src, dest_img_dir, dest_mask_dir, mover=False) if os.path.exists(dst_img): dst_img = nome_disponivel(dest_img_dir, new_base, img_ext) + # se tiver mask2, usa o MESMO new_base final + if mask2_src and dest_mask2_dir: + garantir_pasta(dest_mask2_dir) + mask2_ext = os.path.splitext(mask2_src)[1].lower() + dst_mask2 = os.path.join(dest_mask2_dir, new_base + mask2_ext) + if os.path.exists(dst_mask2): + dst_mask2 = nome_disponivel(dest_mask2_dir, new_base, mask2_ext) + if mover: shutil.move(img_src, dst_img) shutil.move(mask_src, dst_mask) + if mask2_src and dst_mask2: + shutil.move(mask2_src, dst_mask2) else: shutil.copy2(img_src, dst_img) shutil.copy2(mask_src, dst_mask) + if mask2_src and dst_mask2: + shutil.copy2(mask2_src, dst_mask2) - return dst_img, dst_mask + return dst_img, dst_mask, dst_mask2 # ====================== Pipeline principal ====================== def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT, validar_dim=True, estrito=False, labelmap_bgr=False): # carrega config/paths - MODELO, pasta_base, labelmap_path, PASTA_NEW_IMAGES, PASTA_NEW_MASKS, PASTA_FINAL = carregar_config_e_paths() + MODELO, pasta_base, labelmap_path, PASTA_NEW_IMAGES, PASTA_NEW_MASKS, PASTA_NEW_MASKS2, USE_MASKS2, PASTA_FINAL = carregar_config_e_paths() # carrega labelmap completo cor_para_id, _colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path) ignore_id = inferir_ignore_id(ignore_rgb, cor_para_id) + #print(f"cor_para_id: {cor_para_id}, _colormap_rgb: {_colormap_rgb}, id_para_nome: {id_para_nome}") # garante pastas garantir_pasta(PASTA_NEW_IMAGES) garantir_pasta(PASTA_NEW_MASKS) garantir_pasta(PASTA_FINAL) + usar_masks2 = USE_MASKS2 and os.path.isdir(PASTA_NEW_MASKS2) + if usar_masks2: + garantir_pasta(PASTA_NEW_MASKS2) + print(f"[INFO] masks2 detectada: {PASTA_NEW_MASKS2}") # indexa máscaras mapa_masks = mapear_masks_por_base(PASTA_NEW_MASKS) + mapa_masks2 = mapear_masks2_por_base(PASTA_NEW_MASKS2) if usar_masks2 else {} registros = [] - totais = {"total_masks":0, "processados":0, "pulados":0, "sem_imagem":0, - "dim_mismatch":0, "erros":0} + totais = {"total_masks":0, "processados":0, "pulados":0, "sem_imagem":0, "sem_mask2":0, "dim_mismatch":0, "erros":0} por_grupo = {} for base, mask_path in sorted(mapa_masks.items()): @@ -228,6 +283,11 @@ def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT, continue try: + mask2_path = mapa_masks2.get(base) if usar_masks2 else None + if usar_masks2 and not mask2_path: + totais["sem_mask2"] += 1 + print(f"[AVISO] masks2 existe, mas não achei mask2 para base '{base}' (vou agrupar mesmo).") + if validar_dim: try: img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) @@ -245,6 +305,24 @@ def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT, continue else: print(msg + " → copiando mesmo assim.") + + if mask2_path: + m2 = cv2.imread(mask2_path, cv2.IMREAD_UNCHANGED) + if m2 is None: + print(f"[AVISO] Falha ao abrir mask2: {mask2_path} → ignorando mask2.") + mask2_path = None + else: + h2, w2 = m2.shape[:2] + if (hi, wi) != (h2, w2): + totais["dim_mismatch"] += 1 + msg2 = f"[AVISO] Dimensões diferem (img {wi}x{hi} vs mask2 {w2}x{h2}) para base '{base}'" + if estrito: + print(msg2 + " → pulando (por mask2).") + totais["pulados"] += 1 + continue + else: + print(msg2 + " → copiando mesmo assim.") + except Exception as e_dim: print(f"[AVISO] Falha ao validar dimensões: {e_dim} → copiando mesmo assim.") @@ -260,14 +338,17 @@ def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT, dest_base = os.path.join(PASTA_FINAL, grupo) dest_img_dir = os.path.join(dest_base, "images") dest_mask_dir = os.path.join(dest_base, "masks") + dest_mask2_dir = os.path.join(dest_base, "masks2") if usar_masks2 else None - dst_img, dst_mask = copiar_ou_mover(img_path, mask_path, dest_img_dir, dest_mask_dir, mover=mover) + dst_img, dst_mask, dst_mask2 = copiar_ou_mover(img_path, mask_path, dest_img_dir, dest_mask_dir, mover=mover, mask2_src=mask2_path, dest_mask2_dir=dest_mask2_dir) totais["processados"] += 1 - registros.append([img_path, mask_path, dst_img, dst_mask, grupo]) + registros.append([img_path, mask_path, mask2_path or "", dst_img, dst_mask, dst_mask2 or "", grupo]) por_grupo[grupo] = por_grupo.get(grupo, 0) + 1 print(f"[OK] {os.path.basename(dst_img)} → grupo: {grupo}") + extra = " +mask2" if dst_mask2 else "" + print(f"[OK] {os.path.basename(dst_img)}{extra} → grupo: {grupo}") except Exception as e: totais["erros"] += 1 @@ -277,7 +358,7 @@ def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT, if manifesto and registros: with open(manifesto, "w", newline="", encoding="utf-8") as f: w = csv.writer(f) - w.writerow(["src_image", "src_mask", "dst_image", "dst_mask", "grupo"]) + w.writerow(["src_image", "src_mask", "src_mask2", "dst_image", "dst_mask", "dst_mask2", "grupo"]) w.writerows(registros) print(f"[MANIFESTO] {manifesto} salvo ({len(registros)} entradas).") diff --git a/Python/OAK/datasets/_4_group_images_by_class_raw.py b/Python/OAK/datasets/_4_group_images_by_class_raw.py index 9c07c92d9..4a19f1bf7 100644 --- a/Python/OAK/datasets/_4_group_images_by_class_raw.py +++ b/Python/OAK/datasets/_4_group_images_by_class_raw.py @@ -49,11 +49,16 @@ MANIFESTO_DEFAULT = "manifest.csv" def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True): """ Extração ultra-rápida de IDs presentes na máscara. + Estratégia: - 1) Tenta máscara 1 canal → np.unique (instrantâneo) - 2) Amostragem em grid → converter_mask_rgb_para_ids → unique - 3) Early exit quando todas classes forem encontradas - 4) Fallback full-scan apenas se necessário (raríssimo) + 1) Máscara 1 canal → np.unique (instantâneo). + 2) Máscara 3 canais: + - se assume_rgb=True: labelmap está em RGB, + convertemos a imagem BGR -> RGB para casar com cor_para_id. + - se assume_rgb=False: labelmap está em BGR, + mantemos a imagem em BGR. + 3) Amostragem em grid + converter_mask_rgb_para_ids. + 4) Fallback full-scan se necessário. """ m = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED) if m is None: @@ -63,39 +68,34 @@ def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True): # CASO 1: máscara indexada (1 canal) — instantâneo # ------------------------------------------------------- if len(m.shape) == 2 or (len(m.shape) == 3 and m.shape[2] == 1): - ids = np.unique(m).tolist() + ids = np.unique(m) return set(int(v) for v in ids) # ------------------------------------------------------- - # CASO 2: máscara RGB — LUT + amostragem + # CASO 2: máscara RGB # ------------------------------------------------------- - h, w, _ = m.shape - - # monta LUT RGB->ID - # assume_rgb=True => cor=(R,G,B) if assume_rgb: - mapa_rgb = {cor: cid for cor, cid in cor_para_id.items()} + # Labelmap em RGB, OpenCV em BGR -> converte + img = cv2.cvtColor(m, cv2.COLOR_BGR2RGB) else: - mapa_rgb = {(b, g, r): cid for (r, g, b), cid in cor_para_id.items()} + # Labelmap em BGR, OpenCV já em BGR -> usa direto + img = m + # Agora cor_para_id e img estão no MESMO espaço de cor + mapa_rgb = cor_para_id max_classes = len(cor_para_id) # ---------- AMOSTRAGEM ---------- - step = 8 - amostra = m[::step, ::step] # reduz drasticamente o custo + step = 8 # pode virar 4 se quiser mais precisão, 16 se quiser mais velocidade + amostra = img[::step, ::step] amostra_ids = converter_mask_rgb_para_ids(amostra, mapa_rgb, ignore_id=255) ids = set(int(x) for x in np.unique(amostra_ids) if x != 255) - if len(ids) == max_classes: - return ids - - # ---------- EARLY EXIT ---------- - # Se já veio todas as classes possíveis, não precisamos do full-scan if len(ids) >= max_classes: return ids # ---------- FULL-SCAN OTIMIZADO (último caso) ---------- - full_ids = converter_mask_rgb_para_ids(m, mapa_rgb, ignore_id=255) + full_ids = converter_mask_rgb_para_ids(img, mapa_rgb, ignore_id=255) ids = set(int(x) for x in np.unique(full_ids) if x != 255) return ids @@ -146,10 +146,14 @@ def inferir_ignore_id(ignore_rgb, cor_para_id): return None def montar_nome_grupo(ids_presentes, id_para_nome): + """ + Constrói o nome do grupo respeitando a ordem natural dos IDs do labelmap. + Ex: {0,1,2} -> chao_cana_obstaculo + """ + if not ids_presentes: + return "sem_classe" nomes = [id_para_nome.get(cid, str(cid)) for cid in sorted(ids_presentes)] - prefer = {"chao": 0, "erva": 1, "cana": 2} - nomes = sorted(nomes, key=lambda n: (prefer.get(n, 99), n)) - return "_".join(nomes) if nomes else "sem_classe" + return "_".join(nomes) # ============================================================ diff --git a/Python/OAK/datasets/_5_augmentation.py b/Python/OAK/datasets/_5_augmentation.py index 40886d499..6cc449a7e 100644 --- a/Python/OAK/datasets/_5_augmentation.py +++ b/Python/OAK/datasets/_5_augmentation.py @@ -32,6 +32,7 @@ import argparse with open("config.json", "r", encoding="utf-8") as f: config = json.load(f) MODELO = config.get("camera", ".") +USE_MASKS2 = config.get("dual_head", False) # Pastas base DATASET_BASE = os.path.join(MODELO, "dataset") @@ -41,12 +42,15 @@ AUG_GROUP_ROOT = os.path.join(DATASET_BASE, "augmented", "group") # Fallback (modo antigo, sem grupos) ORIG_OLD_IMG = os.path.join(DATASET_BASE, "original", "images") ORIG_OLD_MSK = os.path.join(DATASET_BASE, "original", "masks") +ORIG_OLD_MSK2 = os.path.join(DATASET_BASE, "original", "masks2") AUG_OLD_IMG = os.path.join(DATASET_BASE, "augmented", "images") AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks") +AUG_OLD_MSK2 = os.path.join(DATASET_BASE, "augmented", "masks2") # Extensões aceitas IMG_EXTS = (".jpg", ".jpeg", ".png") MSK_EXTS = (".png", ".jpg", ".jpeg") # manter prioridade PNG quando possível +MSK2_EXTS = (".png", ".jpg", ".jpeg") def garantir_dir(p): os.makedirs(p, exist_ok=True) @@ -86,7 +90,9 @@ train_tf = A.Compose([ A.RandomSunFlare(p=0.10), A.ChannelShuffle(p=0.05), A.CoarseDropout(max_holes=6, max_height=16, max_width=16, p=0.10), -], additional_targets={'mask':'mask'}) +], additional_targets={ + 'mask2': 'mask' +}) def load_rgb(path): # cv2 lê BGR → converte para RGB @@ -131,31 +137,60 @@ def map_masks_by_base(msk_dir): by_base[base] = cand return by_base -def ensure_aug_dirs(group_name=None): - """Cria diretórios de saída para o grupo ou modo antigo.""" +def map_masks2_by_base(msk2_dir): + """Mapeia máscaras2 por base (prioriza .png).""" + by_base = {} + if not os.path.isdir(msk2_dir): + return by_base + for fname in os.listdir(msk2_dir): + f_lower = fname.lower() + if not f_lower.endswith(MSK2_EXTS): + continue + base, ext = os.path.splitext(fname) + cand = os.path.join(msk2_dir, fname) + if base not in by_base: + by_base[base] = cand + else: + cur_ext = os.path.splitext(by_base[base])[1].lower() + if cur_ext != ".png" and ext.lower() == ".png": + by_base[base] = cand + return by_base + +def ensure_aug_dirs(group_name=None, use_masks2=False): + """Cria diretórios de saída para o grupo ou modo antigo. Se use_masks2, cria masks2.""" if group_name: img_out = os.path.join(AUG_GROUP_ROOT, group_name, "images") msk_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks") + msk2_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks2") if use_masks2 else None else: img_out = AUG_OLD_IMG msk_out = AUG_OLD_MSK + msk2_out = AUG_OLD_MSK2 if use_masks2 else None garantir_dir(img_out) garantir_dir(msk_out) - return img_out, msk_out + if use_masks2 and msk2_out: + garantir_dir(msk2_out) + return img_out, msk_out, msk2_out -def augment_pair(img_path, msk_path, img_out_dir, msk_out_dir, copies): +def augment_pair(img_path, msk_path, img_out_dir, msk_out_dir, copies, msk2_path=None, msk2_out_dir=None): base_img, img_ext = os.path.splitext(os.path.basename(img_path)) base_msk, msk_ext = os.path.splitext(os.path.basename(msk_path)) + msk2_ext = os.path.splitext(os.path.basename(msk2_path))[1] if msk2_path else None # padroniza pelo base da imagem base = base_img img = load_rgb(img_path) msk = load_rgb(msk_path) + msk2 = load_rgb(msk2_path) if msk2_path else None gen = 0 for i in range(copies): - aug = train_tf(image=img, mask=msk) + if msk2 is not None and msk2_out_dir: + aug = train_tf(image=img, mask=msk, mask2=msk2) + else: + aug = train_tf(image=img, mask=msk) + img_aug = aug["image"] msk_aug = aug["mask"] @@ -163,6 +198,12 @@ def augment_pair(img_path, msk_path, img_out_dir, msk_out_dir, copies): out_msk = os.path.join(msk_out_dir, f"{base}_aug_{i:02d}{msk_ext}") save_rgb(out_img, img_aug) save_rgb(out_msk, msk_aug) + + if msk2 is not None and msk2_out_dir: + msk2_aug = aug["mask2"] + out_msk2 = os.path.join(msk2_out_dir, f"{base}_aug_{i:02d}{msk2_ext}") + save_rgb(out_msk2, msk2_aug) + gen += 1 return gen @@ -170,13 +211,16 @@ def process_group(group_name, copies): """Processa um grupo único (images/masks dentro de ORIG_GROUP_ROOT//).""" img_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "images") msk_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks") + msk2_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks2") if not (os.path.isdir(img_dir) and os.path.isdir(msk_dir)): print(f"[WARN] Grupo '{group_name}' inválido (sem images/masks). Pulando.") return 0 imgs = [f for f in os.listdir(img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS] msk_map = map_masks_by_base(msk_dir) - img_out_dir, msk_out_dir = ensure_aug_dirs(group_name) + use_masks2 = USE_MASKS2 and os.path.isdir(msk2_dir) + msk2_map = map_masks2_by_base(msk2_dir) if use_masks2 else {} + img_out_dir, msk_out_dir, msk2_out_dir = ensure_aug_dirs(group_name, use_masks2=use_masks2) count = 0 for img_file in sorted(imgs): @@ -185,13 +229,18 @@ def process_group(group_name, copies): if not msk_file: print(f"[WARN] [{group_name}] Máscara não encontrada para {img_file}, pulando.") continue + msk2_file = msk2_map.get(base) if use_masks2 else None + if use_masks2 and not msk2_file: + print(f"[WARN] [{group_name}] mask2 não encontrada para {img_file}, gerando só img+mask.") try: count += augment_pair( os.path.join(img_dir, img_file), msk_file, img_out_dir, msk_out_dir, - copies=copies + copies=copies, + msk2_path=msk2_file, + msk2_out_dir=msk2_out_dir ) except Exception as e: print(f"[ERRO] [{group_name}] {img_file}: {e}") @@ -206,7 +255,9 @@ def process_legacy(copies): imgs = [f for f in os.listdir(ORIG_OLD_IMG) if os.path.splitext(f.lower())[1] in IMG_EXTS] msk_map = map_masks_by_base(ORIG_OLD_MSK) - img_out_dir, msk_out_dir = ensure_aug_dirs(group_name=None) + use_masks2 = USE_MASKS2 and os.path.isdir(ORIG_OLD_MSK2) + msk2_map = map_masks2_by_base(ORIG_OLD_MSK2) if use_masks2 else {} + img_out_dir, msk_out_dir, msk2_out_dir = ensure_aug_dirs(group_name=None, use_masks2=use_masks2) count = 0 for img_file in sorted(imgs): @@ -215,13 +266,18 @@ def process_legacy(copies): if not msk_file: print(f"[WARN] (legacy) Máscara não encontrada para {img_file}, pulando.") continue + msk2_file = msk2_map.get(base) if use_masks2 else None + if use_masks2 and not msk2_file: + print(f"[WARN] (legacy) mask2 não encontrada para {img_file}, gerando só img+mask.") try: count += augment_pair( os.path.join(ORIG_OLD_IMG, img_file), msk_file, img_out_dir, msk_out_dir, - copies=copies + copies=copies, + msk2_path=msk2_file, + msk2_out_dir=msk2_out_dir ) except Exception as e: print(f"[ERRO] (legacy) {img_file}: {e}") diff --git a/Python/OAK/datasets/_6_normalize.py b/Python/OAK/datasets/_6_normalize.py index 549e17186..348f64123 100644 --- a/Python/OAK/datasets/_6_normalize.py +++ b/Python/OAK/datasets/_6_normalize.py @@ -28,6 +28,7 @@ from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids with open("config.json", "r", encoding="utf-8") as f: config = json.load(f) MODELO = config["camera"] +USE_MASKS2 = config["dual_head"] RESOLUCAO = tuple(config["resolucao"]) # [W, H] ou [width, height] pasta_base = os.path.join(MODELO, "dataset") labelmap_path = os.path.join(pasta_base, "labelmap.txt") @@ -43,6 +44,7 @@ FONTES = ["original", "augmented"] # Extensões aceitas IMG_EXTS = (".jpg", ".jpeg", ".png") MSK_EXTS = (".png", ".jpg", ".jpeg") # preferir .png +MSK2_EXTS = (".png", ".jpg", ".jpeg") # idem def infer_ignore_id(ignore_rgb, default_id=255): """ @@ -98,6 +100,25 @@ def map_masks_by_base(msk_dir: str) -> Dict[str, str]: by_base[base] = cand return by_base +def map_masks2_by_base(msk2_dir: str) -> Dict[str, str]: + """Retorna {base: caminho_mask2}, priorizando .png quando houver múltiplas por base.""" + by_base = {} + if not os.path.isdir(msk2_dir): + return by_base + for fname in os.listdir(msk2_dir): + f_lower = fname.lower() + if not f_lower.endswith(MSK2_EXTS): + continue + base, ext = os.path.splitext(fname) + cand = os.path.join(msk2_dir, fname) + if base not in by_base: + by_base[base] = cand + else: + cur_ext = os.path.splitext(by_base[base])[1].lower() + if cur_ext != ".png" and ext.lower() == ".png": + by_base[base] = cand + return by_base + def normalize_pair(caminho_rgb: str, caminho_mask: str, cor_para_id, ignore_id: int, out_img_dir: str, out_msk_dir: str, dim: Tuple[int,int], prefix: str = ""): """Redimensiona e grava a imagem e a máscara (se houver).""" @@ -137,6 +158,43 @@ def normalize_pair(caminho_rgb: str, caminho_mask: str, cor_para_id, ignore_id: return True +def normalize_pair_mask2(caminho_rgb: str, caminho_mask2: str, + out_msk2_dir: str, dim: Tuple[int,int], prefix: str = ""): + """ + Redimensiona e grava máscara2 (corredor binário), assumindo que ela já é uma máscara "pronta". + - Se for RGB/BGR (3 canais): converte para cinza e faz threshold (0/255) antes de redimensionar. + - Se for 1 canal: mantém, faz threshold (0/255). + - Redimensiona com INTER_NEAREST. + Saída sempre .png com o mesmo nome base/prefixo do arquivo de imagem. + """ + if not caminho_mask2 or not os.path.isfile(caminho_mask2): + return False + + nome = os.path.basename(caminho_rgb) + nome_saida = f"{prefix}{nome}" if prefix else nome + for ext in (".jpg", ".jpeg", ".png"): + if nome_saida.lower().endswith(ext): + nome_saida = nome_saida[: -len(ext)] + ".png" + break + + m2 = cv2.imread(caminho_mask2, cv2.IMREAD_UNCHANGED) + if m2 is None: + print(f"[!] Erro ao ler máscara2: {caminho_mask2}") + return False + + if len(m2.shape) == 3: + # BGR/RGB -> gray + m2g = cv2.cvtColor(m2, cv2.COLOR_BGR2GRAY) + else: + m2g = m2 + + # binariza para 0/255 (evita lixo de compressão) + _, m2bin = cv2.threshold(m2g, 127, 255, cv2.THRESH_BINARY) + m2res = cv2.resize(m2bin, dim, interpolation=cv2.INTER_NEAREST) + garantir_dir(out_msk2_dir) + cv2.imwrite(os.path.join(out_msk2_dir, nome_saida), m2res) + return True + def process_group_root(fonte_root: str, fonte_nome: str, cor_para_id, ignore_id: int, groups_except: str = None): """Processa uma raiz do tipo ...//group/ agrupando por cada subpasta de grupo.""" total = 0 @@ -155,13 +213,17 @@ def process_group_root(fonte_root: str, fonte_nome: str, cor_para_id, ignore_id: continue in_img_dir = os.path.join(fonte_root, grupo, "images") in_msk_dir = os.path.join(fonte_root, grupo, "masks") + in_msk2_dir = os.path.join(fonte_root, grupo, "masks2") if not (os.path.isdir(in_img_dir) and os.path.isdir(in_msk_dir)): print(f"[WARN] Grupo inválido (sem images/masks): {grupo}") continue out_img_dir = os.path.join(out_root, grupo, "images") out_msk_dir = os.path.join(out_root, grupo, "masks") + usar_masks2 = USE_MASKS2 and os.path.isdir(in_msk2_dir) + out_msk2_dir = os.path.join(out_root, grupo, "masks2") if usar_masks2 else None msk_map = map_masks_by_base(in_msk_dir) + msk2_map = map_masks2_by_base(in_msk2_dir) if usar_masks2 else {} imgs = [f for f in os.listdir(in_img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS] n = len(imgs) @@ -169,10 +231,19 @@ def process_group_root(fonte_root: str, fonte_nome: str, cor_para_id, ignore_id: base, _ = os.path.splitext(fname) caminho_rgb = os.path.join(in_img_dir, fname) caminho_mask = msk_map.get(base) + caminho_mask2 = msk2_map.get(base) if usar_masks2 else None ok = normalize_pair( caminho_rgb, caminho_mask, cor_para_id, ignore_id, out_img_dir, out_msk_dir, dim, prefix=f"{fonte_nome}_" ) + if usar_masks2 and out_msk2_dir: + if not caminho_mask2: + print(f"[WARN] [{fonte_nome} | {grupo}] masks2 existe, mas não achei mask2 p/ {fname} (vou seguir).") + else: + normalize_pair_mask2( + caminho_rgb, caminho_mask2, + out_msk2_dir, dim, prefix=f"{fonte_nome}_" + ) if ok: total += 1 print(f"[{fonte_nome} | {grupo} | {nome_res}] {i}/{n} → {fname}") @@ -187,18 +258,31 @@ def process_legacy_root(legacy_img: str, legacy_msk: str, fonte_nome: str, cor_p for nome_res, dim in RESOLUCOES.items(): out_img_dir = os.path.join(pasta_base, nome_res, "images") out_msk_dir = os.path.join(pasta_base, nome_res, "masks") + legacy_msk2 = os.path.join(os.path.dirname(legacy_msk), "masks2") + usar_masks2 = USE_MASKS2 and os.path.isdir(legacy_msk2) + out_msk2_dir = os.path.join(pasta_base, nome_res, "masks2") if usar_masks2 else None msk_map = map_masks_by_base(legacy_msk) + msk2_map = map_masks2_by_base(legacy_msk2) if usar_masks2 else {} imgs = [f for f in os.listdir(legacy_img) if os.path.splitext(f.lower())[1] in IMG_EXTS] n = len(imgs) for i, fname in enumerate(sorted(imgs), 1): base, _ = os.path.splitext(fname) caminho_rgb = os.path.join(legacy_img, fname) caminho_mask = msk_map.get(base) + caminho_mask2 = msk2_map.get(base) if usar_masks2 else None ok = normalize_pair( caminho_rgb, caminho_mask, cor_para_id, ignore_id, out_img_dir, out_msk_dir, dim, prefix=f"{fonte_nome}_" ) + if usar_masks2 and out_msk2_dir: + if not caminho_mask2: + print(f"[WARN] [{fonte_nome} | legacy] masks2 existe, mas não achei mask2 p/ {fname} (vou seguir).") + else: + normalize_pair_mask2( + caminho_rgb, caminho_mask2, + out_msk2_dir, dim, prefix=f"{fonte_nome}_" + ) if ok: total += 1 print(f"[{fonte_nome} | legacy | {nome_res}] {i}/{n} → {fname}") @@ -209,6 +293,7 @@ def main(args): # Espera tupla na ordem: (cor_para_id, colormap_rgb, id_para_nome, ignore_rgb) cor_para_id, _colormap_rgb, _id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path) ignore_id = infer_ignore_id(ignore_rgb, default_id=255) + print(cor_para_id, _colormap_rgb, _id_para_nome) total_geral = 0 # === ORIGINAL === diff --git a/Python/OAK/datasets/_7_split.py b/Python/OAK/datasets/_7_split.py index 2c94590ba..61937183d 100644 --- a/Python/OAK/datasets/_7_split.py +++ b/Python/OAK/datasets/_7_split.py @@ -36,6 +36,7 @@ import argparse with open("config.json", "r", encoding="utf-8") as f: config = json.load(f) MODELO = config.get("camera") +USE_MASKS2 = config.get("dual_head", False) RESOLUCAO = tuple(config.get("resolucao")) # Pastas @@ -77,6 +78,10 @@ def mask_from_image_name(img_name): base, _ = os.path.splitext(img_name) return base + MSK_EXT +def mask2_from_image_name(img_name): + base, _ = os.path.splitext(img_name) + return base + MSK_EXT # masks2 também normalizadas em PNG no normalize + def classify_source_and_family(filename_no_ext): """ Retorna (source, family_key) @@ -193,8 +198,11 @@ def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test): return n_train, n_val, n_test -def copiar(nomes, src_img_dir, src_msk_dir, dst_img_dir, dst_msk_dir): +def copiar(nomes, src_img_dir, src_msk_dir, dst_img_dir, dst_msk_dir, src_msk2_dir=None, dst_msk2_dir=None): garantir(dst_img_dir); garantir(dst_msk_dir) + use_msk2 = bool(src_msk2_dir and dst_msk2_dir and os.path.isdir(src_msk2_dir)) + if use_msk2: + garantir(dst_msk2_dir) moved = 0 for nome in nomes: mask_name = mask_from_image_name(nome) @@ -204,12 +212,19 @@ def copiar(nomes, src_img_dir, src_msk_dir, dst_img_dir, dst_msk_dir): continue shutil.copy2(src_img, os.path.join(dst_img_dir, nome)) shutil.copy2(src_msk, os.path.join(dst_msk_dir, mask_name)) + if use_msk2: + m2_name = mask2_from_image_name(nome) + src_m2 = os.path.join(src_msk2_dir, m2_name) + if os.path.exists(src_m2): + shutil.copy2(src_m2, os.path.join(dst_msk2_dir, m2_name)) moved += 1 return moved def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None): src_img_dir = os.path.join(pasta_origem, group_name, "images") src_msk_dir = os.path.join(pasta_origem, group_name, "masks") + src_msk2_dir = os.path.join(pasta_origem, group_name, "masks2") + use_msk2 = USE_MASKS2 and os.path.isdir(src_msk2_dir) familias = build_family_index(src_img_dir, src_msk_dir) # apenas famílias que têm ORIGINAL para participar de val/test @@ -272,9 +287,13 @@ def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None): dest_test_img = os.path.join(pasta_destino, "test", "group", group_name, "images") dest_test_msk = os.path.join(pasta_destino, "test", "group", group_name, "masks") - m_train = copiar(nomes_train, src_img_dir, src_msk_dir, dest_train_img, dest_train_msk) - m_val = copiar(nomes_val, src_img_dir, src_msk_dir, dest_val_img, dest_val_msk) - m_test = copiar(nomes_test, src_img_dir, src_msk_dir, dest_test_img, dest_test_msk) + dest_train_msk2 = os.path.join(pasta_destino, "train", "group", group_name, "masks2") if use_msk2 else None + dest_val_msk2 = os.path.join(pasta_destino, "val", "group", group_name, "masks2") if use_msk2 else None + dest_test_msk2 = os.path.join(pasta_destino, "test", "group", group_name, "masks2") if use_msk2 else None + + m_train = copiar(nomes_train, src_img_dir, src_msk_dir, dest_train_img, dest_train_msk, src_msk2_dir, dest_train_msk2) + m_val = copiar(nomes_val, src_img_dir, src_msk_dir, dest_val_img, dest_val_msk, src_msk2_dir, dest_val_msk2) + m_test = copiar(nomes_test, src_img_dir, src_msk_dir, dest_test_img, dest_test_msk, src_msk2_dir, dest_test_msk2) print(f"[{group_name}] famílias={total_familias} → train(imgs)={m_train}, val(imgs)={m_val}, test(imgs)={m_test}") return {"train": m_train, "val": m_val, "test": m_test, "familias": total_familias} diff --git a/Python/OAK/datasets/_8_train_segformer_b3.py b/Python/OAK/datasets/_8_train_segformer_b3.py new file mode 100644 index 000000000..94cdd18bf --- /dev/null +++ b/Python/OAK/datasets/_8_train_segformer_b3.py @@ -0,0 +1,481 @@ +#python _8_train_segformer_b3.py --epochs 110 --batch 2 --lr 3e-5 --wd 0.01 --num_workers 4 --amp --amp_val --grad_accum 2 --class_weights auto --main_class navegavel --resume + +# _8_train_segformer_b3.py (PATCH) +import os + +# (Opcional) ajuda com fragmentação em algumas máquinas. +# Idealmente isso deveria vir ANTES de importar torch, mas já ajuda quando setado fora também. +os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128") + +import json +import math +import time +import argparse +from typing import Dict, List, Optional, Tuple, Any + +import numpy as np +import torch +import torch.nn as nn +from torch.utils.data import DataLoader + +# >>> troque AMP deprecated +from torch.amp import autocast, GradScaler + +from transformers import SegformerForSemanticSegmentation + +# >>>>>>>>> AJUSTE AQUI <<<<<<<<< +from _8_train_fastscnn_v2 import ROISegDataset # troque se necessário + + +def set_seed(seed: int = 42): + import random + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _ids_by_names(class_names: Any, wanted_names: List[str]) -> List[int]: + if class_names is None: + return [] + wanted_names = [str(x).lower() for x in wanted_names] + + if isinstance(class_names, list): + low = [c.lower() for c in class_names] + return [low.index(n) for n in wanted_names if n in low] + + if isinstance(class_names, dict): + if all(isinstance(k, int) for k in class_names.keys()): + inv = {str(v).lower(): int(k) for k, v in class_names.items()} + return [inv[n] for n in wanted_names if n in inv] + inv = {str(k).lower(): int(v) for k, v in class_names.items()} + return [inv[n] for n in wanted_names if n in inv] + + return [] + + +@torch.no_grad() +def update_confusion_matrix(cm: torch.Tensor, preds: torch.Tensor, labels: torch.Tensor, num_classes: int, ignore_index: int = 255): + preds = preds.view(-1) + labels = labels.view(-1) + + valid = labels != ignore_index + preds = preds[valid] + labels = labels[valid] + + idx = labels * num_classes + preds + bins = torch.bincount(idx, minlength=num_classes * num_classes) + cm += bins.view(num_classes, num_classes) + + +@torch.no_grad() +def compute_iou_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> Tuple[float, List[float]]: + cm = cm.float() + tp = torch.diag(cm) + fp = cm.sum(0) - tp + fn = cm.sum(1) - tp + denom = tp + fp + fn + eps + iou = (tp / denom).cpu().tolist() + miou = float(np.mean(iou)) + return miou, iou + + +@torch.no_grad() +def compute_pixel_acc_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> float: + cm = cm.float() + acc = (torch.diag(cm).sum() / (cm.sum() + eps)).item() + return acc + + +IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) +IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) +def normalize_img(img: torch.Tensor) -> torch.Tensor: + return (img - IMAGENET_MEAN.to(img.device)) / IMAGENET_STD.to(img.device) + + +def default_collate(batch): + imgs = [] + masks = [] + for item in batch: + if isinstance(item, dict): + img = item["image"] + mask = item["mask"] + else: + img, mask = item + + if isinstance(img, np.ndarray): + img = torch.from_numpy(img) + if isinstance(mask, np.ndarray): + mask = torch.from_numpy(mask) + + if img.ndim == 3 and img.shape[-1] == 3: + img = img.permute(2, 0, 1) + + if img.dtype != torch.float32: + img = img.float() + if img.max() > 1.5: + img = img / 255.0 + + imgs.append(img) + masks.append(mask.long()) + + return torch.stack(imgs, 0), torch.stack(masks, 0) + + +def estimate_class_weights(ds: ROISegDataset, num_classes: int, ignore_index: int = 255, max_samples: int = 800) -> torch.Tensor: + n = min(len(ds), max_samples) + idxs = np.random.choice(len(ds), size=n, replace=False) + + counts = np.zeros(num_classes, dtype=np.float64) + for i in idxs: + item = ds[i] + mask = item["mask"] if isinstance(item, dict) else item[1] + m = mask.cpu().numpy() if isinstance(mask, torch.Tensor) else np.array(mask) + + m = m.reshape(-1) + m = m[m != ignore_index] + if m.size == 0: + continue + + counts += np.bincount(m, minlength=num_classes)[:num_classes] + + freq = counts / (counts.sum() + 1e-12) + freq = np.clip(freq, 1e-12, 1.0) + weights = 1.0 / np.log(1.02 + freq) + weights = weights / weights.mean() + return torch.tensor(weights, dtype=torch.float32) + + +def save_checkpoint(path: str, + model: nn.Module, + optimizer: torch.optim.Optimizer, + scaler: Optional[GradScaler], + epoch: int, + best_miou: float, + best_main_iou: float, + extra: Optional[Dict[str, Any]] = None): + ckpt = { + "epoch": epoch, + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "best_miou": best_miou, + "best_main_iou": best_main_iou, + } + if scaler is not None: + ckpt["scaler"] = scaler.state_dict() + if extra: + ckpt["extra"] = extra + torch.save(ckpt, path) + + +def load_checkpoint(path: str, + model: nn.Module, + optimizer: Optional[torch.optim.Optimizer] = None, + scaler: Optional[GradScaler] = None, + map_location: str = "cpu") -> Dict[str, Any]: + ckpt = torch.load(path, map_location=map_location, weights_only=False) + model.load_state_dict(ckpt["model"], strict=True) + if optimizer is not None and "optimizer" in ckpt: + optimizer.load_state_dict(ckpt["optimizer"]) + if scaler is not None and "scaler" in ckpt: + scaler.load_state_dict(ckpt["scaler"]) + return ckpt + + +def run_one_epoch(model: nn.Module, + loader: DataLoader, + optimizer: Optional[torch.optim.Optimizer], + device: torch.device, + num_classes: int, + ignore_index: int, + criterion: nn.Module, + amp: bool, + scaler: Optional[GradScaler], + train: bool, + grad_accum: int = 1) -> Dict[str, Any]: + + model.train(train) + + total_loss = 0.0 + cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device) + + t0 = time.time() + n_batches = 0 + + # >>>>> ESSENCIAL: desliga grad no val + with torch.set_grad_enabled(train): + if train and optimizer is not None: + optimizer.zero_grad(set_to_none=True) + + for step, (imgs, masks) in enumerate(loader): + imgs = imgs.to(device, non_blocking=True) + masks = masks.to(device, non_blocking=True) + + imgs = normalize_img(imgs) + + # AMP tanto em train quanto (opcionalmente) em val + with autocast(device_type="cuda", enabled=amp and device.type == "cuda"): + out = model(pixel_values=imgs) + logits = out.logits + + if logits.shape[-2:] != masks.shape[-2:]: + logits = torch.nn.functional.interpolate( + logits, size=masks.shape[-2:], mode="bilinear", align_corners=False + ) + + loss = criterion(logits, masks) + if train and grad_accum > 1: + loss = loss / grad_accum + + if train and optimizer is not None: + if amp and scaler is not None and device.type == "cuda": + scaler.scale(loss).backward() + if ((step + 1) % grad_accum) == 0: + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + else: + loss.backward() + if ((step + 1) % grad_accum) == 0: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + total_loss += float(loss.item()) * (grad_accum if train and grad_accum > 1 else 1.0) + n_batches += 1 + + preds = torch.argmax(logits, dim=1) + update_confusion_matrix(cm, preds, masks, num_classes=num_classes, ignore_index=ignore_index) + + dt = time.time() - t0 + avg_loss = total_loss / max(n_batches, 1) + miou, iou_per_class = compute_iou_from_cm(cm) + acc = compute_pixel_acc_from_cm(cm) + + return {"loss": avg_loss, "miou": miou, "iou_per_class": iou_per_class, "acc": acc, "time_s": dt} + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--config", default="config.json") + parser.add_argument("--epochs", type=int, default=120) + parser.add_argument("--batch", type=int, default=1) # <<< default seguro pra 8GB + parser.add_argument("--lr", type=float, default=6e-5) + parser.add_argument("--wd", type=float, default=0.01) + parser.add_argument("--num_workers", type=int, default=4) + + parser.add_argument("--save_every", type=int, default=10) + parser.add_argument("--resume", action="store_true") + + parser.add_argument("--amp", action="store_true") + parser.add_argument("--amp_val", action="store_true") # <<< novo: AMP no val + parser.add_argument("--grad_accum", type=int, default=1) # <<< novo + + parser.add_argument("--grad_ckpt", action="store_true") # <<< novo: checkpointing + parser.add_argument("--ignore_index", type=int, default=255) + parser.add_argument("--class_weights", type=str, default="auto") + parser.add_argument("--main_class", type=str, default=None) + parser.add_argument("--es_classes", type=str, default="") + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + set_seed(args.seed) + + with open(args.config, "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"] + MAIN_CLASS_NAME = str(config.get("main_class_name", "erva")).lower() + BACKBONE = config["backbone"] + if args.main_class is not None: + MAIN_CLASS_NAME = args.main_class.lower() + + dataset_path = os.path.join(MODELO, "dataset") + labelmap_path = os.path.join(dataset_path, "labelmap.txt") + + save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME) + os.makedirs(save_path, exist_ok=True) + last_ckpt_path = os.path.join(save_path, "last.pt") + best_miou_path = os.path.join(save_path, "best_miou.pt") + best_main_path = os.path.join(save_path, "best_main.pt") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Device: {device}") + + # Datasets + ds_train = ROISegDataset( + os.path.join(dataset_path, "split", "train"), + save_path, ROI_INICIO, ROI_TAMANHO, + RESOLUCAO[0], RESOLUCAO[1], labelmap_path + ) + ds_val = ROISegDataset( + os.path.join(dataset_path, "split", "val"), + save_path, ROI_INICIO, ROI_TAMANHO, + RESOLUCAO[0], RESOLUCAO[1], labelmap_path + ) + + CLASS_NAMES = getattr(ds_train, "classes", None) + if CLASS_NAMES is None: + raise RuntimeError("ROISegDataset precisa expor .classes (list ou dict).") + + if isinstance(CLASS_NAMES, list): + num_classes = len(CLASS_NAMES) + class_id_by_name = {n.lower(): i for i, n in enumerate(CLASS_NAMES)} + class_name_by_id = {i: n for i, n in enumerate(CLASS_NAMES)} + elif isinstance(CLASS_NAMES, dict): + if all(isinstance(k, int) for k in CLASS_NAMES.keys()): + num_classes = len(CLASS_NAMES) + class_name_by_id = {int(k): str(v) for k, v in CLASS_NAMES.items()} + class_id_by_name = {str(v).lower(): int(k) for k, v in CLASS_NAMES.items()} + else: + class_id_by_name = {str(k).lower(): int(v) for k, v in CLASS_NAMES.items()} + num_classes = len(class_id_by_name) + class_name_by_id = {v: k for k, v in class_id_by_name.items()} + else: + raise RuntimeError("Formato de .classes não reconhecido.") + + main_class_id = class_id_by_name.get(MAIN_CLASS_NAME, None) + if main_class_id is None: + print(f"[WARN] main_class_name='{MAIN_CLASS_NAME}' não encontrado. best_main_iou usa mIoU.") + else: + print(f"Main class: '{MAIN_CLASS_NAME}' -> id={main_class_id}") + + es_names = [x.strip().lower() for x in args.es_classes.split(",") if x.strip()] + ES_CLASS_IDS = _ids_by_names(CLASS_NAMES, es_names) + + dl_train = DataLoader( + ds_train, batch_size=args.batch, shuffle=True, + num_workers=args.num_workers, pin_memory=True, + collate_fn=default_collate, drop_last=True + ) + dl_val = DataLoader( + ds_val, batch_size=1, shuffle=False, # <<< val com batch 1 é mais estável + num_workers=max(2, args.num_workers // 2), pin_memory=True, + collate_fn=default_collate, drop_last=False + ) + + model = SegformerForSemanticSegmentation.from_pretrained( + BACKBONE, + num_labels=num_classes, + ignore_mismatched_sizes=True, + use_safetensors=True + ) + + if args.grad_ckpt: + try: + # Nem toda versão/classe suporta + model.gradient_checkpointing_enable() + print("[OK] gradient checkpointing enabled") + except Exception as e: + print(f"[WARN] gradient checkpointing não suportado aqui ({type(e).__name__}: {e}). Seguindo sem.") + + model.to(device) + + # Loss weights + if args.class_weights.lower() == "none": + weights = None + elif args.class_weights.lower() == "auto": + w = estimate_class_weights(ds_train, num_classes=num_classes, ignore_index=args.ignore_index) + weights = w.to(device) + print("Class weights (auto):", w.cpu().numpy().round(3).tolist()) + else: + parts = [float(x) for x in args.class_weights.split(",")] + if len(parts) != num_classes: + raise ValueError(f"class_weights manual precisa ter {num_classes} valores, recebeu {len(parts)}.") + weights = torch.tensor(parts, dtype=torch.float32, device=device) + + criterion = nn.CrossEntropyLoss(weight=weights, ignore_index=args.ignore_index) + + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd) + + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode="min", factor=0.5, patience=6, threshold=1e-4, verbose=True + ) + + scaler = GradScaler(enabled=args.amp and device.type == "cuda") + + start_epoch = 1 + best_miou = -1.0 + best_main_iou = -1.0 + + if args.resume and os.path.exists(best_miou_path): + ckpt = load_checkpoint(best_miou_path, model, optimizer, scaler=scaler, map_location="cpu") + start_epoch = int(ckpt["epoch"]) + 1 + best_miou = float(ckpt.get("best_miou", -1.0)) + best_main_iou = float(ckpt.get("best_main_iou", -1.0)) + print(f"[RESUME] epoch={start_epoch} best_miou={best_miou:.4f} best_main_iou={best_main_iou:.4f}") + + def pretty_iou(iou_list): + return " | ".join([f"{class_name_by_id.get(cid,cid)}:{v:.3f}" for cid, v in enumerate(iou_list)]) + + for epoch in range(start_epoch, args.epochs + 1): + lr_now = optimizer.param_groups[0]["lr"] + print(f"\n==== Epoch {epoch}/{args.epochs} | lr={lr_now:.2e} ====") + + if device.type == "cuda": + torch.cuda.empty_cache() + + tr = run_one_epoch( + model=model, loader=dl_train, optimizer=optimizer, + device=device, num_classes=num_classes, ignore_index=args.ignore_index, + criterion=criterion, amp=args.amp, scaler=scaler, train=True, + grad_accum=max(1, args.grad_accum) + ) + + if device.type == "cuda": + torch.cuda.empty_cache() + + va = run_one_epoch( + model=model, loader=dl_val, optimizer=None, + device=device, num_classes=num_classes, ignore_index=args.ignore_index, + criterion=criterion, amp=args.amp_val, scaler=None, train=False, + grad_accum=1 + ) + + scheduler.step(va["loss"]) + + main_iou = va["miou"] if main_class_id is None else va["iou_per_class"][main_class_id] + + print(f"TRAIN: loss={tr['loss']:.4f} acc={tr['acc']:.4f} miou={tr['miou']:.4f} (t={tr['time_s']:.1f}s)") + print(f"VAL : loss={va['loss']:.4f} acc={va['acc']:.4f} miou={va['miou']:.4f} main_iou={main_iou:.4f} (t={va['time_s']:.1f}s)") + print("IoU per class:", pretty_iou(va["iou_per_class"])) + + if ES_CLASS_IDS: + mini = " | ".join([f"{class_name_by_id[cid]}:{va['iou_per_class'][cid]:.3f}" for cid in ES_CLASS_IDS]) + print("ES classes:", mini) + + save_checkpoint( + last_ckpt_path, model, optimizer, scaler=scaler, + epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, + extra={"val_loss": va["loss"], "val_miou": va["miou"], "val_main_iou": float(main_iou)} + ) + + if args.save_every > 0 and (epoch % args.save_every == 0): + save_checkpoint( + os.path.join(save_path, f"epoch_{epoch:04d}.pt"), + model, optimizer, scaler=scaler, + epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou + ) + + if va["miou"] > best_miou: + best_miou = va["miou"] + save_checkpoint( + best_miou_path, model, optimizer, scaler=scaler, + epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou + ) + print(f"[BEST mIoU] {best_miou:.4f} -> saved: {best_miou_path}") + + if float(main_iou) > best_main_iou: + best_main_iou = float(main_iou) + save_checkpoint( + best_main_path, model, optimizer, scaler=scaler, + epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou + ) + print(f"[BEST MAIN] {best_main_iou:.4f} -> saved: {best_main_path}") + + +if __name__ == "__main__": + main() diff --git a/Python/OAK/datasets/_8_train_segformer_b3_dual.py b/Python/OAK/datasets/_8_train_segformer_b3_dual.py new file mode 100644 index 000000000..376a3b2d2 --- /dev/null +++ b/Python/OAK/datasets/_8_train_segformer_b3_dual.py @@ -0,0 +1,965 @@ +# _8_train_segformer_b3_dual_v2.py +# Treino SegFormer-B3 com 2 cabeças: +# - Head 1: segmentação semântica (3 classes + ignore) +# - Head 2: corredor binário (0/1 + ignore) +# +# Features: +# - class weights (opcional, estimado do train) +# - pos_weight (opcional, estimado do train p/ BCE do corredor) +# - warmup + cosine (opcional) + ReduceLROnPlateau +# - early stopping com métrica agregada +# - resume completo (model + corridor_head + optimizer + scaler) +# +# Observação: este script assume que seu ROISegDataset retorna dict com: +# {"image": , "mask": , ...} +# e que você já acertou para também expor "image_path"/"mask_path" OU +# que seu DualMaskROISegDataset consiga derivar paths corretamente. + +#python _8_train_segformer_b3_dual.py --epochs 80 --batch 2 --num_workers 2 --amp --amp_val --use_weights --auto_pos_weight --warmup_epochs 2 --cosine_epochs 15 --lambda_corr 0.5 --corr_dice_max 0.10 --corr_dice_ramp_epochs 10 --es_metric harmonic --es_patience 12 + +#python _8_train_segformer_b3_dual.py --epochs 120 --batch 2 --num_workers 4 --amp --amp_val --grad_accum 1 --use_weights --auto_pos_weight --warmup_epochs 2 --cosine_epochs 20 --lambda_corr 0.7 --corr_dice_max 0.15 --corr_dice_ramp_epochs 10 --es_metric harmonic --es_patience 15 + +import os +import re +import time +import math +import json +import random +import argparse +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +from PIL import Image + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader + +from torch.amp import autocast +from torch.cuda.amp import GradScaler + +# HuggingFace transformers +from transformers import SegformerForSemanticSegmentation + +# Seu dataset (o mesmo que você já usa no projeto) +from roi_seg_dataset import ROISegDataset + + +# ----------------------------- +# Utils +# ----------------------------- +def seed_everything(seed: int = 42): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) +IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) + + +def normalize_img(img: torch.Tensor) -> torch.Tensor: + return (img - IMAGENET_MEAN.to(img.device)) / IMAGENET_STD.to(img.device) + + +@torch.no_grad() +def update_confusion_matrix(cm: torch.Tensor, preds: torch.Tensor, labels: torch.Tensor, num_classes: int, ignore_index: int = 255): + preds = preds.view(-1) + labels = labels.view(-1) + + valid = labels != ignore_index + preds = preds[valid] + labels = labels[valid] + + idx = labels * num_classes + preds + bins = torch.bincount(idx, minlength=num_classes * num_classes) + cm += bins.view(num_classes, num_classes) + + +@torch.no_grad() +def compute_iou_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> Tuple[float, List[float]]: + cm = cm.float() + tp = torch.diag(cm) + fp = cm.sum(0) - tp + fn = cm.sum(1) - tp + denom = tp + fp + fn + eps + iou = (tp / denom).cpu().tolist() + miou = float(np.mean(iou)) + return miou, iou + + +@torch.no_grad() +def compute_pixel_acc_from_cm(cm: torch.Tensor, eps: float = 1e-6) -> float: + cm = cm.float() + acc = (torch.diag(cm).sum() / (cm.sum() + eps)).item() + return acc + + +def default_collate_dual(batch): + imgs, m1s, m2s = [], [], [] + for item in batch: + # item pode vir como (img, mask1, mask2) ou dict etc. + if isinstance(item, dict): + img = item["image"] + mask1 = item["mask1"] + mask2 = item["mask2"] + else: + img, mask1, mask2 = item + + if isinstance(img, np.ndarray): + img = torch.from_numpy(img) + if isinstance(mask1, np.ndarray): + mask1 = torch.from_numpy(mask1) + if isinstance(mask2, np.ndarray): + mask2 = torch.from_numpy(mask2) + + if img.ndim == 3 and img.shape[-1] == 3: + img = img.permute(2, 0, 1) + + if img.dtype != torch.float32: + img = img.float() + if img.max() > 1.5: + img = img / 255.0 + + imgs.append(img) + m1s.append(mask1.long()) + m2s.append(mask2.long()) + + return torch.stack(imgs, 0), torch.stack(m1s, 0), torch.stack(m2s, 0) + + +def estimate_class_weights_from_ds(ds: ROISegDataset, num_classes: int, ignore_index: int = 255, max_samples: int = 800) -> torch.Tensor: + n = min(len(ds), max_samples) + idxs = np.random.choice(len(ds), size=n, replace=False) + + counts = np.zeros(num_classes, dtype=np.float64) + for i in idxs: + item = ds[i] + mask = item["mask"] if isinstance(item, dict) else item[1] + m = mask.cpu().numpy() if isinstance(mask, torch.Tensor) else np.array(mask) + + m = m.reshape(-1) + m = m[m != ignore_index] + if m.size == 0: + continue + counts += np.bincount(m, minlength=num_classes)[:num_classes] + + freq = counts / (counts.sum() + 1e-12) + freq = np.clip(freq, 1e-12, 1.0) + weights = 1.0 / np.log(1.02 + freq) + weights = weights / weights.mean() + return torch.tensor(weights, dtype=torch.float32) + + +def estimate_pos_weight_from_masks2(ds_dual, ignore_index2: int = 255, max_samples: int = 800) -> Optional[torch.Tensor]: + """ + Estima pos_weight = neg/pos para BCEWithLogitsLoss. + Considera targets binários {0,1} ignorando ignore_index2. + """ + n = min(len(ds_dual), max_samples) + if n <= 0: + return None + + idxs = np.random.choice(len(ds_dual), size=n, replace=False) + pos = 0 + neg = 0 + for i in idxs: + it = ds_dual[i] + m2 = it["mask2"] if isinstance(it, dict) else it[2] + + if isinstance(m2, torch.Tensor): + m2 = m2.cpu().numpy() + m2 = np.array(m2).reshape(-1) + m2 = m2[m2 != ignore_index2] + if m2.size == 0: + continue + pos += int((m2 == 1).sum()) + neg += int((m2 == 0).sum()) + + if pos <= 0: + return None + pw = float(neg) / float(pos) + # clamp leve pra não explodir treino + pw = float(np.clip(pw, 1.0, 50.0)) + return torch.tensor([pw], dtype=torch.float32) + + +def dice_loss_with_logits(logits: torch.Tensor, targets: torch.Tensor, ignore_index: int = 255, eps: float = 1e-6) -> torch.Tensor: + """ + logits: [B,1,H,W] + targets: [B,H,W] com {0,1} e possivelmente ignore_index + """ + probs = torch.sigmoid(logits) + t = targets.float().unsqueeze(1) + + if ignore_index is not None: + valid = (targets != ignore_index).unsqueeze(1) + probs = probs[valid] + t = t[valid] + + probs = probs.reshape(-1) + t = t.reshape(-1) + + inter = (probs * t).sum() + denom = probs.sum() + t.sum() + eps + dice = (2.0 * inter + eps) / denom + return 1.0 - dice + + +@torch.no_grad() +def binary_iou_and_acc_from_logits(logits: torch.Tensor, targets: torch.Tensor, thr: float = 0.5, ignore_index: int = 255, eps: float = 1e-6) -> Tuple[float, float]: + probs = torch.sigmoid(logits) + preds = (probs >= thr).long().squeeze(1) + + if logits.shape[-2:] != targets.shape[-2:]: + preds = F.interpolate(preds.unsqueeze(1).float(), size=targets.shape[-2:], mode="nearest").long().squeeze(1) + + valid = (targets != ignore_index) + if valid.sum().item() == 0: + return 0.0, 0.0 + + p = preds[valid] + t = targets[valid].long() + + tp = ((p == 1) & (t == 1)).sum().float() + fp = ((p == 1) & (t == 0)).sum().float() + fn = ((p == 0) & (t == 1)).sum().float() + + iou = (tp / (tp + fp + fn + eps)).item() + acc = ((p == t).sum().float() / (t.numel() + eps)).item() + return iou, acc + + +def _pretty_iou(iou_list: List[float], class_name_by_id: Dict[int, str]) -> str: + parts = [] + for cid, v in enumerate(iou_list): + parts.append(f"{class_name_by_id.get(cid, str(cid))}:{v:.3f}") + return " | ".join(parts) + + +def save_checkpoint(path: str, + base_model: nn.Module, + corridor_head: nn.Module, + optimizer: torch.optim.Optimizer, + scaler: Optional[GradScaler], + epoch: int, + best_miou: float, + best_main_iou: float, + best_corr_iou: float, + extra: Optional[Dict[str, Any]] = None): + ckpt = { + "epoch": epoch, + "model": base_model.state_dict(), + "corridor_head": corridor_head.state_dict(), + "optimizer": optimizer.state_dict(), + "best_miou": best_miou, + "best_main_iou": best_main_iou, + "best_corr_iou": best_corr_iou, + } + if scaler is not None: + ckpt["scaler"] = scaler.state_dict() + if extra: + ckpt["extra"] = extra + torch.save(ckpt, path) + + +def load_checkpoint(path: str, + base_model: nn.Module, + corridor_head: nn.Module, + optimizer: Optional[torch.optim.Optimizer] = None, + scaler: Optional[GradScaler] = None, + map_location: str = "cpu") -> Dict[str, Any]: + ckpt = torch.load(path, map_location=map_location, weights_only=False) + base_model.load_state_dict(ckpt["model"], strict=True) + corridor_head.load_state_dict(ckpt["corridor_head"], strict=True) + if optimizer is not None and "optimizer" in ckpt: + optimizer.load_state_dict(ckpt["optimizer"]) + if scaler is not None and "scaler" in ckpt: + scaler.load_state_dict(ckpt["scaler"]) + return ckpt + + +# ----------------------------- +# Dual Dataset Wrapper +# ----------------------------- +class DualMaskROISegDataset(torch.utils.data.Dataset): + def __init__(self, base_ds: ROISegDataset, split_root: str, ignore_index2: int = 255): + self.base_ds = base_ds + self.split_root = split_root + self.ignore_index2 = ignore_index2 + + # NÃO falha aqui por pasta. Vamos resolver por caminho da mask1. + self.masks2_dir = os.path.join(split_root, "masks2") + self.has_flat_masks2 = os.path.isdir(self.masks2_dir) + + def __len__(self): + return len(self.base_ds) + + def _get_mask1_path(self, item: Dict[str, Any], idx: int) -> str: + # 1) se vier no dict, usa + for k in ("mask_path", "mask_file", "maskname", "mask_name"): + if k in item and item[k]: + return str(item[k]) + + # 2) fallback: usa o índice no ROISegDataset (ele tem msk_paths) + if hasattr(self.base_ds, "msk_paths"): + return str(self.base_ds.msk_paths[idx]) + + raise RuntimeError( + "ROISegDataset não retornou mask_path/mask_file e não possui base_ds.msk_paths. " + "Não dá pra localizar a máscara para buscar a masks2." + ) + + def _load_mask2_from_mask1_path(self, mask1_path: str, target_shape_hw: Tuple[int, int]) -> torch.Tensor: + # 1) modo robusto: espelha estrutura trocando /masks/ -> /masks2/ + norm = os.path.normpath(mask1_path) + parts = norm.split(os.sep) + try: + i = parts.index("masks") + parts[i] = "masks2" + path2 = os.sep.join(parts) + except ValueError: + path2 = "" + + # 2) fallback: modo “flat” (split_root/masks2/) + if (not path2) or (not os.path.exists(path2)): + if not self.has_flat_masks2: + raise FileNotFoundError( + f"Mask2 não encontrada espelhando a mask1.\n" + f"mask1: {mask1_path}\n" + f"tentativa: {path2}\n" + f"E também não existe: {self.masks2_dir}" + ) + filename = os.path.basename(mask1_path) + path2 = os.path.join(self.masks2_dir, filename) + if not os.path.exists(path2): + raise FileNotFoundError(f"Mask2 não encontrada: {path2}") + + m = np.array(Image.open(path2)) + if m.ndim == 3: + m = m[..., 0] + + if m.max() > 1: + m = (m >= 128).astype(np.uint8) + + H, W = target_shape_hw + if m.shape[0] != H or m.shape[1] != W: + m = np.array(Image.fromarray(m).resize((W, H), resample=Image.NEAREST)) + + return torch.from_numpy(m).long() + + def __getitem__(self, idx): + item = self.base_ds[idx] + # caso venha (dict,) ou [dict] + if isinstance(item, (list, tuple)) and len(item) == 1 and isinstance(item[0], dict): + item = item[0] + + if isinstance(item, dict): + pass + elif isinstance(item, (list, tuple)) and len(item) >= 2: + img, mask = item[0], item[1] + item = {"image": img, "mask": mask} + else: + raise RuntimeError( + f"ROISegDataset retornou formato inesperado: type={type(item)} value={item}" + ) + + img = item["image"] + mask1 = item["mask"] + + if isinstance(mask1, torch.Tensor): + h, w = int(mask1.shape[-2]), int(mask1.shape[-1]) + else: + m = np.array(mask1) + h, w = m.shape[0], m.shape[1] + + mask1_path = self._get_mask1_path(item, idx) + mask2 = self._load_mask2_from_mask1_path(mask1_path, (h, w)) + + item["mask1"] = mask1 + item["mask2"] = mask2 + return item + + +# ----------------------------- +# Corridor head +# ----------------------------- +class CorridorHead(nn.Module): + """ + Head simples: + feat (B,C,H,W) + logits_seg (B,K,H,W) -> concat -> convs -> 1 canal (logit corredor) + """ + def __init__(self, feat_ch: int, num_classes: int, hidden: int = 256, dropout: float = 0.1): + super().__init__() + in_ch = feat_ch + num_classes + self.in_ch = in_ch + self.net = nn.Sequential( + nn.Conv2d(in_ch, hidden, kernel_size=3, padding=1), + nn.BatchNorm2d(hidden), + nn.ReLU(inplace=True), + nn.Dropout2d(dropout), + nn.Conv2d(hidden, hidden, kernel_size=3, padding=1), + nn.BatchNorm2d(hidden), + nn.ReLU(inplace=True), + nn.Dropout2d(dropout), + nn.Conv2d(hidden, 1, kernel_size=1) + ) + + def forward(self, feat: torch.Tensor, logits_seg: torch.Tensor) -> torch.Tensor: + x = torch.cat([feat, logits_seg], dim=1) + return self.net(x) + + +# ----------------------------- +# One epoch +# ----------------------------- +def run_one_epoch_dual(base_model: nn.Module, + corridor_head: nn.Module, + loader: DataLoader, + optimizer: Optional[torch.optim.Optimizer], + device: torch.device, + num_classes: int, + ignore_index1: int, + ignore_index2: int, + criterion_seg: nn.Module, + bce_corr: nn.Module, + lambda_corr: float, + amp: bool, + scaler: Optional[GradScaler], + train: bool, + grad_accum: int = 1, + corridor_thr: float = 0.5, + corr_dice_mix: float = 0.0) -> Dict[str, Any]: + + base_model.train(train) + corridor_head.train(train) + + total_loss = 0.0 + total_loss_seg = 0.0 + total_loss_corr = 0.0 + + cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device) + + corr_iou_accum = 0.0 + corr_acc_accum = 0.0 + + t0 = time.time() + n_batches = 0 + + with torch.set_grad_enabled(train): + if train and optimizer is not None: + optimizer.zero_grad(set_to_none=True) + + for step, (imgs, masks1, masks2) in enumerate(loader, start=1): + imgs = imgs.to(device, non_blocking=True) + masks1 = masks1.to(device, non_blocking=True) + masks2 = masks2.to(device, non_blocking=True) + + #print("masks2 unique:", torch.unique(masks2)[:10]) + #print("masks2 max:", masks2.max().item(), "min:", masks2.min().item()) + + imgs = normalize_img(imgs) + + with autocast(device_type="cuda", enabled=amp and device.type == "cuda"): + out = base_model(pixel_values=imgs) + logits1 = out.logits # [B,K,H,W] + + if logits1.shape[-2:] != masks1.shape[-2:]: + logits1 = F.interpolate(logits1, size=masks1.shape[-2:], mode="bilinear", align_corners=False) + + loss_seg = criterion_seg(logits1, masks1) + + # feat do encoder: pega hidden_states (se disponível), senão reusa logits como proxy + feat = None + if hasattr(out, "hidden_states") and out.hidden_states is not None: + feat = out.hidden_states[-1] # [B, C, h, w] + else: + # fallback: usa logits como feature (não ideal, mas mantém vivo) + feat = logits1 + + if feat.shape[-2:] != logits1.shape[-2:]: + feat = F.interpolate(feat, size=logits1.shape[-2:], mode="bilinear", align_corners=False) + + logits_corr = corridor_head(feat, logits1) # [B,1,H,W] + + if logits_corr.shape[-2:] != masks2.shape[-2:]: + logits_corr = F.interpolate(logits_corr, size=masks2.shape[-2:], mode="bilinear", align_corners=False) + + #with torch.no_grad(): + # p = torch.sigmoid(logits_corr) + # print("corr prob mean:", p.mean().item(), "min:", p.min().item(), "max:", p.max().item()) + + # BCE (com ignore via masking manual) + valid = (masks2 != ignore_index2) + if valid.sum().item() > 0: + tgt = masks2.float().unsqueeze(1) + bce_val = bce_corr(logits_corr[valid.unsqueeze(1)], tgt[valid.unsqueeze(1)]) + if corr_dice_mix > 0: + d_val = dice_loss_with_logits(logits_corr, masks2, ignore_index=ignore_index2) + loss_corr = (1.0 - corr_dice_mix) * bce_val + corr_dice_mix * d_val + else: + loss_corr = bce_val + else: + loss_corr = torch.zeros([], device=device, dtype=loss_seg.dtype) + + loss = loss_seg + lambda_corr * loss_corr + + if train and grad_accum > 1: + loss = loss / grad_accum + + if train and optimizer is not None: + if amp and scaler is not None and device.type == "cuda": + scaler.scale(loss).backward() + if (step % grad_accum) == 0: + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + else: + loss.backward() + if (step % grad_accum) == 0: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + total_loss += float(loss.item()) * (grad_accum if (train and grad_accum > 1) else 1.0) + total_loss_seg += float(loss_seg.item()) + total_loss_corr += float(loss_corr.item()) + n_batches += 1 + + # metrics seg + preds1 = torch.argmax(logits1.detach(), dim=1) + update_confusion_matrix(cm, preds1, masks1, num_classes, ignore_index=ignore_index1) + + # metrics corr + ciou, cacc = binary_iou_and_acc_from_logits(logits_corr.detach(), masks2, thr=corridor_thr, ignore_index=ignore_index2) + corr_iou_accum += ciou + corr_acc_accum += cacc + + miou, iou_per_class = compute_iou_from_cm(cm) + acc = compute_pixel_acc_from_cm(cm) + + out = { + "loss": total_loss / max(1, n_batches), + "loss_seg": total_loss_seg / max(1, n_batches), + "loss_corr": total_loss_corr / max(1, n_batches), + "acc": acc, + "miou": miou, + "iou_per_class": iou_per_class, + "corr_iou": corr_iou_accum / max(1, n_batches), + "corr_acc": corr_acc_accum / max(1, n_batches), + "time_s": time.time() - t0 + } + return out + + +# ----------------------------- +# LR / EarlyStop helpers +# ----------------------------- +@dataclass +class EarlyStopState: + best: float = -1e9 + bad_epochs: int = 0 + stopped: bool = False + + +def agg_metric(es_metric: str, miou: float, main_iou: float, corr_iou: float) -> float: + if es_metric == "miou": + return float(miou) + if es_metric == "main": + return float(main_iou) + if es_metric == "corr": + return float(corr_iou) + if es_metric == "harmonic": + eps = 1e-6 + a = max(eps, float(main_iou)) + b = max(eps, float(corr_iou)) + return float(2.0 * a * b / (a + b + eps)) + # default + return float(miou) + + +def apply_warmup(optimizer, base_lr: float, epoch: int, warmup_epochs: int): + if warmup_epochs <= 0: + return + if epoch <= warmup_epochs: + wf = epoch / float(warmup_epochs) + for g in optimizer.param_groups: + g["lr"] = base_lr * wf + + +# ----------------------------- +# Main +# ----------------------------- +def main(): + ap = argparse.ArgumentParser() + + ap.add_argument("--config", default="config.json") + ap.add_argument("--epochs", type=int, default=120) + ap.add_argument("--batch", type=int, default=4) + ap.add_argument("--num_workers", type=int, default=4) + ap.add_argument("--seed", type=int, default=42) + + ap.add_argument("--lr", type=float, default=6e-5) + ap.add_argument("--wd", type=float, default=1e-2) + ap.add_argument("--min_lr", type=float, default=6e-6) + + ap.add_argument("--amp", action="store_true") + ap.add_argument("--amp_val", action="store_true") + ap.add_argument("--grad_accum", type=int, default=1) + + ap.add_argument("--ignore_index", type=int, default=255) + ap.add_argument("--ignore_index2", type=int, default=255) + + ap.add_argument("--lambda_corr", type=float, default=0.5) + ap.add_argument("--corr_thr", type=float, default=0.5) + + # class weights + ap.add_argument("--use_weights", action="store_true") + ap.add_argument("--cw_max_samples", type=int, default=800) + + # pos_weight corredor + ap.add_argument("--auto_pos_weight", action="store_true") + ap.add_argument("--pw_max_samples", type=int, default=800) + + # mistura dice no corredor (rampa) + ap.add_argument("--corr_dice_max", type=float, default=0.15) + ap.add_argument("--corr_dice_ramp_epochs", type=int, default=10) + + # schedulers + ap.add_argument("--warmup_epochs", type=int, default=0) + ap.add_argument("--cosine_epochs", type=int, default=0, help="se >0: usa cosine por N épocas e depois troca p/ plateau") + ap.add_argument("--plateau_patience", type=int, default=3) + ap.add_argument("--plateau_factor", type=float, default=0.6) + ap.add_argument("--plateau_cooldown", type=int, default=1) + + # early stopping + ap.add_argument("--es_metric", default="harmonic", choices=["miou", "main", "corr", "harmonic"]) + ap.add_argument("--es_patience", type=int, default=12) + ap.add_argument("--es_min_delta", type=float, default=2e-4) + + # logs/ckpt + ap.add_argument("--save_every", type=int, default=0) + ap.add_argument("--resume", action="store_true") + + ap.add_argument("--main_class", type=str, default=None) + + args = ap.parse_args() + + seed_everything(args.seed) + + with open(args.config, "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"] + BACKBONE = config["backbone"] + MAIN_CLASS_NAME = str(config.get("main_class_name", "erva")).lower() + if args.main_class is not None: + MAIN_CLASS_NAME = args.main_class.lower() + + dataset_path = os.path.join(MODELO, "dataset") + labelmap_path = os.path.join(dataset_path, "labelmap.txt") + + # save paths + save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME + "_dual") + os.makedirs(save_path, exist_ok=True) + last_ckpt_path = os.path.join(save_path, "last.pt") + best_miou_path = os.path.join(save_path, "best_miou.pt") + best_main_path = os.path.join(save_path, "best_main.pt") + best_corr_path = os.path.join(save_path, "best_corr.pt") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print("Device:", device) + + # seus splits estão em train/group/... então o ROISegDataset deve saber ler de lá + split_train = os.path.join(dataset_path, "split", "train") + split_val = os.path.join(dataset_path, "split", "val") + + ds_train_base = ROISegDataset( + split_train, + save_path, ROI_INICIO, ROI_TAMANHO, + RESOLUCAO[0], RESOLUCAO[1], labelmap_path + ) + + ds_val_base = ROISegDataset( + split_val, + save_path, ROI_INICIO, ROI_TAMANHO, + RESOLUCAO[0], RESOLUCAO[1], labelmap_path + ) + + ds_train = DualMaskROISegDataset(ds_train_base, split_train, ignore_index2=args.ignore_index2) + ds_val = DualMaskROISegDataset(ds_val_base, split_val, ignore_index2=args.ignore_index2) + + CLASS_NAMES = getattr(ds_train_base, "classes", None) + if CLASS_NAMES is None: + raise RuntimeError("ROISegDataset precisa expor .classes (list ou dict).") + + if isinstance(CLASS_NAMES, list): + num_classes = len(CLASS_NAMES) + class_id_by_name = {n.lower(): i for i, n in enumerate(CLASS_NAMES)} + class_name_by_id = {i: n for i, n in enumerate(CLASS_NAMES)} + elif isinstance(CLASS_NAMES, dict): + if all(isinstance(k, int) for k in CLASS_NAMES.keys()): + num_classes = len(CLASS_NAMES) + class_name_by_id = {int(k): str(v) for k, v in CLASS_NAMES.items()} + class_id_by_name = {str(v).lower(): int(k) for k, v in CLASS_NAMES.items()} + else: + class_id_by_name = {str(k).lower(): int(v) for k, v in CLASS_NAMES.items()} + num_classes = len(class_id_by_name) + class_name_by_id = {v: k for k, v in class_id_by_name.items()} + else: + raise RuntimeError("Formato de .classes não reconhecido.") + + main_class_id = class_id_by_name.get(MAIN_CLASS_NAME, None) + if main_class_id is None: + print(f"[WARN] main_class_name='{MAIN_CLASS_NAME}' não encontrado. best_main_iou usa mIoU.") + else: + print(f"Main class: '{MAIN_CLASS_NAME}' -> id={main_class_id}") + + + dl_train = DataLoader( + ds_train, + batch_size=args.batch, + shuffle=True, + num_workers=args.num_workers, + pin_memory=(device.type == "cuda"), + collate_fn=default_collate_dual + ) + dl_val = DataLoader( + ds_val, + batch_size=args.batch, + shuffle=False, + num_workers=max(0, args.num_workers // 2), + pin_memory=(device.type == "cuda"), + collate_fn=default_collate_dual + ) + + # model base SegFormer + base_model = SegformerForSemanticSegmentation.from_pretrained( + BACKBONE, + num_labels=num_classes, + ignore_mismatched_sizes=True, + use_safetensors=True, # <- evita torch.load do .bin + ) + base_model.config.output_hidden_states = True + base_model.to(device) + + # descobre feat_ch via dummy forward (mais robusto) + with torch.no_grad(): + dummy = torch.zeros((1, 3, 512, 512), device=device) + out = base_model(pixel_values=dummy) + feat_ch = None + if hasattr(out, "hidden_states") and out.hidden_states is not None: + feat_ch = int(out.hidden_states[-1].shape[1]) + else: + feat_ch = int(out.logits.shape[1]) # fallback + corridor_head = CorridorHead(feat_ch=feat_ch, num_classes=num_classes, hidden=256, dropout=0.1).to(device) + print(f"[corridor_head] in_ch = feat({feat_ch}) + logits_seg({num_classes}) = {corridor_head.in_ch}") + + # losses + class_weights = None + if args.use_weights: + w = estimate_class_weights_from_ds(ds_train_base, num_classes=num_classes, ignore_index=args.ignore_index, max_samples=args.cw_max_samples) + class_weights = w.to(device) + print("[INFO] class_weights:", class_weights.detach().cpu().numpy().round(3).tolist()) + criterion_seg = nn.CrossEntropyLoss(ignore_index=args.ignore_index, weight=class_weights) + + pos_weight = None + if args.auto_pos_weight: + pw = estimate_pos_weight_from_masks2(ds_train, ignore_index2=args.ignore_index2, max_samples=args.pw_max_samples) + if pw is not None: + pos_weight = pw.to(device) + print("[INFO] pos_weight corredor (neg/pos):", float(pos_weight.item())) + else: + print("[INFO] pos_weight não estimável (sem positivos suficientes), usando None.") + bce_corr = nn.BCEWithLogitsLoss(pos_weight=pos_weight) if pos_weight is not None else nn.BCEWithLogitsLoss() + + # optimizer (base + head) + params = list(base_model.parameters()) + list(corridor_head.parameters()) + optimizer = torch.optim.AdamW(params, lr=args.lr, weight_decay=args.wd) + + scaler = GradScaler(enabled=(args.amp and device.type == "cuda")) + + # schedulers + cosine = None + if args.cosine_epochs and args.cosine_epochs > 0: + cosine = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, + T_max=max(1, args.cosine_epochs), + eta_min=args.min_lr + ) + plateau = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, + mode="min", + factor=args.plateau_factor, + patience=args.plateau_patience, + cooldown=args.plateau_cooldown, + min_lr=args.min_lr, + verbose=True + ) + active_sched = "cosine" if cosine is not None else "plateau" + + # resume + start_epoch = 1 + best_miou = -1.0 + best_main_iou = -1.0 + best_corr_iou = -1.0 + + if args.resume and os.path.exists(last_ckpt_path): + ckpt = load_checkpoint(last_ckpt_path, base_model, corridor_head, optimizer=optimizer, scaler=scaler, map_location="cpu") + start_epoch = int(ckpt.get("epoch", 0)) + 1 + best_miou = float(ckpt.get("best_miou", -1.0)) + best_main_iou = float(ckpt.get("best_main_iou", -1.0)) + best_corr_iou = float(ckpt.get("best_corr_iou", -1.0)) + print(f"[RESUME] epoch={start_epoch} best_miou={best_miou:.4f} best_main_iou={best_main_iou:.4f} best_corr_iou={best_corr_iou:.4f}") + + # early stopping + es = EarlyStopState(best=-1e9, bad_epochs=0, stopped=False) + + for epoch in range(start_epoch, args.epochs + 1): + lr_now = optimizer.param_groups[0]["lr"] + + # warmup + apply_warmup(optimizer, args.lr, epoch, args.warmup_epochs) + lr_now = optimizer.param_groups[0]["lr"] + + # corr dice ramp + if args.corr_dice_max > 0 and args.corr_dice_ramp_epochs > 0: + corr_dice_mix = min(args.corr_dice_max, (epoch - 1) / float(args.corr_dice_ramp_epochs) * args.corr_dice_max) + else: + corr_dice_mix = 0.0 + + print(f"\n==== Epoch {epoch}/{args.epochs} | lr={lr_now:.2e} | lambda_corr={args.lambda_corr} | corr_dice={corr_dice_mix:.3f} | sched={active_sched} ====") + + if device.type == "cuda": + torch.cuda.empty_cache() + + tr = run_one_epoch_dual( + base_model=base_model, + corridor_head=corridor_head, + loader=dl_train, + optimizer=optimizer, + device=device, + num_classes=num_classes, + ignore_index1=args.ignore_index, + ignore_index2=args.ignore_index2, + criterion_seg=criterion_seg, + bce_corr=bce_corr, + lambda_corr=args.lambda_corr, + amp=args.amp, + scaler=scaler, + train=True, + grad_accum=max(1, args.grad_accum), + corridor_thr=args.corr_thr, + corr_dice_mix=corr_dice_mix + ) + + if device.type == "cuda": + torch.cuda.empty_cache() + + va = run_one_epoch_dual( + base_model=base_model, + corridor_head=corridor_head, + loader=dl_val, + optimizer=None, + device=device, + num_classes=num_classes, + ignore_index1=args.ignore_index, + ignore_index2=args.ignore_index2, + criterion_seg=criterion_seg, + bce_corr=bce_corr, + lambda_corr=args.lambda_corr, + amp=args.amp_val, + scaler=None, + train=False, + grad_accum=1, + corridor_thr=args.corr_thr, + corr_dice_mix=corr_dice_mix + ) + + # scheduler step + if cosine is not None and epoch <= args.cosine_epochs: + cosine.step() + else: + active_sched = "plateau" + plateau.step(va["loss"]) + + main_iou = va["miou"] if main_class_id is None else va["iou_per_class"][main_class_id] + corr_iou = float(va["corr_iou"]) + + print(f"TRAIN: loss={tr['loss']:.4f} (seg={tr['loss_seg']:.4f} corr={tr['loss_corr']:.4f}) " + f"acc={tr['acc']:.4f} miou={tr['miou']:.4f} corr_iou={tr['corr_iou']:.4f} (t={tr['time_s']:.1f}s)") + print(f"VAL : loss={va['loss']:.4f} (seg={va['loss_seg']:.4f} corr={va['loss_corr']:.4f}) " + f"acc={va['acc']:.4f} miou={va['miou']:.4f} main_iou={float(main_iou):.4f} " + f"corr_iou={va['corr_iou']:.4f} corr_acc={va['corr_acc']:.4f} (t={va['time_s']:.1f}s)") + print("IoU per class:", _pretty_iou(va["iou_per_class"], class_name_by_id)) + + # save last + save_checkpoint( + last_ckpt_path, + base_model, + corridor_head, + optimizer, + scaler=scaler, + epoch=epoch, + best_miou=best_miou, + best_main_iou=best_main_iou, + best_corr_iou=best_corr_iou, + extra={ + "val_loss": va["loss"], + "val_miou": va["miou"], + "val_main_iou": float(main_iou), + "val_corr_iou": float(corr_iou), + "lr": optimizer.param_groups[0]["lr"], + "corr_dice_mix": corr_dice_mix + } + ) + + # save every + if args.save_every > 0 and (epoch % args.save_every == 0): + save_checkpoint( + os.path.join(save_path, f"epoch_{epoch:04d}.pt"), + base_model, corridor_head, optimizer, scaler=scaler, + epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou + ) + + # bests + if va["miou"] > best_miou: + best_miou = va["miou"] + save_checkpoint(best_miou_path, base_model, corridor_head, optimizer, scaler=scaler, + epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou) + print(f"[BEST mIoU] {best_miou:.4f} -> saved: {best_miou_path}") + + if float(main_iou) > best_main_iou: + best_main_iou = float(main_iou) + save_checkpoint(best_main_path, base_model, corridor_head, optimizer, scaler=scaler, + epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou) + print(f"[BEST MAIN] {best_main_iou:.4f} -> saved: {best_main_path}") + + if corr_iou > best_corr_iou: + best_corr_iou = corr_iou + save_checkpoint(best_corr_path, base_model, corridor_head, optimizer, scaler=scaler, + epoch=epoch, best_miou=best_miou, best_main_iou=best_main_iou, best_corr_iou=best_corr_iou) + print(f"[BEST CORR] {best_corr_iou:.4f} -> saved: {best_corr_path}") + + # early stopping + score = agg_metric(args.es_metric, va["miou"], float(main_iou), corr_iou) + if score > es.best + args.es_min_delta: + es.best = score + es.bad_epochs = 0 + print(f"[ES] improved {args.es_metric} -> {score:.6f}") + else: + es.bad_epochs += 1 + print(f"[ES] no improve ({es.bad_epochs}/{args.es_patience}) best={es.best:.6f} now={score:.6f}") + if es.bad_epochs >= args.es_patience: + print(f"🛑 Early stopping acionado (metric={args.es_metric}).") + break + + +if __name__ == "__main__": + main() diff --git a/Python/OAK/datasets/_9_test_fastscnn.py b/Python/OAK/datasets/_9_test_fastscnn.py index c68e1ee26..f466596f4 100644 --- a/Python/OAK/datasets/_9_test_fastscnn.py +++ b/Python/OAK/datasets/_9_test_fastscnn.py @@ -213,7 +213,8 @@ def main(): cv2.destroyAllWindows() else: # === Modo imagens (agrupado + fallback) === - test_root = os.path.join(dataset_path, "split", split_folder) + #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: diff --git a/Python/OAK/datasets/_9_test_segformer_b3.py b/Python/OAK/datasets/_9_test_segformer_b3.py new file mode 100644 index 000000000..68a848e7d --- /dev/null +++ b/Python/OAK/datasets/_9_test_segformer_b3.py @@ -0,0 +1,405 @@ +#!/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). +""" + +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 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)") + 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_b3 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.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() # RGB + 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 + 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 SegFormer (OAK + PyTorch)", cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR)) + if cv2.waitKey(1) & 0xFF == ord('q'): + break + + cv2.destroyAllWindows() + + else: + # === Modo imagens (agrupado + fallback) === + # Igual teu fastscnn: você pode usar dataset/512x288 diretamente + 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: + # fallback clássico + 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) == len(mask_paths) and len(image_paths) > 0, "Nenhuma imagem/máscara encontrada." + + 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] + 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 = 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) + 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) + # Overlay da predição sobre a imagem original + overlay_pred = cv2.addWeighted(img_resized, 0.6, pred_rgb, 0.4, 0.0) + resultado = np.concatenate([img_resized, mask_gt_rgb, 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) + + 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() + + +if __name__ == "__main__": + main() diff --git a/Python/OAK/datasets/_9_test_segformer_b3_dual.py b/Python/OAK/datasets/_9_test_segformer_b3_dual.py new file mode 100644 index 000000000..364486f70 --- /dev/null +++ b/Python/OAK/datasets/_9_test_segformer_b3_dual.py @@ -0,0 +1,570 @@ +import json +import os +import glob +import argparse +from dataclasses import dataclass +from typing import List, Optional + +import cv2 +import numpy as np +import torch +import torch.nn as nn +from transformers import SegformerForSemanticSegmentation + + +# ----------------------------- +# Utils +# ----------------------------- +def find_device(): + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + +def load_labelmap_lut_bgr(labelmap_path: str): + """ + Lê labelmap no formato: + classe:r,g,b + Ignora linhas vazias e comentários '#' + Se existir 'ignore', retorna ignore_id=255 e ignore_bgr. + + Retorna: + lut_bgr: dict[int, tuple(b,g,r)] + id_to_name: dict[int, str] + ignore_id: int (default 255) + """ + lut_bgr = {} + id_to_name = {} + ignore_id = 255 + ignore_bgr = (180, 0, 180) # fallback + + idx = 0 + with open(labelmap_path, "r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line or line.startswith("#"): + continue + parts = line.split(":") + if len(parts) < 2: + continue + name = parts[0].strip() + rgb_str = parts[1].strip() + r, g, b = map(int, rgb_str.split(",")) + bgr = (b, g, r) + + if name.lower() == "ignore": + ignore_bgr = bgr + continue + + lut_bgr[idx] = bgr + id_to_name[idx] = name + idx += 1 + + lut_bgr[ignore_id] = ignore_bgr + return lut_bgr, id_to_name, ignore_id + +def imread_bgr(path: str) -> np.ndarray: + img = cv2.imread(path, cv2.IMREAD_COLOR) + if img is None: + raise RuntimeError(f"Falha ao ler imagem: {path}") + return img + + +def imread_gray(path: str) -> np.ndarray: + m = cv2.imread(path, cv2.IMREAD_GRAYSCALE) + if m is None: + raise RuntimeError(f"Falha ao ler máscara: {path}") + return m + + +def ensure_uint8(x: np.ndarray) -> np.ndarray: + if x.dtype == np.uint8: + return x + x = np.clip(x, 0, 255).astype(np.uint8) + return x + + +def normalize01(x: np.ndarray, eps: float = 1e-6) -> np.ndarray: + x = x.astype(np.float32) + mn, mx = float(x.min()), float(x.max()) + return (x - mn) / (mx - mn + eps) + + +def colorize_seg(mask_ids: np.ndarray, lut_bgr: dict) -> np.ndarray: + h, w = mask_ids.shape[:2] + out = np.zeros((h, w, 3), dtype=np.uint8) + + # pinta cada id presente + for cid, bgr in lut_bgr.items(): + out[mask_ids == cid] = bgr + + return out + + +def overlay_mask(img_bgr: np.ndarray, mask_bgr: np.ndarray, alpha: float = 0.45) -> np.ndarray: + return cv2.addWeighted(img_bgr, 1.0 - alpha, mask_bgr, alpha, 0.0) + + +def put_hud(img: np.ndarray, lines: List[str]) -> np.ndarray: + out = img.copy() + y = 22 + for s in lines: + cv2.putText(out, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (15, 15, 15), 3, cv2.LINE_AA) + cv2.putText(out, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (245, 245, 245), 1, cv2.LINE_AA) + y += 22 + return out + + +# ----------------------------- +# Corridor Head (igual ao treino) +# ----------------------------- +class CorridorHead(nn.Module): + """ + Head simples: + feat (B,C,H,W) + logits_seg (B,K,H,W) -> concat -> convs -> 1 canal (logit corredor) + """ + def __init__(self, feat_ch: int, num_classes: int, hidden: int = 256, dropout: float = 0.1): + super().__init__() + in_ch = feat_ch + num_classes + self.in_ch = in_ch + self.net = nn.Sequential( + nn.Conv2d(in_ch, hidden, kernel_size=3, padding=1), + nn.BatchNorm2d(hidden), + nn.ReLU(inplace=True), + nn.Dropout2d(dropout), + nn.Conv2d(hidden, hidden, kernel_size=3, padding=1), + nn.BatchNorm2d(hidden), + nn.ReLU(inplace=True), + nn.Dropout2d(dropout), + nn.Conv2d(hidden, 1, kernel_size=1) + ) + + def forward(self, feat: torch.Tensor, logits_seg: torch.Tensor) -> torch.Tensor: + x = torch.cat([feat, logits_seg], dim=1) + return self.net(x) + + +# ----------------------------- +# Dataset discovery (split/test/group/**) +# ----------------------------- +@dataclass +class Sample: + img_path: str + mask_path: Optional[str] + mask2_path: Optional[str] + group_name: str + filename: str + + +IMG_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp") + + +def discover_samples(split_root: str) -> List[Sample]: + """ + Espera estrutura: + split/test/group//images/*.png + split/test/group//masks/*.png + split/test/group//masks2/*.png (opcional, mas recomendado) + """ + group_root = os.path.join(split_root, "group") + if not os.path.isdir(group_root): + raise RuntimeError(f"Não achei pasta: {group_root}") + + samples: List[Sample] = [] + + # pega qualquer images/ em qualquer subpasta de group + img_dirs = glob.glob(os.path.join(group_root, "**", "images"), recursive=True) + img_dirs = [d for d in img_dirs if os.path.isdir(d)] + + for idir in img_dirs: + base = os.path.dirname(idir) # .../group/ + group_name = os.path.relpath(base, group_root).replace("\\", "/") + + mdir = os.path.join(base, "masks") + m2dir = os.path.join(base, "masks2") + + img_paths = [] + for ext in IMG_EXTS: + img_paths.extend(glob.glob(os.path.join(idir, f"*{ext}"))) + img_paths = sorted(img_paths) + + def find_mask(dir_path: str, filename: str): + stem, _ = os.path.splitext(filename) + for ext in [".png", ".jpg", ".jpeg", ".bmp", ".tif"]: + p = os.path.join(dir_path, stem + ext) + if os.path.exists(p): + return p + return None + + for ip in img_paths: + fn = os.path.basename(ip) + + mask_path = find_mask(mdir, fn) + mask2_path = find_mask(m2dir, fn) + + samples.append(Sample( + img_path=ip, + mask_path=mask_path, + mask2_path=mask2_path, + group_name=group_name, + filename=fn + )) + + if len(samples) == 0: + raise RuntimeError(f"Nenhuma imagem encontrada em: {group_root}/**/images") + return samples + + +# ----------------------------- +# Model loading +# ----------------------------- +def build_models(num_classes: int, device: torch.device, backbone: str): + base_model = SegformerForSemanticSegmentation.from_pretrained( + backbone, + num_labels=num_classes, + ignore_mismatched_sizes=True, + use_safetensors=True, # <- evita torch.load do .bin + ) + base_model.to(device) + base_model.config.output_hidden_states = True + + # descobre feat_ch real a partir de um forward fake + with torch.no_grad(): + dummy = torch.zeros((1, 3, 512, 512), device=device) + out = base_model(pixel_values=dummy, output_hidden_states=True) + feat_ch = int(out.hidden_states[-1].shape[1]) + + in_ch = feat_ch + num_classes + corridor_head = CorridorHead(feat_ch=feat_ch, num_classes=num_classes, hidden=256, dropout=0.1) + print(f"[corridor_head] in_ch = feat({feat_ch}) + logits_seg({num_classes}) = {in_ch}") + return base_model, corridor_head + + +def make_corridor_head(in_ch: int, mid: int = 256): + """ + Cria CorridorHead independente de qual assinatura a classe tem no script. + Suporta: + - CorridorHead(in_ch=..., mid=...) + - CorridorHead(in_channels=..., hidden=...) + - CorridorHead(feat_ch=..., num_classes=...) (quando in_ch = feat_ch + num_classes) + """ + # 1) assinatura direta in_ch/mid + try: + return CorridorHead(in_ch=in_ch, mid=mid) + except TypeError: + pass + + # 2) variações comuns + try: + return CorridorHead(in_channels=in_ch, mid=mid) + except TypeError: + pass + + try: + return CorridorHead(in_channels=in_ch, hidden=mid) + except TypeError: + pass + + # 3) assinatura "feat_ch + num_classes" + # aqui só funciona se for exatamente "feat_ch, num_classes" + # e se in_ch for decomponível (ex: 6 = 3+3 ou 515 = 512+3) + for feat_ch_guess, num_classes_guess in [(512, 3), (3, 3)]: + if feat_ch_guess + num_classes_guess == in_ch: + try: + return CorridorHead(feat_ch=feat_ch_guess, num_classes=num_classes_guess, hidden=mid) + except TypeError: + try: + return CorridorHead(feat_ch=feat_ch_guess, num_classes=num_classes_guess) + except TypeError: + pass + + raise RuntimeError( + f"Não consegui instanciar CorridorHead para in_ch={in_ch}. " + f"Verifique a assinatura do __init__() da classe CorridorHead no script." + ) + +def load_dual_checkpoint(ckpt_path, base_model, corridor_head, device): + ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) + + base_model.load_state_dict(ckpt["model"], strict=True) + + # >>> pega quantos canais o head do checkpoint espera + sd = ckpt["corridor_head"] + in_ch_ckpt = int(sd["net.0.weight"].shape[1]) # ex: 6 ou 515 + + # >>> recria head se o in_ch não bater + in_ch_model = int(corridor_head.net[0].in_channels) + if in_ch_model != in_ch_ckpt: + print(f"[corridor_head] rebuild: in_ch_model={in_ch_model} -> in_ch_ckpt={in_ch_ckpt}") + corridor_head = make_corridor_head(in_ch=in_ch_ckpt, mid=256) + + corridor_head.load_state_dict(sd, strict=True) + + base_model.to(device).eval() + corridor_head.to(device).eval() + meta = ckpt.get("meta", {}) + return meta, corridor_head + + +# ----------------------------- +# Inference +# ----------------------------- +def run_corridor_head(corridor_head, x_rgb, feat, logits_seg): + import inspect + import torch.nn.functional as F + + n_args = len(inspect.signature(corridor_head.forward).parameters) + + # descobre in_ch esperado pela primeira conv do head + in_ch = None + if hasattr(corridor_head, "net"): + try: + in_ch = int(corridor_head.net[0].in_channels) + except Exception: + in_ch = None + + # prepara entradas conforme o head + if in_ch == 6: + # RGB(3) + probs(3) + probs = torch.softmax(logits_seg, dim=1) + if probs.shape[-2:] != x_rgb.shape[-2:]: + probs = F.interpolate(probs, size=x_rgb.shape[-2:], mode="bilinear", align_corners=False) + feat_in = x_rgb + logits_in = probs + else: + # feat(512) + logits(3) + if logits_seg.shape[-2:] != feat.shape[-2:]: + logits_in = F.interpolate(logits_seg, size=feat.shape[-2:], mode="bilinear", align_corners=False) + else: + logits_in = logits_seg + feat_in = feat + + # chama do jeito que o head espera + if n_args == 2: + return corridor_head(feat_in, logits_in) + elif n_args == 1: + x2 = torch.cat([feat_in, logits_in], dim=1) + return corridor_head(x2) + + raise RuntimeError(f"Assinatura inesperada: {inspect.signature(corridor_head.forward)}") + +@torch.no_grad() +def infer_dual(base_model, corridor_head, img_bgr: np.ndarray, device: torch.device, + input_size: int = 512, corridor_thr: float = 0.5): + """ + Retorna: + pred_ids_full: [H,W] uint8 + prob_corr_full: [H,W] float32 (0..1) + bin_corr_full: [H,W] uint8 (0/255) + """ + h0, w0 = img_bgr.shape[:2] + + img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) + img_res = cv2.resize(img_rgb, (input_size, input_size), interpolation=cv2.INTER_AREA) + + x = torch.from_numpy(img_res).float() / 255.0 + x = x.permute(2, 0, 1).unsqueeze(0).to(device) + + out = base_model(pixel_values=x, output_hidden_states=True) + logits_seg = out.logits # [B, 3, H, W] + feat = out.hidden_states[-1] # [B, 512, h, w] (não vamos usar no modo 6) + + logits_corr = run_corridor_head(corridor_head, x, feat, logits_seg) + + prob_corr = torch.sigmoid(logits_corr)[0, 0].detach().float().cpu().numpy() + bin_corr = (prob_corr >= corridor_thr).astype(np.uint8) * 255 + + pred_ids = torch.argmax(logits_seg, dim=1)[0].detach().cpu().numpy().astype(np.uint8) + + # volta pro tamanho original + pred_ids_full = cv2.resize(pred_ids, (w0, h0), interpolation=cv2.INTER_NEAREST) + prob_corr_full = cv2.resize(prob_corr, (w0, h0), interpolation=cv2.INTER_LINEAR) + bin_corr_full = cv2.resize(bin_corr, (w0, h0), interpolation=cv2.INTER_NEAREST) + + return pred_ids_full, prob_corr_full, bin_corr_full + + +# ----------------------------- +# Viewer +# ----------------------------- +def make_panel(sample: Sample , lut_bgr: dict, img_bgr: np.ndarray, + gt1: Optional[np.ndarray], + pred1: np.ndarray, + gt2: Optional[np.ndarray], + prob2: np.ndarray, + bin2: np.ndarray, + corridor_thr: float): + h, w = img_bgr.shape[:2] + + pred1_col = colorize_seg(pred1, lut_bgr) + pred1_ov = overlay_mask(img_bgr, pred1_col, 0.45) + + if gt1 is not None: + gt1_col = colorize_seg(gt1, lut_bgr) + gt1_ov = overlay_mask(img_bgr, gt1_col, 0.45) + else: + gt1_ov = img_bgr.copy() + gt1_col = np.zeros_like(img_bgr) + + # --- PRED CORR (prob) sólido: verde=navegável, vermelho=bloqueado + prob = np.clip(prob2, 0.0, 1.0).astype(np.float32) + g = (prob * 255.0).astype(np.uint8) + r = ((1.0 - prob) * 255.0).astype(np.uint8) + b = np.zeros_like(g, dtype=np.uint8) + prob_solid = cv2.merge([b, g, r]) # BGR + + # (opcional) desenha contorno do binário por cima do prob + edges = cv2.Canny(bin2, 50, 150) + prob_solid[edges > 0] = (255, 255, 255) + + # --- PRED CORR (bin) sólido: branco=navegável, preto=bloqueado + bin2_solid = cv2.cvtColor(bin2, cv2.COLOR_GRAY2BGR) + + if gt2 is not None: + gt2_vis = (gt2.copy()).astype(np.uint8) + gt2_vis[gt2_vis == 1] = 255 + gt2_bgr = cv2.cvtColor(gt2_vis, cv2.COLOR_GRAY2BGR) + gt2_ov = overlay_mask(img_bgr, gt2_bgr, 0.35) + else: + gt2_ov = img_bgr.copy() + + # monta grid + tile_w = 520 + tile_h = int(tile_w * h / w) + + def fit(im): + return cv2.resize(im, (tile_w, tile_h), interpolation=cv2.INTER_AREA) + + row1 = np.concatenate([fit(img_bgr), fit(gt1_ov), fit(pred1_ov)], axis=1) + row2 = np.concatenate([fit(gt2_ov), fit(prob_solid), fit(bin2_solid)], axis=1) + panel = np.concatenate([row1, row2], axis=0) + + hud = [ + f"{sample.group_name}/{sample.filename}", + f"Corr thr={corridor_thr:.2f} | A/D navega | Q sai", + "Topo: IMG | GT SEG | PRED SEG", + "Baixo: GT CORR | PRED CORR (prob) | PRED CORR (bin)" + ] + panel = put_hud(panel, hud) + return panel + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ckpt", type=str, default=None, help="Caminho do .pt (se não informar, usa best_miou.pt)") + ap.add_argument("--split", default="test", choices=["test", "val"], help="qual split usar") + ap.add_argument("--input_size", type=int, default=512) + ap.add_argument("--corr_thr", type=float, default=0.5) + args = ap.parse_args() + + device = find_device() + print("Device:", device) + + with open("config.json", "r") as f: + config = json.load(f) + MODELO = config["camera"] + MODEL_NAME = config["model_name"] # ex: "segformer_b3" + BACKBONE = config["backbone"] + modelo_folder = config["modelo"] # ex: "dif" + dataset_path = os.path.join(MODELO, "dataset") + + split_root = os.path.join(dataset_path, "split", args.split) + if not os.path.isdir(split_root): + print(f"[WARN] split/{args.split} não existe. Vou usar split/val.") + split_root = os.path.join(dataset_path, "split", "val") + + samples = discover_samples(split_root) + print(f"[INFO] samples: {len(samples)} | split_root={split_root}") + + labelmap_path = os.path.join(dataset_path, "labelmap.txt") # ajusta se seu arquivo estiver em outro lugar + lut_bgr, id_to_name, ignore_id = load_labelmap_lut_bgr(labelmap_path) + print("[labelmap]", id_to_name) + + # ---- Descobre checkpoint ---- + if args.ckpt is not None: + ckpt_path = args.ckpt + else: + # Caminho padrão igual ao treino: + # save_path = MODELO/backup/modelo/model_name/raw4 + save_path = os.path.join(MODELO, "backup", modelo_folder, f"{MODEL_NAME}_dual") + ckpt_path = os.path.join(save_path, "best_miou.pt") + + if not os.path.isfile(ckpt_path): + raise SystemExit(f"Checkpoint não encontrado em: {ckpt_path}") + + print(f"[model] ckpt = {ckpt_path}") + + num_classes = len(id_to_name) + base_model, corridor_head = build_models(num_classes=num_classes, device=device, backbone=BACKBONE) + meta, corridor_head = load_dual_checkpoint(ckpt_path, base_model, corridor_head, device=device) + print(f"... epoch={meta.get('epoch')} best_miou={meta.get('best_miou')} best_main={meta.get('best_main_iou')} best_corr={meta.get('best_corr_iou')}") + + idx = 0 + win = "SegFormer Dual Viewer" + cv2.namedWindow(win, cv2.WINDOW_NORMAL) + + while True: + s = samples[idx] + img = imread_bgr(s.img_path) + + gt1 = imread_gray(s.mask_path) if s.mask_path else None + gt2 = imread_gray(s.mask2_path) if s.mask2_path else None + + pred1, prob2, bin2 = infer_dual( + base_model=base_model, + corridor_head=corridor_head, + img_bgr=img, + device=device, + input_size=args.input_size, + corridor_thr=args.corr_thr + ) + + def dbg_mask(name, m): + if m is None: + print(f"[{name}] None") + else: + u = np.unique(m) + print(f"[{name}] shape={m.shape} dtype={m.dtype} unique={u[:20]}{'...' if len(u)>20 else ''}") + + #dbg_mask("gt1(seg)", gt1) + #dbg_mask("gt2(corr)", gt2) + + panel = make_panel( + sample=s, + lut_bgr=lut_bgr, + img_bgr=img, + gt1=gt1, + pred1=pred1, + gt2=gt2, + prob2=prob2, + bin2=bin2, + corridor_thr=args.corr_thr + ) + + cv2.imshow(win, panel) + k = cv2.waitKey(0) & 0xFF + + # Q / ESC + if k in (ord('q'), 27): + break + + # A (volta) + if k in (ord('a'), ord('A')): + idx = (idx - 1) % len(samples) + continue + + # D (avança) + if k in (ord('d'), ord('D')): + idx = (idx + 1) % len(samples) + continue + + # W/S muda threshold + if k in (ord('w'), ord('W')): + args.corr_thr = min(0.95, args.corr_thr + 0.05) + continue + if k in (ord('s'), ord('S')): + args.corr_thr = max(0.05, args.corr_thr - 0.05) + continue + + cv2.destroyAllWindows() + + +if __name__ == "__main__": + main() diff --git a/Python/OAK/datasets/config.json b/Python/OAK/datasets/config.json index a202fffb9..41f04133b 100644 --- a/Python/OAK/datasets/config.json +++ b/Python/OAK/datasets/config.json @@ -2,6 +2,7 @@ "camera": "gal5000", "modelo": "segformer_b0", "model_name": "ndvi_big", + "dual_head": false, "main_class_name": "erva", "es_classes": "", "model_to_use": "geral", diff --git a/Python/OAK/datasets/config_oak.json b/Python/OAK/datasets/config_oak.json new file mode 100644 index 000000000..adb90c7f1 --- /dev/null +++ b/Python/OAK/datasets/config_oak.json @@ -0,0 +1,17 @@ +{ + "camera": "oak-d", + "modelo": "segformer_b0", + "model_name": "nav", + "dual_head": false, + "main_class_name": "navegavel", + "es_classes": "", + "model_to_use": "geral", + "raw_size": [1296, 1028], + "resolucao": [1024, 576], + "roi_inicio": 0.0, + "roi_tamanho": 1.0, + "shaves": 3, + "channels": 3, + "use_ndvi": true, + "backbone": "nvidia/segformer-b0-finetuned-ade-512-512" +} \ No newline at end of file diff --git a/Python/OAK/datasets/gal5000/dataset/labelmap.txt b/Python/OAK/datasets/gal5000/dataset/labelmap.txt new file mode 100644 index 000000000..2436fef30 --- /dev/null +++ b/Python/OAK/datasets/gal5000/dataset/labelmap.txt @@ -0,0 +1,4 @@ +# label:color_rgb:parts:actions +naopulverizar:128,0,0:: +pulverizar:0,128,0:: +ignore:255,255,255:: \ No newline at end of file diff --git a/Python/OAK/datasets/test_fps.py b/Python/OAK/datasets/test_fps.py new file mode 100644 index 000000000..5d4df2366 --- /dev/null +++ b/Python/OAK/datasets/test_fps.py @@ -0,0 +1,113 @@ +from collections import deque +import json +import os +import time +import cv2 +import numpy as np +import torch +from gal5000.gal_service import Gal5000Camera +from raw_segformer_service import RawSegformerService + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +print(f"Device: {device}") + +INFERIR = False + +cfg_path = r"config.json" + +# ---- Carrega config ---- +with open(cfg_path, "r") as f: + config = json.load(f) + +MODELO = config["camera"] # ex: "gal5000" +MODEL_NAME = config["model_name"] # ex: "segformer_b3" +modelo_folder = config["modelo"] # ex: "dif" + +USE_NDVI = bool(config.get("use_ndvi", False)) +CHANNELS = int(config.get("channels", 4)) +RESOLUCAO = config["resolucao"] +W, H = RESOLUCAO[0], RESOLUCAO[1] +preview_w = RESOLUCAO[0] +preview_h = RESOLUCAO[1] +fps_window = 30 + +if INFERIR: + save_path = os.path.join(MODELO, "backup", modelo_folder, MODEL_NAME, f"raw{CHANNELS}") + ckpt_path = os.path.join(save_path, "best_miou.pt") + + # ---- Instancia service do modelo RAW (4 ou 5 canais) ---- + model_svc = RawSegformerService( + config_path=cfg_path, + ckpt_path=ckpt_path, + device=device, + use_amp=True, + ) + +print("[mode] Câmera ao vivo (GAL5000 + SegFormer RAW)") +cam = Gal5000Camera(raw_w=W, raw_h=H, use_auto_exposure=True) + +win = "GAL5000 RAW + SegFormer (Q=quit)" +cv2.namedWindow(win, cv2.WINDOW_NORMAL) + +tq = deque(maxlen=max(5, int(fps_window))) + +with cam: + #cam.configure_scaling_and_fps(bin_h=2, bin_v=2, dec_h=1, dec_v=1, fps=20.0) + cam.configure_fps(fps_window) + cam.start_streaming() + + while True: + t0 = time.time() + + # raw4_base: (4,H,W) float32 0..1 [R,G,IR,B] já redimensionado + raw4_base, dbg = cam.grab_raw4(out_h=preview_h, out_w=preview_w, timeout_ms=2000) + + # Separa canais básicos + r = raw4_base[0] + g = raw4_base[1] + ir = raw4_base[2] + b = raw4_base[3] + + if INFERIR: + # Deixa o service montar a entrada final com 4 ou 5 canais + # (R,G,IR,B) | (R,G,B,NDVI) | (R,G,B,IR,NDVI) + raw_input = model_svc.build_raw_input(r, g, ir, b) # (C,H,W) float32 0..1 + # Inferência + previews via service + pred_ids, rgb, pred_rgb, overlay, t_inf, t_pvw = model_svc.infer_and_preview(raw_input) + else: + rgb = np.stack([r, g, b], axis=0) # (3,H,W) + rgb = (rgb * 255.0).clip(0, 255).astype("uint8") + overlay = np.transpose(rgb, (1, 2, 0)) # (H,W,3) -> formato OpenCV + + overlay = np.ascontiguousarray(overlay) + if overlay.dtype != np.uint8: + overlay = overlay.astype(np.uint8) + + tq.append(time.time() - t0) + fps = 1.0 / (sum(tq) / len(tq)) + + status = cam.get_status() + + shape = dbg["raw_shape"] + lat = dbg["latency_s"] * 1000 + tc = dbg.get("t_capture", 0) * 1000 + tconv = dbg.get("t_convert", 0) * 1000 + t_ae = dbg.get("t_ae", 0) * 1000 + + lines = [ + f"{preview_w}x{preview_h} | FPS~{fps:.1f} | {shape}", + f"EX={status['exp_raw']} G={status['gain_a']}/{status['gain_d']}", + f"C={CHANNELS} NDVI={int(USE_NDVI)}", + f"cap={tc:.1f}ms ae={t_ae:.1f}ms conv={tconv:.1f}ms", + ] + y = 24 + for line in lines: + cv2.putText(overlay, line, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA) + y += 22 + + cv2.imshow(win, cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR)) + k = cv2.waitKey(1) & 0xFF + if k in (ord("q"), ord("Q"), 27): + break + +cv2.destroyAllWindows() \ No newline at end of file diff --git a/Python/OAK/detect_corridor.py b/Python/OAK/detect_corridor.py new file mode 100644 index 000000000..9a0c4a59d --- /dev/null +++ b/Python/OAK/detect_corridor.py @@ -0,0 +1,2210 @@ +import cv2 +import numpy as np +import argparse +import os +from pathlib import Path + + +def list_images(folder: str, exts=(".png", ".jpg", ".jpeg", ".bmp")): + p = Path(folder) + if not p.exists() or not p.is_dir(): + raise FileNotFoundError(f"Pasta inválida: {folder}") + files = [x for x in p.rglob("*") if x.suffix.lower() in exts] + files.sort() + return files + + +def color_mask_bgr(img_bgr: np.ndarray, target_bgr: tuple[int, int, int], tol: int = 40) -> np.ndarray: + """ + Retorna máscara binária (uint8 0/255) onde a cor está próxima de target_bgr. + """ + b, g, r = target_bgr + lower = np.array([max(0, b - tol), max(0, g - tol), max(0, r - tol)], dtype=np.uint8) + upper = np.array([min(255, b + tol), min(255, g + tol), min(255, r + tol)], dtype=np.uint8) + mask = cv2.inRange(img_bgr, lower, upper) + return mask + + +def smooth_1d(x: np.ndarray, k: int = 31) -> np.ndarray: + k = int(k) + if k < 3: + return x.astype(np.float32) + if k % 2 == 0: + k += 1 + kernel = np.ones(k, dtype=np.float32) / k + return np.convolve(x.astype(np.float32), kernel, mode="same") + + +def find_best_pair(score_cana: np.ndarray, + score_chao: np.ndarray, + x_center: int, + min_w: int, + max_w: int, + k_peaks: int = 6) -> tuple[int, int, float] | None: + """ + Escolhe (l, r) maximizando uma função objetivo: + J = cana[l] + cana[r] + mean(chao[l:r]) - penalidades + Retorna (l, r, J) onde l < r. + """ + W = score_cana.shape[0] + left_half = score_cana[:x_center] + right_half = score_cana[x_center:] + + # Picos candidatos (índices das maiores colunas) + left_idx = np.argsort(left_half)[-k_peaks:] + right_idx = np.argsort(right_half)[-k_peaks:] + x_center + + best = None + for l in left_idx: + for r in right_idx: + if r <= l + 10: + continue + width = r - l + if width < min_w or width > max_w: + continue + + corridor_chao = float(np.mean(score_chao[l:r])) + J = float(score_cana[l] + score_cana[r] + corridor_chao) + + # Penaliza se o centro do corredor fugir muito do centro esperado + mid = (l + r) // 2 + J -= 0.002 * abs(mid - x_center) + + # Penaliza largura muito nas bordas do range + if width < (min_w * 1.15): + J -= 0.2 + if width > (max_w * 0.9): + J -= 0.2 + + if best is None or J > best[2]: + best = (int(l), int(r), float(J)) + + return best + + +def refine_edges_from_peaks(score_cana: np.ndarray, l_peak: int, r_peak: int, x_center: int, + drop_ratio: float = 0.55) -> tuple[int, int]: + """ + Ajusta as bordas do corredor a partir dos picos de cana. + Para a esquerda: anda do pico em direção ao centro até cair para drop_ratio * pico. + Para a direita: anda do pico em direção ao centro até cair para drop_ratio * pico. + """ + W = score_cana.shape[0] + + l_val = score_cana[l_peak] + r_val = score_cana[r_peak] + + l_thr = l_val * drop_ratio + r_thr = r_val * drop_ratio + + # Borda esquerda (dentro do corredor): começa no pico e vai para a direita + xL = l_peak + for x in range(l_peak, min(x_center, W - 1)): + if score_cana[x] <= l_thr: + xL = x + break + + # Borda direita (dentro do corredor): começa no pico e vai para a esquerda + xR = r_peak + for x in range(r_peak, max(x_center, 0), -1): + if score_cana[x] <= r_thr: + xR = x + break + + if xR <= xL + 10: + # fallback simples + xL = min(xL, x_center - 20) + xR = max(xR, x_center + 20) + + return int(xL), int(xR) + + +def detect_corridor_trapezoid(mask_bgr: np.ndarray, + bands: int = 7, + y0_frac: float = 0.45, + y1_frac: float = 0.95, + smooth_k: int = 31, + min_w_frac: float = 0.25, + max_w_frac: float = 0.85, + tol_color: int = 40, + debug: bool = False): + """ + Detecta trapézio do corredor por projeção em bandas horizontais. + Retorna: + - pts_trap (4x2 int) ou None + - conf (0..1) + - extras dict + """ + img = mask_bgr + H, W = img.shape[:2] + + # Cores alvo em BGR (OpenCV) + CHAO_BGR = (0, 0, 128) # vermelho + CANA_BGR = (0, 128, 0) # verde + OBST_BGR = (128, 0, 0) # azul + + m_chao = color_mask_bgr(img, CHAO_BGR, tol=tol_color) # 0/255 + m_cana = color_mask_bgr(img, CANA_BGR, tol=tol_color) + m_obst = color_mask_bgr(img, OBST_BGR, tol=tol_color) + m_known = ((m_chao > 0) | (m_cana > 0) | (m_obst > 0)).astype(np.uint8) * 255 + m_ignore = cv2.bitwise_not(m_known) # 255 onde é "branco/ignore" + + # Considera obstáculo como "não chão" na métrica de corredor (mas não impede cana) + # score_chao vai considerar só pixels vermelhos mesmo. + y0 = int(H * y0_frac) + y1 = int(H * y1_frac) + if y1 <= y0 + 20: + y0 = int(H * 0.5) + y1 = int(H * 0.95) + + x_center = W // 2 + min_w = int(W * min_w_frac) + max_w = int(W * max_w_frac) + + band_edges = [] # (y_mid, xL, xR, J) + bands_dbg = [] # lista de dicts para debug, não interfere no resto + band_h = max(10, (y1 - y0) // bands) + + for bi in range(bands): + ya = y0 + bi * band_h + yb = min(y1, ya + band_h) + y_mid = (ya + yb) // 2 + if yb <= ya + 5: + continue + + # recorte da banda + chao_band = (m_chao[ya:yb, :] > 0).astype(np.float32) + cana_band = (m_cana[ya:yb, :] > 0).astype(np.float32) + ignore_band = (m_ignore[ya:yb, :] > 0).astype(np.float32) + + # score por coluna = fração de pixels na banda + score_chao = np.mean(chao_band, axis=0) + score_cana = np.mean(cana_band, axis=0) + score_ignore = np.mean(ignore_band, axis=0) + + # suaviza + score_chao_s = smooth_1d(score_chao, smooth_k) + score_cana_s = smooth_1d(score_cana, smooth_k) + + best = find_best_pair(score_cana_s, score_chao_s, x_center, min_w, max_w, k_peaks=7) + if best is None: + continue + + l_peak, r_peak, J = best + xL, xR = refine_edges_from_peaks(score_cana_s, l_peak, r_peak, x_center, drop_ratio=0.55) + + if debug: + bands_dbg.append({ + "y_mid": int(y_mid), + "xL": int(xL), + "xR": int(xR), + "l_peak": int(l_peak), + "r_peak": int(r_peak), + "J": float(J), + "corridor_chao": float(np.mean(score_chao_s[xL:xR])) if xR > xL else 0.0, + "ok": True, + "ya": int(ya), + "yb": int(yb), + }) + + # validações rápidas + width = xR - xL + if width < min_w or width > max_w: + if debug: + bands_dbg.append({ + "y_mid": int(y_mid), + "xL": None, + "xR": None, + "l_peak": int(l_peak), + "r_peak": int(r_peak), + "J": float(J), + "corridor_chao": 0.0, + "ok": False, + "ya": int(ya), + "yb": int(yb), + }) + continue + + # miolo precisa ter chão razoável (senão pode ser engano) + corridor_chao = float(np.mean(score_chao_s[xL:xR])) if xR > xL else 0.0 + if corridor_chao < 0.10: + if debug: + bands_dbg.append({ + "y_mid": int(y_mid), + "xL": None, + "xR": None, + "l_peak": int(l_peak), + "r_peak": int(r_peak), + "J": float(J), + "corridor_chao": 0.0, + "ok": False, + "ya": int(ya), + "yb": int(yb), + }) + continue + + # corredor não pode ser "vazio" / ignore demais no miolo + ignore_in_corridor = float(np.mean(score_ignore[xL:xR])) if xR > xL else 1.0 + if ignore_in_corridor > 0.10: + if debug: + bands_dbg.append({ + "y_mid": int(y_mid), + "xL": None, + "xR": None, + "l_peak": int(l_peak), + "r_peak": int(r_peak), + "J": float(J), + "corridor_chao": 0.0, + "ok": False, + "ya": int(ya), + "yb": int(yb), + }) + continue + + band_edges.append((int(y_mid), int(xL), int(xR), float(J))) + + if len(band_edges) < max(2, bands // 3): + return None, 0.0, {"reason": "poucas bandas válidas", "bands_ok": len(band_edges)} + + # Remove outliers por largura e por centro + band_edges = sorted(band_edges, key=lambda t: t[0]) + widths = np.array([xR - xL for _, xL, xR, _ in band_edges], dtype=np.float32) + centers = np.array([(xL + xR) / 2.0 for _, xL, xR, _ in band_edges], dtype=np.float32) + + w_med = float(np.median(widths)) + c_med = float(np.median(centers)) + + filtered = [] + for (y, xL, xR, J) in band_edges: + w = xR - xL + c = (xL + xR) / 2.0 + if abs(w - w_med) > 0.35 * w_med: + if debug: + bands_dbg.append({ + "y_mid": int(y_mid), + "xL": None, + "xR": None, + "l_peak": int(l_peak), + "r_peak": int(r_peak), + "J": float(J), + "corridor_chao": 0.0, + "ok": False, + "ya": int(ya), + "yb": int(yb), + }) + continue + if abs(c - c_med) > 0.20 * W: + if debug: + bands_dbg.append({ + "y_mid": int(y_mid), + "xL": None, + "xR": None, + "l_peak": int(l_peak), + "r_peak": int(r_peak), + "J": float(J), + "corridor_chao": 0.0, + "ok": False, + "ya": int(ya), + "yb": int(yb), + }) + continue + filtered.append((y, xL, xR, J)) + + if len(filtered) < 2: + return None, 0.0, {"reason": "outliers demais", "bands_ok": len(band_edges), "bands_filt": len(filtered)} + + ys = np.array([y for y, _, _, _ in filtered], dtype=np.float32) + xLs = np.array([xL for _, xL, _, _ in filtered], dtype=np.float32) + xRs = np.array([xR for _, _, xR, _ in filtered], dtype=np.float32) + + # Fit linear x = a*y + b para esquerda e direita + aL, bL = np.polyfit(ys, xLs, 1) + aR, bR = np.polyfit(ys, xRs, 1) + + y_top = int(ys[0]) + y_bot = int(ys[-1]) + + xL_top = int(aL * y_top + bL) + xL_bot = int(aL * y_bot + bL) + xR_top = int(aR * y_top + bR) + xR_bot = int(aR * y_bot + bR) + + # clamp + xL_top = int(np.clip(xL_top, 0, W - 1)) + xL_bot = int(np.clip(xL_bot, 0, W - 1)) + xR_top = int(np.clip(xR_top, 0, W - 1)) + xR_bot = int(np.clip(xR_bot, 0, W - 1)) + + if xR_top <= xL_top + 10 or xR_bot <= xL_bot + 10: + return None, 0.0, {"reason": "trapézio degenerado"} + + pts = np.array([ + [xL_top, y_top], + [xR_top, y_top], + [xR_bot, y_bot], + [xL_bot, y_bot] + ], dtype=np.int32) + + # Confiança simples: proporção de bandas boas + estabilidade de centro/largura + bands_ok = len(filtered) + bands_total = bands + conf = bands_ok / float(max(1, bands_total)) + + w_std = float(np.std([xR - xL for _, xL, xR, _ in filtered])) + c_std = float(np.std([(xL + xR) / 2.0 for _, xL, xR, _ in filtered])) + conf *= float(np.clip(1.0 - (w_std / (0.35 * W)), 0.0, 1.0)) + conf *= float(np.clip(1.0 - (c_std / (0.25 * W)), 0.0, 1.0)) + + extras = { + "bands_ok": bands_ok, + "bands_total": bands_total, + "y_top": y_top, + "y_bot": y_bot, + "fit_left": (float(aL), float(bL)), + "fit_right": (float(aR), float(bR)), + "conf": float(conf), + } + + if debug: + extras["filtered"] = filtered + extras["bands_dbg"] = bands_dbg + extras["dbg_params"] = {"y0_frac": y0_frac, "y1_frac": y1_frac, "bands": bands} + + return pts, float(conf), extras + + +def draw_trapezoid(img_bgr: np.ndarray, pts: np.ndarray, conf: float): + out = img_bgr.copy() + cv2.polylines(out, [pts], isClosed=True, color=(255, 255, 0), thickness=3) # amarelo/ciano + cx = int(np.mean(pts[:, 0])) + cy = int(np.mean(pts[:, 1])) + cv2.putText(out, f"corridor conf={conf:.2f}", (max(10, cx - 120), max(30, cy)), + cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2, cv2.LINE_AA) + return out + + +def draw_band_debug(img_bgr: np.ndarray, extras: dict, y0_frac: float, y1_frac: float, bands: int): + out = img_bgr.copy() + H, W = out.shape[:2] + + bands_dbg = extras.get("bands_dbg", []) + if not bands_dbg: + return out + + # desenha faixas horizontais + y0 = int(H * y0_frac) + y1 = int(H * y1_frac) + band_h = max(10, (y1 - y0) // bands) + + for bi in range(bands): + ya = y0 + bi * band_h + yb = min(y1, ya + band_h) + if yb <= ya: + continue + cv2.line(out, (0, ya), (W-1, ya), (80, 80, 80), 1) + + # desenha pontos por banda + for d in bands_dbg: + y = int(d["y_mid"]) + ok = bool(d.get("ok", True)) + + # cor: verde ok, vermelho rejeitada + color = (0, 255, 0) if ok else (0, 0, 255) + + # marca picos (cristas cana) + lp = d.get("l_peak", None) + rp = d.get("r_peak", None) + if lp is not None: + cv2.circle(out, (int(lp), y), 5, (0, 255, 255), -1) # amarelo + if rp is not None: + cv2.circle(out, (int(rp), y), 5, (0, 255, 255), -1) + + xL = d.get("xL", None) + xR = d.get("xR", None) + + if xL is not None and xR is not None: + # bordas + cv2.circle(out, (int(xL), y), 6, color, -1) + cv2.circle(out, (int(xR), y), 6, color, -1) + # centro da banda + xc = int((xL + xR) / 2) + cv2.circle(out, (xc, y), 5, (255, 0, 255), -1) # roxo + # linha do corredor nessa banda + cv2.line(out, (int(xL), y), (int(xR), y), color, 2) + + # label curtinho + if ok and xL is not None and xR is not None: + txt = f"J={d.get('J',0):.2f} ch={d.get('corridor_chao',0):.2f}" + else: + txt = f"rej" + cv2.putText(out, txt, (10, max(20, y-5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA) + + return out + + +def detect_corridor_regime(mask_bgr: np.ndarray, + bands: int = 15, + y0_frac: float = 0.1, + y1_frac: float = 0.9, + tol_color: int = 40, + smooth_k: int = 51, + cane_peak_min: float = 0.12, + notfloor_thr: float = 0.25, + floor_open_thr: float = 0.70, + cane_present_thr: float = 0.06): + """ + Classifica o "regime" do frame, sem calcular trapézio. + Retorna: (regime_str, metrics_dict) + + Regimes: + - OPEN_FIELD: só chão (sem paredes) em quase todas as bandas + - ENTERING: topo tem 2 paredes, base não tem + - EXITING: base tem 2 paredes, topo não tem + - IN_CORRIDOR: topo e base com 2 paredes + - ONE_WALL: muitas bandas com 1 parede, poucas com 2 + - UNKNOWN: não bateu com regras + """ + img = mask_bgr + H, W = img.shape[:2] + x_center = W // 2 + + # Cores alvo em BGR + CHAO_BGR = (0, 0, 128) # vermelho + CANA_BGR = (0, 128, 0) # verde + OBST_BGR = (128, 0, 0) # azul + + m_chao = color_mask_bgr(img, CHAO_BGR, tol=tol_color) # 0/255 + m_cana = color_mask_bgr(img, CANA_BGR, tol=tol_color) + m_obst = color_mask_bgr(img, OBST_BGR, tol=tol_color) + m_known = ((m_chao > 0) | (m_cana > 0) | (m_obst > 0)).astype(np.uint8) * 255 + m_ignore = cv2.bitwise_not(m_known) + + y0 = int(H * y0_frac) + y1 = int(H * y1_frac) + if y1 <= y0 + 20: + y0 = int(H * 0.5) + y1 = int(H * 0.95) + + band_h = max(10, (y1 - y0) // max(1, bands)) + + per_band = [] + n0 = n1 = n2 = 0 + + # define "top" e "bottom" pela metade das bandas válidas + for bi in range(bands): + ya = y0 + bi * band_h + yb = min(y1, ya + band_h) + if yb <= ya + 5: + continue + y_mid = (ya + yb) // 2 + + chao_band = (m_chao[ya:yb, :] > 0).astype(np.float32) + cana_band = (m_cana[ya:yb, :] > 0).astype(np.float32) + ignore_band = (m_ignore[ya:yb, :] > 0).astype(np.float32) + obst_band = (m_obst[ya:yb, :] > 0).astype(np.float32) + + score_chao = smooth_1d(np.mean(chao_band, axis=0), smooth_k) + score_cana = smooth_1d(np.mean(cana_band, axis=0), smooth_k) + score_ign = np.mean(ignore_band, axis=0) + score_obst = np.mean(obst_band, axis=0) + + # métricas simples + floor_ratio = float(np.mean(score_chao)) # fração média de chão por coluna + cane_ratio = float(np.mean(score_cana)) + ignore_ratio = float(np.mean(score_ign)) + obst_ratio = float(np.mean(score_obst)) + + # "parede" = pico de cana suficientemente alto em cada metade + left_peak = float(np.max(score_cana[:x_center])) if x_center > 5 else 0.0 + right_peak = float(np.max(score_cana[x_center:])) if (W - x_center) > 5 else 0.0 + + has_left = left_peak >= cane_peak_min + has_right = right_peak >= cane_peak_min + + walls = int(has_left) + int(has_right) + + # sanity: se quase tudo é ignore/obst, essa banda é ruim + bad_band = (ignore_ratio > 0.35) or (obst_ratio > 0.25) + + # contagem final só se banda não estiver "podre" + if not bad_band: + if walls == 0: + n0 += 1 + elif walls == 1: + n1 += 1 + else: + n2 += 1 + + per_band.append({ + "y_mid": y_mid, + "walls": walls, + "bad": bad_band, + "floor_ratio": floor_ratio, + "cane_ratio": cane_ratio, + "left_peak": left_peak, + "right_peak": right_peak, + "ignore_ratio": ignore_ratio, + "obst_ratio": obst_ratio, + }) + + valid_bands = [b for b in per_band if not b["bad"]] + if len(valid_bands) < max(3, bands // 3): + return "UNKNOWN", {"reason": "poucas bandas válidas", "n_valid": len(valid_bands), "n0": n0, "n1": n1, "n2": n2} + + half = max(1, len(valid_bands) // 2) + top_bands = valid_bands[:half] + bot_bands = valid_bands[half:] + + def count_walls(bands_list): + c0 = sum(1 for b in bands_list if b["walls"] == 0) + c1 = sum(1 for b in bands_list if b["walls"] == 1) + c2 = sum(1 for b in bands_list if b["walls"] == 2) + return c0, c1, c2 + + t0, t1, t2 = count_walls(top_bands) + b0, b1, b2 = count_walls(bot_bands) + + # médias globais (ajudam a decidir OPEN_FIELD) + mean_floor = float(np.mean([b["floor_ratio"] for b in valid_bands])) if valid_bands else 0.0 + mean_cane = float(np.mean([b["cane_ratio"] for b in valid_bands])) if valid_bands else 0.0 + + # Regras principais + # 1) OPEN_FIELD: chão alto, cana baixa, quase ninguém com 2 paredes + if mean_floor >= floor_open_thr and mean_cane <= cane_present_thr and n2 <= max(1, len(valid_bands) // 10): + return "OPEN_FIELD", { + "n_valid": len(valid_bands), "n0": n0, "n1": n1, "n2": n2, + "top": (t0, t1, t2), "bot": (b0, b1, b2), + "mean_floor": mean_floor, "mean_cane": mean_cane + } + + # 2) IN_CORRIDOR: topo e base com várias bandas de 2 paredes + if t2 >= max(2, len(top_bands) // 3) and b2 >= max(2, len(bot_bands) // 3): + return "IN_CORRIDOR", { + "n_valid": len(valid_bands), "n0": n0, "n1": n1, "n2": n2, + "top": (t0, t1, t2), "bot": (b0, b1, b2), + "mean_floor": mean_floor, "mean_cane": mean_cane + } + + # 3) ENTERING: topo tem 2 paredes, base não + if t2 >= max(2, len(top_bands) // 3) and b2 <= max(1, len(bot_bands) // 6) and b0 >= max(2, len(bot_bands) // 3): + return "ENTERING", { + "n_valid": len(valid_bands), "n0": n0, "n1": n1, "n2": n2, + "top": (t0, t1, t2), "bot": (b0, b1, b2), + "mean_floor": mean_floor, "mean_cane": mean_cane + } + + # 4) EXITING: base tem 2 paredes, topo não + if b2 >= max(2, len(bot_bands) // 3) and t2 <= max(1, len(top_bands) // 6) and t0 >= max(2, len(top_bands) // 3): + return "EXITING", { + "n_valid": len(valid_bands), "n0": n0, "n1": n1, "n2": n2, + "top": (t0, t1, t2), "bot": (b0, b1, b2), + "mean_floor": mean_floor, "mean_cane": mean_cane + } + + # 5) ONE_WALL: domina 1 parede, e 2 paredes é raro + if n1 >= max(3, len(valid_bands) // 3) and n2 <= max(1, len(valid_bands) // 6): + return "ONE_WALL", { + "n_valid": len(valid_bands), "n0": n0, "n1": n1, "n2": n2, + "top": (t0, t1, t2), "bot": (b0, b1, b2), + "mean_floor": mean_floor, "mean_cane": mean_cane + } + + return "UNKNOWN", { + "n_valid": len(valid_bands), "n0": n0, "n1": n1, "n2": n2, + "top": (t0, t1, t2), "bot": (b0, b1, b2), + "mean_floor": mean_floor, "mean_cane": mean_cane + } + +def _safe_smooth_k(W: int, k: int) -> int: + # k ímpar e <= W (e pelo menos 3) + k = int(k) + if W <= 3: + return 3 + k = min(k, W if (W % 2 == 1) else (W - 1)) + if k < 3: + k = 3 + if k % 2 == 0: + k -= 1 + return k + +def _clip01(x: float) -> float: + return 0.0 if x < 0.0 else (1.0 if x > 1.0 else x) + +def _run_length_from_bottom(vec: list[int], value: int) -> int: + k = 0 + for v in reversed(vec): + if v == value: + k += 1 + else: + break + return k + +def _longest_run_geq(vec: np.ndarray, thr: float) -> int: + """Maior run-length contínuo onde vec >= thr.""" + if vec.size == 0: + return 0 + m = (vec >= thr).astype(np.uint8) + best = cur = 0 + for v in m: + if v: + cur += 1 + best = max(best, cur) + else: + cur = 0 + return int(best) + +def band_has_corridor_valley(score_cana: np.ndarray, + score_chao: np.ndarray, + cane_thr: float = 0.12, + floor_min_ratio: float = 0.40, # teu limiar + valley_min_width_frac: float = 0.18, + peaks_min_sep_frac: float = 0.22, + + # parede real (anti-folha) + wall_run_thr: float = 0.10, + wall_min_run_frac: float = 0.04, + wall_side_margin_frac: float = 0.10, + wall_min_peak: float = 0.12, + + # opcional: evita caso “miolo totalmente verde” + valley_cane_max_ratio: float = 0.90): + """ + Corredor por 'vale': + - duas paredes (run-length em regiões laterais) + - vale largo entre elas + - no vale: chão >= floor_min_ratio (não precisa dominar) + """ + + W = int(score_cana.shape[0]) + if W < 20: + return False, {"reason": "small_W"} + + sc = np.clip(score_cana, 0, 1) + sf = np.clip(score_chao, 0, 1) + + # 1) dois picos separados (só pra achar L e R) + i1 = int(np.argmax(sc)) + p1 = float(sc[i1]) + + min_sep = int(W * peaks_min_sep_frac) + mask_far = np.ones(W, dtype=bool) + mask_far[max(0, i1 - min_sep): min(W, i1 + min_sep)] = False + if not np.any(mask_far): + return False, {"reason": "no_far_region"} + + i2 = int(np.argmax(sc * mask_far)) + p2 = float(sc[i2]) + + if (p1 < cane_thr) or (p2 < cane_thr): + return False, {"reason": "peaks_weak", "p1": p1, "p2": p2} + + L = min(i1, i2) + R = max(i1, i2) + + valley_min_w = int(W * valley_min_width_frac) + if (R - L) < valley_min_w: + return False, {"reason": "valley_too_narrow", "valley_w": (R-L), "valley_min_w": valley_min_w} + + # 2) validar “paredes” com run-length (anti-folha) + margin = int(W * wall_side_margin_frac) + + left_end = max(1, min(L+1, W - margin)) + right_start = min(max(R, margin), W-1) + + left_region = sc[:left_end] + right_region = sc[right_start:] + + min_run = int(W * wall_min_run_frac) + left_run = _longest_run_geq(left_region, wall_run_thr) + right_run = _longest_run_geq(right_region, wall_run_thr) + + left_peak = float(np.max(left_region)) if left_region.size else 0.0 + right_peak = float(np.max(right_region)) if right_region.size else 0.0 + + ok_left_wall = (left_run >= min_run) and (left_peak >= wall_min_peak) + ok_right_wall = (right_run >= min_run) and (right_peak >= wall_min_peak) + + wall_dbg = { + "left_run": left_run, "right_run": right_run, "min_run": min_run, + "left_peak": left_peak, "right_peak": right_peak, + "ok_left_wall": bool(ok_left_wall), + "ok_right_wall": bool(ok_right_wall), + } + + if not (ok_left_wall and ok_right_wall): + return False, { + "reason": "thin_walls", + "left_run": left_run, "right_run": right_run, "min_run": min_run, + "left_peak": left_peak, "right_peak": right_peak, + "iL": L, "iR": R, + **wall_dbg + } + + # 3) medir o vale por proporção de chão + gap = slice(L+1, R-1) + if (R - L) <= 3: + return False, {"reason": "valley_too_small", **wall_dbg} + + valley_floor_ratio = float(np.mean(sf[gap] >= 0.50)) # chão “forte” (ajustável) + valley_cane_ratio = float(np.mean(sc[gap] >= 0.50)) # cana “forte” (debug) + + ok_valley_floor = (valley_floor_ratio >= floor_min_ratio) + ok_valley_cane = (valley_cane_ratio <= valley_cane_max_ratio) + + if not (ok_valley_floor and ok_valley_cane): + return False, { + "reason": "valley_fail", + "valley_floor_ratio": valley_floor_ratio, + "valley_cane_ratio": valley_cane_ratio, + "iL": L, "iR": R, + "pL": float(sc[L]), "pR": float(sc[R]), + **wall_dbg + } + + center = 0.5 * (L + R) / max(1, (W-1)) + width = float(R - L) / max(1, W) + + return True, { + "iL": L, "iR": R, + "pL": float(sc[L]), "pR": float(sc[R]), + "valley_floor_ratio": valley_floor_ratio, + "valley_cane_ratio": valley_cane_ratio, + "left_run": left_run, "right_run": right_run, + "center_frac": center, + "width_frac": width, + **wall_dbg + } + +def band_has_corridor(score_cana: np.ndarray, + score_chao: np.ndarray, + cane_thr: float = 0.12, + # gap: chão precisa "ganhar" da cana no trecho central + gap_floor_thr: float = 0.45, + gap_cane_max: float = 0.18, + # largura mínima do gap (fração da largura da imagem) + gap_min_width_frac: float = 0.18, + # separação mínima entre picos (fração da largura) + peaks_min_sep_frac: float = 0.22): + """ + Decide se existe padrão CANA-CHAO-CANA em uma banda (invariante a deslocamento no X). + Retorna (has_corridor, info_dict). + + Estratégia: + - acha índices onde cana é forte (>= cane_thr) + - separa em região esquerda e direita por um "corte" móvel: usa o pico global e procura um segundo pico distante + - garante que entre picos há um trecho com chão alto e cana baixa + """ + W = score_cana.shape[0] + if W < 20: + return False, {"reason": "small_W"} + + # normaliza segurança (não obrigatório, mas ajuda se tiver ruído) + sc = np.clip(score_cana, 0, 1) + sf = np.clip(score_chao, 0, 1) + + # picos candidatos + idx_strong = np.where(sc >= cane_thr)[0] + if idx_strong.size == 0: + return False, {"reason": "no_strong_cane"} + + # pega o pico mais alto como referência + i1 = int(np.argmax(sc)) + p1 = float(sc[i1]) + + # procura segundo pico longe o suficiente + min_sep = int(W * peaks_min_sep_frac) + # mascara pontos longe do i1 + mask_far = np.ones(W, dtype=bool) + mask_far[max(0, i1 - min_sep): min(W, i1 + min_sep)] = False + + if not np.any(mask_far): + return False, {"reason": "no_far_region"} + + i2 = int(np.argmax(sc * mask_far)) + p2 = float(sc[i2]) + + # ambos precisam ser fortes + if p2 < cane_thr: + return False, {"reason": "second_peak_weak", "p1": p1, "p2": p2} + + # ordena esquerda/direita + L = min(i1, i2) + R = max(i1, i2) + + # gap entre picos + if R <= L + 3: + return False, {"reason": "peaks_too_close"} + + gap = slice(L+1, R-1) + if (R - L) <= 3: + return False, {"reason": "peaks_too_close"} + + gap_floor = float(np.mean(sf[gap])) + gap_cane = float(np.mean(sc[gap])) + + gap_min_w = int(W * gap_min_width_frac) + if (R - L) < gap_min_w: + return False, {"reason": "gap_too_narrow", "gap_w": (R-L), "gap_min_w": gap_min_w} + + # critério: no gap o chão domina e cana é baixa + has_gap = (gap_floor >= gap_floor_thr) and (gap_cane <= gap_cane_max) + + return bool(has_gap), { + "iL": L, "iR": R, + "pL": float(sc[L]), "pR": float(sc[R]), + "gap_floor": gap_floor, + "gap_cane": gap_cane, + "gap_w": (R - L), + "has_gap": has_gap + } + +def corridor_bridge_ok(corr_vec: list[int], max_gap_run: int = 4) -> tuple[bool, dict]: + """ + Retorna True se existe conectividade topo->base permitindo buracos (0s) de até max_gap_run. + """ + idx_ones = [i for i, v in enumerate(corr_vec) if v == 1] + if len(idx_ones) < 2: + return False, {"reason": "few_ones", "n_ones": len(idx_ones)} + + # maior gap entre 1s consecutivos + gaps = [] + for a, b in zip(idx_ones[:-1], idx_ones[1:]): + gaps.append((b - a - 1)) # quantos zeros entre eles + + max_gap = max(gaps) if gaps else 0 + ok = (idx_ones[0] <= 1) and (idx_ones[-1] >= len(corr_vec) - 2) and (max_gap <= max_gap_run) + + return ok, {"max_gap": max_gap, "idx_first": idx_ones[0], "idx_last": idx_ones[-1], "n_ones": len(idx_ones)} + +def corridor_center_consistent(per_band_ok: list[dict], + max_center_jump: float = 0.12) -> tuple[bool, dict]: + """ + Exige que o centro do corredor não pule demais entre bandas OK. + """ + centers = [b["center_frac"] for b in per_band_ok if b.get("center_frac") is not None] + if len(centers) < 2: + return True, {"reason": "few_centers"} # não bloqueia + + jumps = [abs(b - a) for a, b in zip(centers[:-1], centers[1:])] + max_jump = max(jumps) if jumps else 0.0 + return (max_jump <= max_center_jump), {"max_center_jump": max_jump, "n_centers": len(centers)} + + +def is_open_field(mask_bgr: np.ndarray, + bands: int = 15, + y0_frac: float = 0.10, + y1_frac: float = 0.90, + tol_color: int = 40, + smooth_k: int = 51, + cane_peak_min: float = 0.12, + floor_open_thr: float = 0.70, + cane_present_thr: float = 0.06, + max_n2_ratio: float = 0.10, # máximo % de bandas com 2 paredes + min_n0_ratio: float = 0.70, # mínimo % de bandas com 0 paredes + top_corridor_veto_ratio: float = 0.25, # veto se topo já tem muito 2-paredes + bad_ignore_thr: float = 0.35, + bad_obst_thr: float = 0.25, + min_valid_bands_ratio: float = 0.33): + """ + Decide APENAS se o frame está em OPEN_FIELD (campo aberto: só chão, sem paredes). + + Retorna: + (is_open: bool, metrics: dict) + + Robustez extra: + - exige chão alto e cana baixa + - exige dominância de 0 paredes (n0 alto) + - limita 2 paredes (n2 baixo) + - VETA OPEN_FIELD se topo já apresenta padrão de corredor (t2 alto) + - trata bandas ruins (ignore/obst) + """ + + img = mask_bgr + H, W = img.shape[:2] + x_center = W // 2 + smooth_k = _safe_smooth_k(W, smooth_k) + + # cores alvo em BGR (ajuste conforme teu mask) + CHAO_BGR = (0, 0, 128) # vermelho + CANA_BGR = (0, 128, 0) # verde + OBST_BGR = (128, 0, 0) # azul + + m_chao = color_mask_bgr(img, CHAO_BGR, tol=tol_color) + m_cana = color_mask_bgr(img, CANA_BGR, tol=tol_color) + m_obst = color_mask_bgr(img, OBST_BGR, tol=tol_color) + m_known = ((m_chao > 0) | (m_cana > 0) | (m_obst > 0)).astype(np.uint8) * 255 + m_ignore = cv2.bitwise_not(m_known) + + # região vertical analisada + y0 = int(H * y0_frac) + y1 = int(H * y1_frac) + if y1 <= y0 + 20: + y0 = int(H * 0.5) + y1 = int(H * 0.95) + + band_h = max(10, (y1 - y0) // max(1, bands)) + + per_band = [] + n0 = n1 = n2 = 0 + + for bi in range(bands): + ya = y0 + bi * band_h + yb = min(y1, ya + band_h) + if yb <= ya + 5: + continue + + chao_band = (m_chao[ya:yb, :] > 0).astype(np.float32) + cana_band = (m_cana[ya:yb, :] > 0).astype(np.float32) + ignore_band = (m_ignore[ya:yb, :] > 0).astype(np.float32) + obst_band = (m_obst[ya:yb, :] > 0).astype(np.float32) + + score_chao = smooth_1d(np.mean(chao_band, axis=0), smooth_k) + score_cana = smooth_1d(np.mean(cana_band, axis=0), smooth_k) + + ignore_ratio = float(np.mean(ignore_band)) + obst_ratio = float(np.mean(obst_band)) + + floor_ratio = float(np.mean(score_chao)) + cane_ratio = float(np.mean(score_cana)) + + left_peak = float(np.max(score_cana[:x_center])) if x_center > 5 else 0.0 + right_peak = float(np.max(score_cana[x_center:])) if (W - x_center) > 5 else 0.0 + + has_left = left_peak >= cane_peak_min + has_right = right_peak >= cane_peak_min + walls = int(has_left) + int(has_right) + + bad_band = (ignore_ratio > bad_ignore_thr) or (obst_ratio > bad_obst_thr) + + if not bad_band: + if walls == 0: + n0 += 1 + elif walls == 1: + n1 += 1 + else: + n2 += 1 + + per_band.append({ + "bi": bi, + "walls": walls, + "bad": bad_band, + "floor_ratio": floor_ratio, + "cane_ratio": cane_ratio, + "left_peak": left_peak, + "right_peak": right_peak, + "ignore_ratio": ignore_ratio, + "obst_ratio": obst_ratio, + }) + + valid = [b for b in per_band if not b["bad"]] + n_valid = len(valid) + min_valid = max(3, int(bands * min_valid_bands_ratio)) + + if n_valid < min_valid: + return False, { + "reason": "few_valid_bands", + "n_valid": n_valid, "min_valid": min_valid, + "n0": n0, "n1": n1, "n2": n2 + } + + # split topo/base + mid = bands / 2.0 + top = [b for b in valid if b["bi"] < mid] + bot = [b for b in valid if b["bi"] >= mid] + # se por algum motivo um lado ficou vazio (muito ruído), fallback simples: + if len(top) == 0 or len(bot) == 0: + half = max(1, n_valid // 2) + top = valid[:half] + bot = valid[half:] + + t2 = sum(1 for b in top if b["walls"] == 2) + b2 = sum(1 for b in bot if b["walls"] == 2) + + mean_floor = float(np.mean([b["floor_ratio"] for b in valid])) + mean_cane = float(np.mean([b["cane_ratio"] for b in valid])) + + n0_ratio = n0 / max(1, n_valid) + n2_ratio = n2 / max(1, n_valid) + + # VETO: se topo já tem “cara de corredor”, não é campo aberto + # (mesmo que a base ainda seja chão) + top_veto_thr = max(2, int(len(top) * top_corridor_veto_ratio)) + top_has_corridor = (t2 >= top_veto_thr) + + # regra base OPEN_FIELD robusta + ok_floor = (mean_floor >= floor_open_thr) + ok_cane = (mean_cane <= cane_present_thr) + ok_n2 = (n2_ratio <= max_n2_ratio) + ok_n0 = (n0_ratio >= min_n0_ratio) + + is_open = (ok_floor and ok_cane and ok_n2 and ok_n0 and (not top_has_corridor)) + + # 1) chão: acima do thr é bom (vira 1), abaixo cai + floor_strength = _clip01(mean_floor / max(1e-6, floor_open_thr)) + + # 2) cana: abaixo do thr é bom (vira 1), acima cai + cane_strength = _clip01(1.0 - (mean_cane / max(1e-6, cane_present_thr))) + + # 3) dominância de 0 paredes: acima do mínimo é bom + n0_strength = _clip01(n0_ratio / max(1e-6, min_n0_ratio)) + + # 4) poucas bandas com 2 paredes: abaixo do máximo é bom + n2_strength = _clip01(1.0 - (n2_ratio / max(1e-6, max_n2_ratio))) + + # 5) qualidade (quantas bandas válidas) + quality = _clip01(n_valid / max(1, bands)) + + # 6) penalidade forte se topo já parece corredor (veto) + corridor_pen = 1.0 if top_has_corridor else 0.0 + + # mistura ponderada (ajusta pesos se quiser) + base_conf = ( + 0.30 * floor_strength + + 0.25 * cane_strength + + 0.20 * n0_strength + + 0.15 * n2_strength + + 0.10 * quality + ) + + # Gates: se falhou algum critério, confiança cai junto + gate = 1.0 + gate *= 1.0 if ok_floor else 0.30 + gate *= 1.0 if ok_cane else 0.25 + gate *= 1.0 if ok_n0 else 0.35 + gate *= 1.0 if ok_n2 else 0.35 + + # Penalidade forte se topo já tem “cara de corredor” + if top_has_corridor: + gate *= 0.08 + + # Penalidade leve se está com pouca banda válida (passou raspando) + # (só ajuda a deixar conf mais honesta perto do limiar) + min_valid = max(3, int(bands * min_valid_bands_ratio)) + if n_valid <= min_valid + 1: + gate *= 0.70 + + open_conf = 100.0 * _clip01(base_conf) * _clip01(gate) + + # aplica veto como penalidade (derruba bem se topo tem corredor) + open_conf = 100.0 * _clip01(base_conf) * (0.15 if corridor_pen > 0 else 1.0) + + return is_open, { + "n_valid": n_valid, + "n0": n0, "n1": n1, "n2": n2, + "n0_ratio": n0_ratio, + "n2_ratio": n2_ratio, + "t2": t2, "b2": b2, + "top_veto_thr": top_veto_thr, + "top_has_corridor": top_has_corridor, + "mean_floor": mean_floor, + "mean_cane": mean_cane, + "ok_floor": ok_floor, + "ok_cane": ok_cane, + "ok_n2": ok_n2, + "ok_n0": ok_n0, + + "open_conf": round(open_conf, 2), + "conf_dbg": { + "floor_strength": round(floor_strength, 3), + "cane_strength": round(cane_strength, 3), + "n0_strength": round(n0_strength, 3), + "n2_strength": round(n2_strength, 3), + "quality": round(quality, 3), + "base_conf": round(base_conf, 3), + "gate": round(gate, 3), + }, + } + +def is_entering_corridor(mask_bgr: np.ndarray, + bands: int = 15, + y0_frac: float = 0.10, + y1_frac: float = 0.90, + tol_color: int = 40, + smooth_k: int = 51, + cane_peak_min: float = 0.12, + + # topo: precisa ter "corredor" (via cana-chão-cana) + top_corr_ratio: float = 0.35, + top_min_corr_bands: int = 2, + + # base: precisa ainda ter "aberto" (fronteira móvel) + min_open_bottom_bands: int = 1, # >=1 banda sem corredor no fundo + max_bot_corr_ratio: float = 0.30, # base não pode estar "cheia" de corredor + + # ainda usamos walls pra vetar ONE_WALL e IN_CORRIDOR + bot_b1_ratio_max: float = 0.20, + bot_max_bands_1wall: int = 2, + in_corridor_bot_t2_ratio: float = 0.33, + top_t2_ratio: float = 0.40, + top_min_bands_2walls: int = 2, + + # qualidade + bad_ignore_thr: float = 0.35, + bad_obst_thr: float = 0.25, + min_valid_bands_ratio: float = 0.33, + + # parâmetros do auxiliar (gap) + gap_floor_thr: float = 0.45, + gap_cane_max: float = 0.18, + gap_min_width_frac: float = 0.18, + peaks_min_sep_frac: float = 0.22, + ): + """ + ENTERING (transição para dentro do corredor, com fronteira móvel): + - Existe corredor no horizonte (muitas bandas superiores com corr_ok) + - Ainda sobra chão aberto no fundo (open_run_bottom >= 1, ou >= limiar) + - Base não está "dominada" por corredor (evita confundir com IN_CORRIDOR) + - Vetos: ONE_WALL consistente na base, ou já IN_CORRIDOR + + Retorna: + (is_entering: bool, metrics: dict) + metrics inclui: + - enter_progress 0..100 (quanto já está entrando) + - enter_conf 0..100 (confiança do ENTERING) + """ + + img = mask_bgr + H, W = img.shape[:2] + x_center = W // 2 + smooth_k = _safe_smooth_k(W, smooth_k) + + CHAO_BGR = (0, 0, 128) + CANA_BGR = (0, 128, 0) + OBST_BGR = (128, 0, 0) + + m_chao = color_mask_bgr(img, CHAO_BGR, tol=tol_color) + m_cana = color_mask_bgr(img, CANA_BGR, tol=tol_color) + m_obst = color_mask_bgr(img, OBST_BGR, tol=tol_color) + m_known = ((m_chao > 0) | (m_cana > 0) | (m_obst > 0)).astype(np.uint8) * 255 + m_ignore = cv2.bitwise_not(m_known) + + y0 = int(H * y0_frac) + y1 = int(H * y1_frac) + if y1 <= y0 + 20: + y0 = int(H * 0.5) + y1 = int(H * 0.95) + + band_h = max(10, (y1 - y0) // max(1, bands)) + + per_band = [] + for bi in range(bands): + ya = y0 + bi * band_h + yb = min(y1, ya + band_h) + if yb <= ya + 5: + continue + + chao_band = (m_chao[ya:yb, :] > 0).astype(np.float32) + cana_band = (m_cana[ya:yb, :] > 0).astype(np.float32) + ignore_band = (m_ignore[ya:yb, :] > 0).astype(np.float32) + obst_band = (m_obst[ya:yb, :] > 0).astype(np.float32) + + score_chao = smooth_1d(np.mean(chao_band, axis=0), smooth_k) + score_cana = smooth_1d(np.mean(cana_band, axis=0), smooth_k) + + ignore_ratio = float(np.mean(ignore_band)) + obst_ratio = float(np.mean(obst_band)) + + # walls (secundário, mas útil p/ vetos) + left_peak = float(np.max(score_cana[:x_center])) if x_center > 5 else 0.0 + right_peak = float(np.max(score_cana[x_center:])) if (W - x_center) > 5 else 0.0 + has_left = left_peak >= cane_peak_min + has_right = right_peak >= cane_peak_min + walls = int(has_left) + int(has_right) + + side = "none" + if walls == 1: + side = "L" if has_left else "R" + elif walls == 2: + side = "LR" + + # corredor invariante no X (principal) + corr_ok, corr_info = band_has_corridor( + score_cana, score_chao, + cane_thr=cane_peak_min, + gap_floor_thr=gap_floor_thr, + gap_cane_max=gap_cane_max, + gap_min_width_frac=gap_min_width_frac, + peaks_min_sep_frac=peaks_min_sep_frac, + ) + + bad_band = (ignore_ratio > bad_ignore_thr) or (obst_ratio > bad_obst_thr) + + per_band.append({ + "bi": bi, + "bad": bad_band, + "walls": walls, + "side": side, + "corr_ok": bool(corr_ok), + "left_peak": left_peak, + "right_peak": right_peak, + "ignore_ratio": ignore_ratio, + "obst_ratio": obst_ratio, + "corr_info": corr_info, + }) + + valid = [b for b in per_band if not b["bad"]] + n_valid = len(valid) + min_valid = max(3, int(bands * min_valid_bands_ratio)) + if n_valid < min_valid: + return False, { + "reason": "few_valid_bands", + "n_valid": n_valid, "min_valid": min_valid + } + + # ordena por bi (de cima -> baixo) + valid_sorted = sorted(valid, key=lambda b: b["bi"]) + corr_vec = [1 if b["corr_ok"] else 0 for b in valid_sorted] + n_corr_total = sum(corr_vec) + corr_ratio_total = n_corr_total / max(1, len(corr_vec)) + + # fronteira móvel: quantas bandas "abertas" ainda sobram no fundo? + open_run_bottom = _run_length_from_bottom(corr_vec, 0) + corr_run_bottom = _run_length_from_bottom(corr_vec, 1) + + # progresso (0..100): 0 = começou a aparecer corredor lá em cima; 100 = sobra só 1 banda aberta no fundo + N = len(corr_vec) + enter_progress = 100.0 * (1.0 - _clip01((open_run_bottom - 1) / max(1, (N - 1)))) + + # ainda mantemos topo/base para métricas e vetos (secundário) + mid = bands / 2.0 + top = [b for b in valid_sorted if b["bi"] < mid] + bot = [b for b in valid_sorted if b["bi"] >= mid] + if len(top) == 0 or len(bot) == 0: + half = max(1, n_valid // 2) + top = valid_sorted[:half] + bot = valid_sorted[half:] + + def counts_walls(lst): + c0 = sum(1 for b in lst if b["walls"] == 0) + c1 = sum(1 for b in lst if b["walls"] == 1) + c2 = sum(1 for b in lst if b["walls"] == 2) + return c0, c1, c2 + + def count_corr(lst): + return sum(1 for b in lst if b["corr_ok"]) + + t0, t1, t2 = counts_walls(top) + b0, b1, b2 = counts_walls(bot) + + tCorr = count_corr(top) + bCorr = count_corr(bot) + + t2_ratio = t2 / max(1, len(top)) + b2_ratio = b2 / max(1, len(bot)) + b1_ratio = b1 / max(1, len(bot)) + + tCorr_ratio = tCorr / max(1, len(top)) + bCorr_ratio = bCorr / max(1, len(bot)) + + # Regras principais + ok_top_corr = (tCorr >= max(top_min_corr_bands, int(len(top) * top_corr_ratio))) + + # Base aberta por fronteira móvel (quebra corte fixo!) + ok_bot_open_by_frontier = (open_run_bottom >= min_open_bottom_bands) + + # Base não pode estar "dominada" por corredor (evita IN_CORRIDOR travestido) + ok_bot_not_full_corridor = (bCorr_ratio <= max_bot_corr_ratio) + + # evitar ONE_WALL na base + ok_bot_no_onewall = (b1 <= bot_max_bands_1wall) and (b1_ratio <= bot_b1_ratio_max) + + # veto IN_CORRIDOR (secundário por walls2, ainda útil) + veto_in_corridor = (b2_ratio >= in_corridor_bot_t2_ratio) and (t2_ratio >= top_t2_ratio) + + # debug extra: topo com 2 paredes suficiente (selo secundário) + ok_top_walls2_dbg = (t2 >= max(top_min_bands_2walls, int(len(top) * top_t2_ratio))) + + # veto one-wall consistente na base (mesmo lado) + bot_onewall_sides = [b["side"] for b in bot if b["walls"] == 1] + veto_bot_onewall_consistent = (len(bot_onewall_sides) >= 3) and ( + bot_onewall_sides.count("R") / len(bot_onewall_sides) >= 0.85 or + bot_onewall_sides.count("L") / len(bot_onewall_sides) >= 0.85 + ) + + is_entering = ( + ok_top_corr and + ok_bot_open_by_frontier and + ok_bot_not_full_corridor and + ok_bot_no_onewall and + (not veto_in_corridor) and + (not veto_bot_onewall_consistent) + ) + + # ---------------- CONF (0..100) com gates + penalidades ---------------- + # “Força” do topo: precisa ter corredor e também bater mínimo absoluto (top_min_corr_bands) + top_ratio_strength = _clip01(tCorr_ratio / max(1e-6, top_corr_ratio)) + top_abs_strength = _clip01(tCorr / max(1, top_min_corr_bands)) + top_strength = min(top_ratio_strength, top_abs_strength) + + # “Força” do aberto no fundo: não satura tão cedo + # Queremos: 0 quando open_run_bottom < min_open_bottom_bands, + # 1 quando open_run_bottom >= (min_open_bottom_bands + 3) por exemplo. + open_margin = 3 # ajusta (2 a 4 costuma ficar bom) + open_strength = _clip01((open_run_bottom - min_open_bottom_bands) / max(1, open_margin)) + + # Base não “dominada” por corredor: 1 bom, 0 ruim + base_ok = _clip01(1.0 - (bCorr_ratio / max(1e-6, max_bot_corr_ratio))) + + # One-wall na base: 1 bom, 0 ruim (usa tanto contagem quanto razão) + onewall_ratio_ok = _clip01(1.0 - (b1_ratio / max(1e-6, bot_b1_ratio_max))) + onewall_abs_ok = _clip01(1.0 - (b1 / max(1, bot_max_bands_1wall))) + no_onewall_strength = min(onewall_ratio_ok, onewall_abs_ok) + + # Qualidade + quality = _clip01(n_valid / max(1, bands)) + + # Score base (sem veto) + score = ( + 0.45 * top_strength + + 0.25 * open_strength + + 0.20 * base_ok + + 0.10 * quality + ) + + # Gates: se alguma condição principal falha, a confiança cai junto + gate = 1.0 + gate *= 1.0 if ok_top_corr else 0.35 + gate *= 1.0 if ok_bot_open_by_frontier else 0.25 + gate *= 1.0 if ok_bot_not_full_corridor else 0.35 + gate *= 1.0 if ok_bot_no_onewall else 0.40 + + # Penalidades de veto (derruba forte) + if veto_in_corridor: + gate *= 0.08 + if veto_bot_onewall_consistent: + gate *= 0.12 + + enter_conf = 100.0 * _clip01(score) * _clip01(gate) + + return is_entering, { + "n_valid": n_valid, + + # métricas globais + "corr_ratio_total": corr_ratio_total, + "open_run_bottom": open_run_bottom, + "corr_run_bottom_dbg": corr_run_bottom, + + # novo: progressão + confiança + "enter_progress": round(enter_progress, 2), + "enter_conf": round(enter_conf, 2), + "conf_dbg": { + "top_strength": round(top_strength, 3), + "open_strength": round(open_strength, 3), + "base_ok": round(base_ok, 3), + "no_onewall_strength": round(no_onewall_strength, 3), + "quality": round(quality, 3), + "gate": round(gate, 3), + "score": round(score, 3), + }, + + # métricas topo/base (secundárias) + "top": (t0, t1, t2), + "bot": (b0, b1, b2), + "t2_ratio": t2_ratio, + "b1_ratio": b1_ratio, + "b2_ratio": b2_ratio, + "tCorr": tCorr, + "bCorr": bCorr, + "tCorr_ratio": tCorr_ratio, + "bCorr_ratio": bCorr_ratio, + + # regras + "ok_top_corr": ok_top_corr, + "ok_bot_open_by_frontier": ok_bot_open_by_frontier, + "ok_bot_not_full_corridor": ok_bot_not_full_corridor, + "ok_bot_no_onewall": ok_bot_no_onewall, + "veto_in_corridor": veto_in_corridor, + "veto_bot_onewall_consistent": veto_bot_onewall_consistent, + + # debug extra + "ok_top_walls2_dbg": ok_top_walls2_dbg, + } + +def is_exiting_corridor(mask_bgr: np.ndarray, + bands: int = 15, + y0_frac: float = 0.10, + y1_frac: float = 0.90, + tol_color: int = 40, + smooth_k: int = 51, + cane_peak_min: float = 0.12, + + # topo: precisa já estar "aberto" (pouco corredor no horizonte) + top_corr_ratio_max: float = 0.15, + top_max_corr_bands: int = 1, + top_t0_ratio: float = 0.50, # ajuda a confirmar "aberto" + + # base: ainda tem corredor (fronteira móvel) + min_corr_bottom_bands: int = 1, # >=1 banda com corredor no fundo + min_bot_corr_ratio: float = 0.30, # base não pode estar "vazia" de corredor + + # vetos (secundários) + in_corridor_top_corr_ratio: float = 0.30, # se topo ainda tem corredor forte => IN_CORRIDOR + in_corridor_bot_corr_ratio: float = 0.30, + + # evitar confundir com ONE_WALL no topo (secundário) + top_t1_ratio_max: float = 0.30, + top_max_bands_1wall: int = 2, + + # qualidade + bad_ignore_thr: float = 0.35, + bad_obst_thr: float = 0.25, + min_valid_bands_ratio: float = 0.33, + + # parâmetros do auxiliar (gap) + gap_floor_thr: float = 0.45, + gap_cane_max: float = 0.18, + gap_min_width_frac: float = 0.18, + peaks_min_sep_frac: float = 0.22, + ): + """ + EXITING (transição saindo do corredor, com fronteira móvel): + - Base: ainda há corredor embaixo (corr_run_bottom >= limiar) + - Topo: corredor já sumiu (tCorr_ratio baixo) + - Vetos: se topo ainda tem corredor forte junto com base forte => IN_CORRIDOR + Retorna: + (is_exiting: bool, metrics: dict) + metrics inclui: + - exit_progress 0..100 (quanto já está saindo; 100 = quase fora, sobra 1 banda de corredor no fundo) + - exit_conf 0..100 (confiança do EXITING) + """ + + img = mask_bgr + H, W = img.shape[:2] + x_center = W // 2 + smooth_k = _safe_smooth_k(W, smooth_k) + + CHAO_BGR = (0, 0, 128) + CANA_BGR = (0, 128, 0) + OBST_BGR = (128, 0, 0) + + m_chao = color_mask_bgr(img, CHAO_BGR, tol=tol_color) + m_cana = color_mask_bgr(img, CANA_BGR, tol=tol_color) + m_obst = color_mask_bgr(img, OBST_BGR, tol=tol_color) + m_known = ((m_chao > 0) | (m_cana > 0) | (m_obst > 0)).astype(np.uint8) * 255 + m_ignore = cv2.bitwise_not(m_known) + + y0 = int(H * y0_frac) + y1 = int(H * y1_frac) + if y1 <= y0 + 20: + y0 = int(H * 0.5) + y1 = int(H * 0.95) + + band_h = max(10, (y1 - y0) // max(1, bands)) + + per_band = [] + for bi in range(bands): + ya = y0 + bi * band_h + yb = min(y1, ya + band_h) + if yb <= ya + 5: + continue + + chao_band = (m_chao[ya:yb, :] > 0).astype(np.float32) + cana_band = (m_cana[ya:yb, :] > 0).astype(np.float32) + ignore_band = (m_ignore[ya:yb, :] > 0).astype(np.float32) + obst_band = (m_obst[ya:yb, :] > 0).astype(np.float32) + + score_chao = smooth_1d(np.mean(chao_band, axis=0), smooth_k) + score_cana = smooth_1d(np.mean(cana_band, axis=0), smooth_k) + + ignore_ratio = float(np.mean(ignore_band)) + obst_ratio = float(np.mean(obst_band)) + + # walls (secundário, mas útil p/ depuração e veto de one-wall) + left_peak = float(np.max(score_cana[:x_center])) if x_center > 5 else 0.0 + right_peak = float(np.max(score_cana[x_center:])) if (W - x_center) > 5 else 0.0 + has_left = left_peak >= cane_peak_min + has_right = right_peak >= cane_peak_min + walls = int(has_left) + int(has_right) + + side = "none" + if walls == 1: + side = "L" if has_left else "R" + elif walls == 2: + side = "LR" + + # corredor invariante no X (principal) + corr_ok, corr_info = band_has_corridor( + score_cana, score_chao, + cane_thr=cane_peak_min, + gap_floor_thr=gap_floor_thr, + gap_cane_max=gap_cane_max, + gap_min_width_frac=gap_min_width_frac, + peaks_min_sep_frac=peaks_min_sep_frac, + ) + + bad_band = (ignore_ratio > bad_ignore_thr) or (obst_ratio > bad_obst_thr) + + per_band.append({ + "bi": bi, + "bad": bad_band, + "walls": walls, + "side": side, + "corr_ok": bool(corr_ok), + "left_peak": left_peak, + "right_peak": right_peak, + "ignore_ratio": ignore_ratio, + "obst_ratio": obst_ratio, + "corr_info": corr_info, + }) + + valid = [b for b in per_band if not b["bad"]] + n_valid = len(valid) + min_valid = max(3, int(bands * min_valid_bands_ratio)) + if n_valid < min_valid: + return False, { + "reason": "few_valid_bands", + "n_valid": n_valid, "min_valid": min_valid + } + + # ordena por bi (de cima -> baixo) + valid_sorted = sorted(valid, key=lambda b: b["bi"]) + corr_vec = [1 if b["corr_ok"] else 0 for b in valid_sorted] + n_corr_total = sum(corr_vec) + corr_ratio_total = n_corr_total / max(1, len(corr_vec)) + + # fronteira móvel: quantas bandas ainda têm corredor no fundo? + corr_run_bottom = _run_length_from_bottom(corr_vec, 1) + open_run_bottom = _run_length_from_bottom(corr_vec, 0) + + # progresso (0..100): 0 = começou a abrir no topo; 100 = sobra só 1 banda de corredor no fundo (quase fora) + N = len(corr_vec) + exit_progress = 100.0 * (1.0 - _clip01((corr_run_bottom - 1) / max(1, (N - 1)))) + + # topo/base para métricas e vetos (secundário) + mid = bands / 2.0 + top = [b for b in valid_sorted if b["bi"] < mid] + bot = [b for b in valid_sorted if b["bi"] >= mid] + if len(top) == 0 or len(bot) == 0: + half = max(1, n_valid // 2) + top = valid_sorted[:half] + bot = valid_sorted[half:] + + def counts_walls(lst): + c0 = sum(1 for b in lst if b["walls"] == 0) + c1 = sum(1 for b in lst if b["walls"] == 1) + c2 = sum(1 for b in lst if b["walls"] == 2) + return c0, c1, c2 + + def count_corr(lst): + return sum(1 for b in lst if b["corr_ok"]) + + t0, t1, t2 = counts_walls(top) + b0, b1, b2 = counts_walls(bot) + + tCorr = count_corr(top) + bCorr = count_corr(bot) + + t0_ratio = t0 / max(1, len(top)) + t1_ratio = t1 / max(1, len(top)) + + tCorr_ratio = tCorr / max(1, len(top)) + bCorr_ratio = bCorr / max(1, len(bot)) + + # Regras principais + ok_top_no_corr = (tCorr <= top_max_corr_bands) and (tCorr_ratio <= top_corr_ratio_max) + ok_top_open = (t0_ratio >= top_t0_ratio) + + ok_bot_corr_by_frontier = (corr_run_bottom >= min_corr_bottom_bands) + ok_bot_not_empty_corr = (bCorr_ratio >= min_bot_corr_ratio) + + # evitar confundir com ONE_WALL persistente no topo (secundário) + ok_top_no_onewall = (t1 <= top_max_bands_1wall) and (t1_ratio <= top_t1_ratio_max) + + # veto IN_CORRIDOR: se topo ainda tem corredor forte e base tem corredor forte, então é IN_CORRIDOR + veto_in_corridor = (tCorr_ratio >= in_corridor_top_corr_ratio) and (bCorr_ratio >= in_corridor_bot_corr_ratio) + + is_exiting = ( + ok_top_no_corr and + ok_top_open and + ok_top_no_onewall and + ok_bot_corr_by_frontier and + ok_bot_not_empty_corr and + (not veto_in_corridor) + ) + + # ---------------- CONF (0..100) com gates + penalidades ---------------- + + # 1) topo aberto (precisa bater ratio e também "pouco corredor" no topo) + top_open_strength = _clip01(t0_ratio / max(1e-6, top_t0_ratio)) + + # topo sem corredor: 1 bom, 0 ruim (cai rápido se passa do limiar) + top_no_corr_strength = _clip01(1.0 - (tCorr_ratio / max(1e-6, top_corr_ratio_max))) + + # 2) ainda tem corredor no fundo (mas não saturar cedo) + # 0 quando corr_run_bottom < min_corr_bottom_bands + # 1 quando corr_run_bottom >= min + margem + corr_margin = 3 # 2..4 costuma ficar bom + corr_bottom_strength = _clip01((corr_run_bottom - min_corr_bottom_bands) / max(1, corr_margin)) + + # 3) base não pode estar "vazia" de corredor (bCorr_ratio >= min_bot_corr_ratio) + bot_corr_strength = _clip01(bCorr_ratio / max(1e-6, min_bot_corr_ratio)) + + # 4) qualidade + quality = _clip01(n_valid / max(1, bands)) + + score = ( + 0.30 * top_open_strength + + 0.30 * top_no_corr_strength + + 0.25 * corr_bottom_strength + + 0.10 * bot_corr_strength + + 0.05 * quality + ) + + # Gates (amarrar no booleano real) + gate = 1.0 + gate *= 1.0 if ok_top_no_corr else 0.20 + gate *= 1.0 if ok_top_open else 0.35 + gate *= 1.0 if ok_top_no_onewall else 0.45 + gate *= 1.0 if ok_bot_corr_by_frontier else 0.25 + gate *= 1.0 if ok_bot_not_empty_corr else 0.35 + + # Penalidade forte se parece IN_CORRIDOR + if veto_in_corridor: + gate *= 0.08 + + exit_conf = 100.0 * _clip01(score) * _clip01(gate) + + return is_exiting, { + "n_valid": n_valid, + + # métricas globais + "corr_ratio_total": corr_ratio_total, + "corr_run_bottom": corr_run_bottom, + "open_run_bottom_dbg": open_run_bottom, + + # novo: progressão + confiança + "exit_progress": round(exit_progress, 2), + "exit_conf": round(exit_conf, 2), + "conf_dbg": { + "top_open_strength": round(top_open_strength, 3), + "top_no_corr_strength": round(top_no_corr_strength, 3), + "corr_bottom_strength": round(corr_bottom_strength, 3), + "bot_corr_strength": round(bot_corr_strength, 3), + "quality": round(quality, 3), + "gate": round(gate, 3), + "score": round(score, 3), + }, + + # métricas topo/base (secundárias) + "top": (t0, t1, t2), + "bot": (b0, b1, b2), + "t0_ratio": t0_ratio, + "t1_ratio": t1_ratio, + "tCorr": tCorr, + "bCorr": bCorr, + "tCorr_ratio": tCorr_ratio, + "bCorr_ratio": bCorr_ratio, + + # regras + "ok_top_no_corr": ok_top_no_corr, + "ok_top_open": ok_top_open, + "ok_top_no_onewall": ok_top_no_onewall, + "ok_bot_corr_by_frontier": ok_bot_corr_by_frontier, + "ok_bot_not_empty_corr": ok_bot_not_empty_corr, + "veto_in_corridor": veto_in_corridor, + } + +def is_in_corridor(mask_bgr: np.ndarray, + bands: int = 15, + y0_frac: float = 0.10, + y1_frac: float = 0.90, + tol_color: int = 40, + smooth_k: int = 51, + cane_peak_min: float = 0.12, + + # modo "estrito" (cana adulta) + min_corr_ratio_total: float = 0.60, + min_top_corr_ratio: float = 0.55, + min_bot_corr_ratio: float = 0.55, + max_open_run_bottom: int = 1, # senão parece ENTERING + max_open_run_top: int = 1, # senão parece EXITING + + # modo "ponte" (cana jovem / falhas no meio) + bridge_max_gap_run: int = 4, # buraco máximo (0s seguidos) entre evidências + bridge_top_margin: int = 1, # quantos índices do topo contam como “topo” + bridge_bot_margin: int = 1, # quantos índices do fundo contam como “fundo” + max_center_jump: float = 0.12, # consistência do corredor em x + + # qualidade + bad_ignore_thr: float = 0.35, + bad_obst_thr: float = 0.25, + min_valid_bands_ratio: float = 0.33, + + # parâmetros do valley detector (repasse se quiser tunar) + floor_min_ratio: float = 0.40, + valley_min_width_frac: float = 0.18, + peaks_min_sep_frac: float = 0.22, + wall_run_thr: float = 0.10, + wall_min_run_frac: float = 0.04, + wall_side_margin_frac: float = 0.10, + wall_min_peak: float = 0.12, + valley_cane_max_ratio: float = 0.90, + + min_end_corr_ratio: float = 0.55, # extremidade clara + min_end_thickwall_ratio: float = 0.55, # fallback: pelo menos 1 parede grossa numa extremidade “ruim” + max_open_run_relaxed: int = 7, # (opcional) exige que não pareça transição demais + ): + """ + IN_CORRIDOR (v2): + - Estrito: corr_ratio_total/top/bot + sem "rabos abertos" (evita ENTERING/EXITING) + - Robusto: existe ponte topo->base (bridge_ok) permitindo falhas no meio + + centro do corredor consistente (center_ok) + - Confiança: combina evidência total + evidência topo/base + ponte + penalidade de transição + + Retorna (is_in: bool, metrics: dict) + """ + + img = mask_bgr + H, W = img.shape[:2] + x_center = W // 2 + smooth_k = _safe_smooth_k(W, smooth_k) + + CHAO_BGR = (0, 0, 128) + CANA_BGR = (0, 128, 0) + OBST_BGR = (128, 0, 0) + + m_chao = color_mask_bgr(img, CHAO_BGR, tol=tol_color) + m_cana = color_mask_bgr(img, CANA_BGR, tol=tol_color) + m_obst = color_mask_bgr(img, OBST_BGR, tol=tol_color) + m_known = ((m_chao > 0) | (m_cana > 0) | (m_obst > 0)).astype(np.uint8) * 255 + m_ignore = cv2.bitwise_not(m_known) + + y0 = int(H * y0_frac) + y1 = int(H * y1_frac) + if y1 <= y0 + 20: + y0 = int(H * 0.5) + y1 = int(H * 0.95) + + band_h = max(10, (y1 - y0) // max(1, bands)) + + per_band = [] + for bi in range(bands): + ya = y0 + bi * band_h + yb = min(y1, ya + band_h) + if yb <= ya + 5: + continue + + chao_band = (m_chao[ya:yb, :] > 0).astype(np.float32) + cana_band = (m_cana[ya:yb, :] > 0).astype(np.float32) + ignore_band = (m_ignore[ya:yb, :] > 0).astype(np.float32) + obst_band = (m_obst[ya:yb, :] > 0).astype(np.float32) + + score_chao = smooth_1d(np.mean(chao_band, axis=0), smooth_k) + score_cana = smooth_1d(np.mean(cana_band, axis=0), smooth_k) + + ignore_ratio = float(np.mean(ignore_band)) + obst_ratio = float(np.mean(obst_band)) + bad_band = (ignore_ratio > bad_ignore_thr) or (obst_ratio > bad_obst_thr) + + # walls (secundário) + left_peak = float(np.max(score_cana[:x_center])) if x_center > 5 else 0.0 + right_peak = float(np.max(score_cana[x_center:])) if (W - x_center) > 5 else 0.0 + has_left = left_peak >= cane_peak_min + has_right = right_peak >= cane_peak_min + walls = int(has_left) + int(has_right) + + # corredor por banda (vale) + corr_ok, corr_info = band_has_corridor_valley( + score_cana, score_chao, + cane_thr=cane_peak_min, + floor_min_ratio=floor_min_ratio, + valley_min_width_frac=valley_min_width_frac, + peaks_min_sep_frac=peaks_min_sep_frac, + wall_run_thr=wall_run_thr, + wall_min_run_frac=wall_min_run_frac, + wall_side_margin_frac=wall_side_margin_frac, + wall_min_peak=wall_min_peak, + valley_cane_max_ratio=valley_cane_max_ratio + ) + + ci = corr_info if isinstance(corr_info, dict) else {} + ok_left_wall = bool(ci.get("ok_left_wall", False)) + ok_right_wall = bool(ci.get("ok_right_wall", False)) + has_thick_wall = ok_left_wall or ok_right_wall + + per_band.append({ + "bi": bi, + "bad": bad_band, + "walls": walls, + "ignore_ratio": ignore_ratio, + "obst_ratio": obst_ratio, + "corr_ok": bool(corr_ok), + "corr_info": corr_info if isinstance(corr_info, dict) else {}, + "thick_wall": bool(has_thick_wall), + }) + + valid = [b for b in per_band if not b["bad"]] + n_valid = len(valid) + min_valid = max(3, int(bands * min_valid_bands_ratio)) + if n_valid < min_valid: + return False, {"reason": "few_valid_bands", "n_valid": n_valid, "min_valid": min_valid} + + # ordena por bi (topo->base) + valid_sorted = sorted(valid, key=lambda b: b["bi"]) + corr_vec = [1 if b["corr_ok"] else 0 for b in valid_sorted] + wall_vec = [1 if b.get("thick_wall") else 0 for b in valid_sorted] + N = len(corr_vec) + + n_corr = sum(corr_vec) + corr_ratio_total = n_corr / max(1, N) + + # topo/base (metade) + mid = N / 2.0 + top_idx = [i for i in range(N) if i < mid] + bot_idx = [i for i in range(N) if i >= mid] + if len(top_idx) == 0 or len(bot_idx) == 0: + half = max(1, N // 2) + top_idx = list(range(half)) + bot_idx = list(range(half, N)) + + tCorr_ratio = sum(corr_vec[i] for i in top_idx) / max(1, len(top_idx)) + bCorr_ratio = sum(corr_vec[i] for i in bot_idx) / max(1, len(bot_idx)) + tWall_ratio = sum(wall_vec[i] for i in top_idx) / max(1, len(top_idx)) + bWall_ratio = sum(wall_vec[i] for i in bot_idx) / max(1, len(bot_idx)) + + # fronteiras (transição) + open_run_bottom = _run_length_from_bottom(corr_vec, 0) + + open_run_top = 0 + for v in corr_vec: + if v == 0: + open_run_top += 1 + else: + break + + # ---------- 1) modo estrito ---------- + ok_total = (corr_ratio_total >= min_corr_ratio_total) + ok_top = (tCorr_ratio >= min_top_corr_ratio) + ok_bot = (bCorr_ratio >= min_bot_corr_ratio) + ok_no_entering = (open_run_bottom <= max_open_run_bottom) + ok_no_exiting = (open_run_top <= max_open_run_top) + strict_ok = (ok_total and ok_top and ok_bot and ok_no_entering and ok_no_exiting) + + # ---------- 2) modo ponte ---------- + # regra: primeiro 1 tem que estar “perto do topo”, último 1 “perto do fundo” e max_gap <= thr + bridge_ok, bridge_dbg = corridor_bridge_ok(corr_vec, max_gap_run=bridge_max_gap_run) + + # Ajusta os critérios topo/fundo com margens (pra não ficar dependente de ser idx<=1 etc) + # Ex.: se o 1 começa até bridge_top_margin e termina até bridge_bot_margin do fundo, ok. + if bridge_ok: + idx_first = bridge_dbg.get("idx_first", 999) + idx_last = bridge_dbg.get("idx_last", -999) + bridge_ok = (idx_first <= bridge_top_margin) and (idx_last >= (N - 1 - bridge_bot_margin)) + + # consistência de centro (só nas bandas OK) + ok_bands = [b for b in valid_sorted if b["corr_ok"] and ("center_frac" in b["corr_info"])] + center_ok, center_dbg = corridor_center_consistent(ok_bands, max_center_jump=max_center_jump) + + # robusto passa se: ponte + centro consistente + não tem “cara forte” de entering/exiting + # (aqui a gente só penaliza se tá MUITO aberto em cima/baixo) + robust_trans_ok = (open_run_bottom <= max_open_run_bottom + 4) and (open_run_top <= max_open_run_top + 4) + bridge_mode_ok = (bridge_ok and center_ok and robust_trans_ok) + + # ---------- 3) modo ENDWALL (uma ponta clara + parede grossa na outra) ---------- + top_clear = (tCorr_ratio >= min_end_corr_ratio) + bot_clear = (bCorr_ratio >= min_end_corr_ratio) + + top_has_wall = (tWall_ratio >= min_end_thickwall_ratio) + bot_has_wall = (bWall_ratio >= min_end_thickwall_ratio) + + endwall_ok = ((top_clear and bot_has_wall) or (bot_clear and top_has_wall)) + + # transição relaxada para endwall (permite base/topo meio falhando) + endwall_trans_ok = (open_run_bottom <= max_open_run_relaxed) and (open_run_top <= max_open_run_relaxed) + + endwall_mode_ok = bool(endwall_ok and endwall_trans_ok) + + # ---------- decisão final ---------- + is_in = bool(strict_ok or bridge_mode_ok or endwall_mode_ok) + + mode = "NO" + if strict_ok: + mode = "STRICT" + elif bridge_mode_ok: + mode = "BRIDGE" + elif endwall_mode_ok: + mode = "ENDWALL" + + # ---------- CONF 0..100 ---------- + # forças + total_strength = _clip01(corr_ratio_total / max(1e-6, min_corr_ratio_total)) + top_strength = _clip01(tCorr_ratio / max(1e-6, min_top_corr_ratio)) + bot_strength = _clip01(bCorr_ratio / max(1e-6, min_bot_corr_ratio)) + quality = _clip01(n_valid / max(1, bands)) + + # ponte fortalece quando strict falha + max_gap = float(bridge_dbg.get("max_gap", bridge_max_gap_run + 1)) + bridge_strength = _clip01(1.0 - (max_gap / max(1.0, bridge_max_gap_run))) if bridge_ok else 0.0 + center_strength = _clip01(1.0 - (float(center_dbg.get("max_center_jump", max_center_jump + 1.0)) / max(1e-6, max_center_jump))) if center_ok else 0.0 + bridge_combo = 0.6 * bridge_strength + 0.4 * center_strength + + # transição: quanto mais rabo aberto, mais parece entering/exiting + enter_pen = _clip01(open_run_bottom / max(1, max_open_run_bottom + 1)) + exit_pen = _clip01(open_run_top / max(1, max_open_run_top + 1)) + trans_pen = 0.5 * (enter_pen + exit_pen) # 0 bom, 1 ruim + + # score base + # - se strict tá forte: depende de total/top/bot + # - se strict tá fraco: ponte entra pra salvar + score = ( + 0.35 * total_strength + + 0.20 * min(top_strength, bot_strength) + + 0.20 * (1.0 - trans_pen) + + 0.15 * quality + + 0.10 * bridge_combo + ) + + # gate: amarra no boolean final, mas sem matar quando é “modo ponte” + gate = 1.0 + if strict_ok: + gate *= 1.0 + elif bridge_mode_ok: + gate *= 0.55 + elif endwall_mode_ok: + gate *= 0.50 # levemente abaixo do bridge (ajustável) + else: + gate *= 0.18 + + # penaliza se parece MUITO entering/exiting + if (open_run_bottom > (max_open_run_bottom + 6)) or (open_run_top > (max_open_run_top + 6)): + gate *= 0.35 + + in_conf = 100.0 * _clip01(score) * _clip01(gate) + + cap = 1.0 + if mode == "ENDWALL": + cap = 0.85 + elif mode == "BRIDGE": + cap = 0.90 + + in_conf = 100.0 * _clip01(score) * _clip01(gate) * cap + + # ---------- métricas extras ---------- + centers = [] + widths = [] + reasons = {} + for b in valid_sorted: + ci = b.get("corr_info", {}) + if isinstance(ci, dict): + if "center_frac" in ci: + centers.append(float(ci["center_frac"])) + if "width_frac" in ci: + widths.append(float(ci["width_frac"])) + r = ci.get("reason", "OK" if b.get("corr_ok") else "NO_REASON") + else: + r = "NO_INFO" + reasons[r] = reasons.get(r, 0) + 1 + + center_mean = float(np.mean(centers)) if len(centers) else None + width_mean = float(np.mean(widths)) if len(widths) else None + + return is_in, { + "n_valid": n_valid, + + "corr_ratio_total": round(corr_ratio_total, 3), + "tCorr_ratio": round(tCorr_ratio, 3), + "bCorr_ratio": round(bCorr_ratio, 3), + + "open_run_bottom": int(open_run_bottom), + "open_run_top": int(open_run_top), + + "bridge_ok": bool(bridge_ok), + "bridge_mode_ok": bool(bridge_mode_ok), + "bridge_dbg": bridge_dbg, + "center_ok": bool(center_ok), + "center_dbg": center_dbg, + + "center_mean_frac": None if center_mean is None else round(center_mean, 3), + "width_mean_frac": None if width_mean is None else round(width_mean, 3), + + "strict_ok": bool(strict_ok), + "ok_total": bool(ok_total), + "ok_top": bool(ok_top), + "ok_bot": bool(ok_bot), + "ok_no_entering": bool(ok_no_entering), + "ok_no_exiting": bool(ok_no_exiting), + + "in_conf": round(in_conf, 2), + "conf_dbg": { + "total_strength": round(float(total_strength), 3), + "top_strength": round(float(top_strength), 3), + "bot_strength": round(float(bot_strength), 3), + "bridge_combo": round(float(bridge_combo), 3), + "trans_pen": round(float(trans_pen), 3), + "quality": round(float(quality), 3), + "gate": round(float(gate), 3), + "score": round(float(score), 3), + }, + "reasons": reasons, + + "mode": mode, + + "tWall_ratio": round(tWall_ratio, 3), + "bWall_ratio": round(bWall_ratio, 3), + + "top_clear": bool(top_clear), + "bot_clear": bool(bot_clear), + "top_has_wall": bool(top_has_wall), + "bot_has_wall": bool(bot_has_wall), + + "endwall_ok": bool(endwall_ok), + "endwall_mode_ok": bool(endwall_mode_ok), + } + + + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--img", default=None, help="Path de 1 máscara PNG (vermelho=chão, verde=cana, azul=obstáculo)") + ap.add_argument("--dir", default=None, help="Pasta com várias máscaras (png/jpg). Usa A/D para navegar.") + ap.add_argument("--start", type=int, default=0, help="índice inicial ao abrir uma pasta") + ap.add_argument("--tol", type=int, default=40, help="tolerância de cor (0-255)") + ap.add_argument("--bands", type=int, default=15, help="quantidade de bandas horizontais a serem analisadas") + ap.add_argument("--y0_frac", type=float, default=0.1, help="percentual inicial para ROI vertical") + ap.add_argument("--y1_frac", type=float, default=0.9, help="percentual final para ROI vertical") + ap.add_argument("--smooth_k", type=int, default=51, help="suavização 1D por coluna") + ap.add_argument("--min_w_frac", type=float, default=0.25, help="largura mínima do corredor (fração da imagem)") + ap.add_argument("--max_w_frac", type=float, default=1.0, help="largura máxima do corredor (fração da imagem)") + ap.add_argument("--debug", action="store_true", help="mostra debug_bands") + args = ap.parse_args() + + # monta lista de imagens + files = [] + if args.dir: + files = list_images(args.dir) + if not files: + raise FileNotFoundError(f"Nenhuma imagem encontrada em: {args.dir}") + elif args.img: + files = [Path(args.img)] + else: + raise ValueError("Use --img ou --dir") + + idx = int(np.clip(args.start, 0, len(files) - 1)) + + def process_one(path: Path): + img = cv2.imread(str(path), cv2.IMREAD_COLOR) + if img is None: + return None, None, 0.0, {"reason": "falha ao abrir"} + + pts, conf, extras = detect_corridor_trapezoid( + img, + bands=args.bands, + y0_frac=args.y0_frac, + y1_frac=args.y1_frac, + smooth_k=args.smooth_k, + min_w_frac=args.min_w_frac, + max_w_frac=args.max_w_frac, + tol_color=args.tol, + debug=args.debug, + ) + + vis = img.copy() + if pts is not None: + vis = draw_trapezoid(vis, pts, conf) + else: + cv2.putText(vis, "SEM CORREDOR", (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 2, cv2.LINE_AA) + + # overlay com nome e índice + label = f"[{idx+1}/{len(files)}] {path.name} conf={conf:.2f}" + cv2.putText(vis, label, (20, vis.shape[0]-20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2, cv2.LINE_AA) + + dbg = None + if args.debug and extras is not None and pts is not None: + params = extras.get("dbg_params", {"y0_frac": args.y0_frac, "y1_frac": args.y1_frac, "bands": args.bands}) + dbg = draw_band_debug(img, extras, params["y0_frac"], params["y1_frac"], params["bands"]) + + #regime, regm = detect_corridor_regime( + # img, + # bands=args.bands, + # y0_frac=args.y0_frac, + # y1_frac=args.y1_frac, + # tol_color=args.tol, + # smooth_k=args.smooth_k, + #) + #print(f"REGIME={regime} | n0={regm.get('n0')} n1={regm.get('n1')} n2={regm.get('n2')} " + # f"| top={regm.get('top')} bot={regm.get('bot')} " + # f"| floor={regm.get('mean_floor',0):.2f} cane={regm.get('mean_cane',0):.2f}") + is_open, m = is_open_field(img, bands=args.bands, y0_frac=args.y0_frac, y1_frac=args.y1_frac, tol_color=args.tol, smooth_k=args.smooth_k) + print(f"OPEN_FIELD = {is_open} | conf={m.get('open_conf')} | n0={m.get('n0')} n1={m.get('n1')} n2={m.get('n2')} " + f"| floor={m.get('mean_floor',0):.2f} cane={m.get('mean_cane',0):.2f} " + f"| t2={m.get('t2')} veto={m.get('top_has_corridor')} " + f"| conf_dbg={m.get('conf_dbg')}") + is_ent, m = is_entering_corridor(img, bands=args.bands, y0_frac=args.y0_frac, y1_frac=args.y1_frac, tol_color=args.tol, smooth_k=args.smooth_k) + print(f"ENTERING = {is_ent} | conf={m.get('enter_conf')} | prog={m.get('enter_progress')} " + f"| corr_total={m.get('corr_ratio_total')} open_run_bot={m.get('open_run_bottom')} " + f"| top={m.get('top')} bot={m.get('bot')} " + f"| ok_top_corr={m.get('ok_top_corr')} ok_bot_open={m.get('ok_bot_open_by_frontier')} " + f"| ok_bot_not_full={m.get('ok_bot_not_full_corridor')} ok_bot_no1={m.get('ok_bot_no_onewall')} " + f"| veto_in={m.get('veto_in_corridor')} veto_onewall={m.get('veto_bot_onewall_consistent')} " + f"| conf_dbg={m.get('conf_dbg')}") + is_exit, m = is_exiting_corridor(img, bands=args.bands, y0_frac=args.y0_frac, y1_frac=args.y1_frac, tol_color=args.tol, smooth_k=args.smooth_k) + print(f"EXITING = {is_exit} | conf={m.get('exit_conf')} | prog={m.get('exit_progress')} " + f"| corr_total={m.get('corr_ratio_total')} corr_run_bot={m.get('corr_run_bottom')} " + f"| top={m.get('top')} bot={m.get('bot')} " + f"| ok_top_no_corr={m.get('ok_top_no_corr')} ok_top_open={m.get('ok_top_open')} " + f"| ok_top_no1={m.get('ok_top_no_onewall')} ok_bot_corr={m.get('ok_bot_corr_by_frontier')} " + f"| ok_bot_not_empty={m.get('ok_bot_not_empty_corr')} veto_in={m.get('veto_in_corridor')} " + f"| conf_dbg={m.get('conf_dbg')}") + is_in, m = is_in_corridor(img, bands=args.bands, y0_frac=args.y0_frac, y1_frac=args.y1_frac, tol_color=args.tol, smooth_k=args.smooth_k) + bd = m.get("bridge_dbg", {}) + cd = m.get("center_dbg", {}) + print(f"IN_CORRIDOR= {is_in} | conf={m.get('in_conf')} " + f"| corr={m.get('corr_ratio_total')} topCorr={m.get('tCorr_ratio')} botCorr={m.get('bCorr_ratio')} " + f"| topWall={m.get('tWall_ratio'):.3f} botWall={m.get('bWall_ratio'):.3f}" + f"| openBot={m.get('open_run_bottom')} openTop={m.get('open_run_top')} " + f"| strict={m.get('strict_ok')} bridgeMode={m.get('bridge_mode_ok')} bridge={m.get('bridge_ok')} " + f"| maxGap={bd.get('max_gap')} maxJump={cd.get('max_center_jump')} " + f"| center={m.get('center_mean_frac')} width={m.get('width_mean_frac')} " + f"| conf_dbg={m.get('conf_dbg')}") + + # se quiser ver estatística de reason + print("IN_CORRIDOR reasons:", m.get("reasons")) + + return img, vis, conf, extras, dbg + + # loop interativo + cv2.namedWindow("corridor_trapezoid", cv2.WINDOW_NORMAL) + cv2.namedWindow("mask", cv2.WINDOW_NORMAL) + if args.debug: + cv2.namedWindow("debug_bands", cv2.WINDOW_NORMAL) + + while True: + path = files[idx] + img, vis, conf, extras, dbg = process_one(path) + + if img is None: + # se falhou abrir, pula + print(f"[ERRO] Não abriu: {path}") + else: + print(f"[{idx+1}/{len(files)}] {path} -> conf={conf:.3f} | {extras.get('bands_ok','?')}/{extras.get('bands_total','?')}") + cv2.imshow("mask", img) + cv2.imshow("corridor_trapezoid", vis) + if args.debug and dbg is not None: + cv2.imshow("debug_bands", dbg) + + key = cv2.waitKey(0) & 0xFF + + # sair + if key in (27, ord('q'), ord('Q')): # ESC / Q + break + + # próximo + if key in (ord('d'), ord('D'), 83): # D ou seta direita (83 geralmente no Windows) + idx = min(len(files) - 1, idx + 1) + continue + + # anterior + if key in (ord('a'), ord('A'), 81): # A ou seta esquerda (81 geralmente no Windows) + idx = max(0, idx - 1) + continue + + # recomputar (mesma imagem) + if key in (ord('r'), ord('R')): + continue + + cv2.destroyAllWindows() + + +if __name__ == "__main__": + main() diff --git a/Python/OAK/detect_corridor_edges.py b/Python/OAK/detect_corridor_edges.py new file mode 100644 index 000000000..28771e996 --- /dev/null +++ b/Python/OAK/detect_corridor_edges.py @@ -0,0 +1,465 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Detecta "corredor entre evidências de cana" (moitas ou parede contínua) +e desenha os limites (bordas internas) encontrados ao longo das bandas. + +Uso: + python corridor_edges_debug.py --img path/para/mask.png --out out.png --show + +Obs: +- Espera máscara BGR (cana=(0,128,0), chão=(0,0,128), obst=(128,0,0)) igual as tuas. +- Funciona mesmo quando a cana é falhada (moitas): ele acha borda interna por banda, + e depois "ponteia" gaps interpolando. +""" + +import argparse +import cv2 +import numpy as np + + +# ------------------------- helpers ------------------------- + +def color_mask_bgr(img_bgr: np.ndarray, bgr: tuple[int, int, int], tol: int = 40) -> np.ndarray: + """Retorna máscara binária (0/255) para pixels próximos da cor BGR.""" + target = np.array(bgr, dtype=np.int16).reshape(1, 1, 3) + img16 = img_bgr.astype(np.int16) + diff = np.abs(img16 - target) + ok = (diff[:, :, 0] <= tol) & (diff[:, :, 1] <= tol) & (diff[:, :, 2] <= tol) + return (ok.astype(np.uint8) * 255) + + +def _safe_smooth_k(W: int, k: int) -> int: + k = int(k) + if k < 3: + return 3 + if k % 2 == 0: + k += 1 + if k >= W: + k = W - 1 if (W - 1) % 2 == 1 else W - 2 + return max(3, k) + + +def smooth_1d(x: np.ndarray, k: int) -> np.ndarray: + """Suaviza 1D com janela uniforme.""" + k = int(k) + if k <= 1: + return x + pad = k // 2 + xp = np.pad(x, (pad, pad), mode="edge") + kernel = np.ones(k, dtype=np.float32) / float(k) + y = np.convolve(xp, kernel, mode="valid") + return y.astype(np.float32) + + +def run_lengths(mask_bool: np.ndarray) -> tuple[int, int, int]: + """ + Retorna (best_len, best_start, best_end_inclusive) do maior run True. + Robusto para runs encostando nas bordas e para qualquer mismatch starts/ends. + """ + if mask_bool is None or mask_bool.size == 0: + return 0, -1, -1 + + a = (mask_bool.astype(np.uint8) > 0).astype(np.uint8) + + # padding garante fechamento de runs + ap = np.pad(a, (1, 1), mode="constant", constant_values=0) + d = np.diff(ap) + + starts = np.where(d == 1)[0] + ends = np.where(d == -1)[0] - 1 # inclusive, ainda no espaço de ap + + # defensivo: se por algum bug ends vier menor que starts, força fechar no final + if starts.size > ends.size: + ends = np.concatenate([ends, np.array([len(ap) - 2], dtype=ends.dtype)]) # -2 pq último índice válido do miolo + + if ends.size > starts.size: + ends = ends[:starts.size] + + if starts.size == 0: + return 0, -1, -1 + + lens = (ends - starts + 1) + j = int(np.argmax(lens)) + + best_len = int(lens[j]) + best_start = int(starts[j] - 1) # remove padding + best_end = int(ends[j] - 1) + + # clamp no range original + best_start = max(0, min(best_start, len(a) - 1)) + best_end = max(0, min(best_end, len(a) - 1)) + if best_end < best_start: + return 0, -1, -1 + + return best_len, best_start, best_end + + +def runs_from_bool(b: np.ndarray): + """ + Retorna lista de runs (start, end, len) onde b é True. + Robusto pra array vazio / sem True. + """ + b = np.asarray(b, dtype=bool) + n = b.size + if n == 0: + return [] + + # detecta mudanças + x = b.astype(np.int8) + d = np.diff(x) + + starts = list(np.where(d == 1)[0] + 1) + ends = list(np.where(d == -1)[0]) + + if b[0]: + starts = [0] + starts + if b[-1]: + ends = ends + [n - 1] + + if len(starts) != len(ends): + # safety net (não deveria acontecer, mas já vi em debug) + k = min(len(starts), len(ends)) + starts, ends = starts[:k], ends[:k] + + runs = [] + for s, e in zip(starts, ends): + if e >= s: + runs.append((int(s), int(e), int(e - s + 1))) + return runs + + +def pick_left_inner_edge(left_bool: np.ndarray, min_run: int): + """ + Parede esquerda: quer a borda interna, então escolhe o run cuja ponta direita (end) + é mais próxima do centro (maior end), desde que len >= min_run. + Retorna (x_end, best_len) ou (None, 0). + """ + runs = runs_from_bool(left_bool) + cand = [r for r in runs if r[2] >= min_run] + if not cand: + return None, 0 + # escolhe o run mais interno: maior end + s, e, L = max(cand, key=lambda r: r[1]) + return e, L + + +def pick_right_inner_edge(right_bool: np.ndarray, min_run: int): + """ + Parede direita: quer a borda interna, então escolhe o run cujo início (start) + é mais próximo do centro (menor start), desde que len >= min_run. + Retorna (x_start, best_len) ou (None, 0). + """ + runs = runs_from_bool(right_bool) + cand = [r for r in runs if r[2] >= min_run] + if not cand: + return None, 0 + # escolhe o run mais interno: menor start + s, e, L = min(cand, key=lambda r: r[0]) + return s, L + + +# ------------------------- core ------------------------- + +def detect_inner_edges_per_band( + mask_bgr: np.ndarray, + bands: int = 15, + y0_frac: float = 0.10, + y1_frac: float = 0.90, + tol_color: int = 40, + smooth_k: int = 51, + cane_thr: float = 0.10, + side_min_run_frac: float = 0.04, + side_min_density: float = 0.12, + edge_density_win_frac: float = 0.06, + center_margin_frac: float = 0.08, +): + H, W = mask_bgr.shape[:2] + smooth_k = _safe_smooth_k(W, smooth_k) + + CANA_BGR = (0, 128, 0) + m_cana = color_mask_bgr(mask_bgr, CANA_BGR, tol=tol_color) + m_cana_f = (m_cana > 0).astype(np.float32) + + y0 = int(H * y0_frac) + y1 = int(H * y1_frac) + if y1 <= y0 + 20: + y0 = int(H * 0.5) + y1 = int(H * 0.95) + + band_h = max(10, (y1 - y0) // max(1, bands)) + x_mid = W // 2 + + prev_center = None + prev_xL = None + prev_xR = None + + cm = int(W * center_margin_frac) + left_search_end = max(1, x_mid - cm) + right_search_start = min(W - 1, x_mid + cm) + + min_run = max(2, int(W * side_min_run_frac)) + win = max(5, int(W * edge_density_win_frac)) + + bands_info = [] + for bi in range(bands): + ya = y0 + bi * band_h + yb = min(y1, ya + band_h) + if yb <= ya + 5: + continue + + # 1) define referência do centro (corredor) para esta banda + if prev_center is None: + x_ref = W // 2 + else: + x_ref = int(prev_center) + + cm = int(W * center_margin_frac) + + left_search_end = max(1, x_ref - cm) + right_search_start = min(W - 1, x_ref + cm) + + band = m_cana_f[ya:yb, :] + cane_score = np.mean(band, axis=0) + cane_score = smooth_1d(cane_score, smooth_k) + + # --------- ESQUERDA (pega run mais interno) ---------- + left_region = cane_score[:left_search_end] + left_bool = left_region >= cane_thr + + left_e, left_best_len = pick_left_inner_edge(left_bool, min_run=min_run) + xL = int(left_e) if left_e is not None else None + + left_density_ok = False + if xL is not None: + # densidade medida em cane_score (mais consistente) + x0 = max(0, xL - win + 1) + x1 = xL + 1 + left_density = float(np.mean(cane_score[x0:x1])) + left_density_ok = (left_density >= side_min_density) + if not left_density_ok: + xL = None + + # --------- DIREITA (pega run mais interno) ---------- + right_region = cane_score[right_search_start:] + right_bool = right_region >= cane_thr + + right_s, right_best_len = pick_right_inner_edge(right_bool, min_run=min_run) + xR = (right_search_start + int(right_s)) if right_s is not None else None + + right_density_ok = False + if xR is not None: + x0 = xR + x1 = min(W, xR + win) + right_density = float(np.mean(cane_score[x0:x1])) + right_density_ok = (right_density >= side_min_density) + if not right_density_ok: + xR = None + + # corredor válido na banda + band_mid_y = int(0.5 * (ya + yb)) + ok_band = (xL is not None) and (xR is not None) and (xR > xL + 10) + + if xL is not None and xR is not None: + prev_xL = xL + prev_xR = xR + prev_center = 0.5 * (xL + xR) + elif xL is not None and prev_xR is not None: + prev_xL = xL + prev_center = 0.5 * (xL + prev_xR) + elif xR is not None and prev_xL is not None: + prev_xR = xR + prev_center = 0.5 * (prev_xL + xR) + + width_frac = None + center_frac = None + if ok_band: + width_frac = float(xR - xL) / max(1, W) + center_frac = float(0.5 * (xL + xR)) / max(1, (W - 1)) + + bands_info.append({ + "bi": bi, + "ya": ya, + "yb": yb, + "y": band_mid_y, + "xL": xL, + "xR": xR, + "ok": bool(ok_band), + "center_frac": center_frac, + "width_frac": width_frac, + "dbg": { + "left_best_len": int(left_best_len), + "right_best_len": int(right_best_len), + "min_run": int(min_run), + "left_density_ok": bool(left_density_ok), + "right_density_ok": bool(right_density_ok), + "left_search_end": int(left_search_end), + "right_search_start": int(right_search_start), + } + }) + + return bands_info, m_cana + + +def bridge_edges(bands_info: list[dict], max_gap_bands: int = 4): + """ + Interpola gaps (bandas onde faltou xL/xR) se a lacuna entre evidências + for <= max_gap_bands. + """ + # índices válidos para cada lado + ys = [b["y"] for b in bands_info] + + def bridge_key(key: str): + xs = [b[key] for b in bands_info] + # positions with value + idx = [i for i, v in enumerate(xs) if v is not None] + if len(idx) < 2: + return xs # nada pra ponteiar + xs2 = xs[:] + for a, b in zip(idx[:-1], idx[1:]): + gap = b - a - 1 + if gap <= 0: + continue + if gap > max_gap_bands: + continue + xa = xs[a] + xb = xs[b] + if xa is None or xb is None: + continue + for t, i in enumerate(range(a + 1, b)): + alpha = (t + 1) / float(gap + 1) + xs2[i] = int(round((1 - alpha) * xa + alpha * xb)) + return xs2 + + xL_br = bridge_key("xL") + xR_br = bridge_key("xR") + + out = [] + for i, b in enumerate(bands_info): + bb = dict(b) + bb["xL_br"] = xL_br[i] + bb["xR_br"] = xR_br[i] + bb["ok_br"] = (bb["xL_br"] is not None) and (bb["xR_br"] is not None) and (bb["xR_br"] > bb["xL_br"] + 10) + if bb["ok_br"]: + bb["center_br"] = float(0.5 * (bb["xL_br"] + bb["xR_br"])) / max(1, (bb.get("yb", 0) * 0 + (bands_info[0].get("yb", 1))) ) # dummy (não usamos) + out.append(bb) + return out + + +def draw_debug(mask_bgr: np.ndarray, bands_br: list[dict], out_path: str, show: bool = False): + vis = mask_bgr.copy() + H, W = vis.shape[:2] + + # desenha linhas por banda + ptsL = [] + ptsR = [] + ptsL_br = [] + ptsR_br = [] + + for b in bands_br: + y = int(b["y"]) + + # linha horizontal da banda (bem leve) + cv2.line(vis, (0, y), (W - 1, y), (40, 40, 40), 1, cv2.LINE_AA) + + xL = b.get("xL", None) + xR = b.get("xR", None) + if xL is not None: + ptsL.append((int(xL), y)) + cv2.circle(vis, (int(xL), y), 4, (255, 255, 255), -1, cv2.LINE_AA) + if xR is not None: + ptsR.append((int(xR), y)) + cv2.circle(vis, (int(xR), y), 4, (255, 255, 255), -1, cv2.LINE_AA) + + # ponteado + xLb = b.get("xL_br", None) + xRb = b.get("xR_br", None) + if xLb is not None: + ptsL_br.append((int(xLb), y)) + cv2.circle(vis, (int(xLb), y), 3, (0, 255, 255), -1, cv2.LINE_AA) + if xRb is not None: + ptsR_br.append((int(xRb), y)) + cv2.circle(vis, (int(xRb), y), 3, (0, 255, 255), -1, cv2.LINE_AA) + + if (xLb is not None) and (xRb is not None) and (xRb > xLb + 10): + # desenha o "corredor" por banda + cv2.line(vis, (int(xLb), y), (int(xRb), y), (0, 255, 255), 2, cv2.LINE_AA) + + # desenha polilinhas (limites) + if len(ptsL_br) >= 2: + cv2.polylines(vis, [np.array(ptsL_br, dtype=np.int32)], False, (0, 255, 255), 2, cv2.LINE_AA) + if len(ptsR_br) >= 2: + cv2.polylines(vis, [np.array(ptsR_br, dtype=np.int32)], False, (0, 255, 255), 2, cv2.LINE_AA) + + # original (não-ponteado) em cinza claro, pra comparar + if len(ptsL) >= 2: + cv2.polylines(vis, [np.array(ptsL, dtype=np.int32)], False, (200, 200, 200), 1, cv2.LINE_AA) + if len(ptsR) >= 2: + cv2.polylines(vis, [np.array(ptsR, dtype=np.int32)], False, (200, 200, 200), 1, cv2.LINE_AA) + + #cv2.imwrite(out_path, vis) + #print(f"[OK] Salvo: {out_path}") + + if show: + cv2.imshow("corridor_edges_debug", vis) + cv2.waitKey(0) + cv2.destroyAllWindows() + + +# ------------------------- main ------------------------- + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--img", required=True, help="Caminho da imagem (mask BGR).") + ap.add_argument("--out", default="corridor_debug_out.png", help="Saída anotada.") + ap.add_argument("--show", action="store_true", help="Abre janela com resultado.") + ap.add_argument("--bands", type=int, default=15) + ap.add_argument("--y0", type=float, default=0.0) + ap.add_argument("--y1", type=float, default=1.0) + ap.add_argument("--tol", type=int, default=40) + ap.add_argument("--smooth_k", type=int, default=51) + ap.add_argument("--cane_thr", type=float, default=0.10) + ap.add_argument("--side_min_run_frac", type=float, default=0.04) + ap.add_argument("--side_min_density", type=float, default=0.12) + ap.add_argument("--edge_density_win_frac", type=float, default=0.06) + ap.add_argument("--center_margin_frac", type=float, default=0.05) + ap.add_argument("--bridge_max_gap_bands", type=int, default=4) + + args = ap.parse_args() + + img = cv2.imread(args.img, cv2.IMREAD_COLOR) + if img is None: + raise SystemExit(f"Falhou ao ler imagem: {args.img}") + + bands_info, m_cana = detect_inner_edges_per_band( + img, + bands=args.bands, + y0_frac=args.y0, + y1_frac=args.y1, + tol_color=args.tol, + smooth_k=args.smooth_k, + cane_thr=args.cane_thr, + side_min_run_frac=args.side_min_run_frac, + side_min_density=args.side_min_density, + edge_density_win_frac=args.edge_density_win_frac, + center_margin_frac=args.center_margin_frac, + ) + + # logs rápidos + n = len(bands_info) + ok = sum(1 for b in bands_info if b["ok"]) + left_present = sum(1 for b in bands_info if b["xL"] is not None) + right_present = sum(1 for b in bands_info if b["xR"] is not None) + print(f"[bands] n={n} ok_both={ok} left_present={left_present} right_present={right_present}") + + # ponteia + bands_br = bridge_edges(bands_info, max_gap_bands=args.bridge_max_gap_bands) + ok_br = sum(1 for b in bands_br if b["ok_br"]) + print(f"[bridge] ok_both_after_bridge={ok_br} (max_gap_bands={args.bridge_max_gap_bands})") + + # desenha + draw_debug(img, bands_br, args.out, show=args.show) + + +if __name__ == "__main__": + main()