Ajustado scripts de treinamento e health worker

This commit is contained in:
Diego Freitas 2026-01-28 14:52:31 -03:00
parent 4a1561fb16
commit cb0ebd8b1d
27 changed files with 6587 additions and 293 deletions

View File

@ -2496,40 +2496,64 @@ namespace AgroBase.Models
public enum AutonomiaCorredorStatus public enum AutonomiaCorredorStatus
{ {
Desconhecido = 0, Desconhecido = 0,
OkCompleto, // Bateria + herbicida suficientes pra entrar e pulverizar OkCompleto, // Bateria + herbicida suficientes pra entrar e pulverizar
OkSomenteTransito, // Bateria ok, herbicida insuficiente pra pulverizar tudo OkSomenteTransito, // Bateria ok, herbicida insuficiente pra pulverizar tudo (mas pode atravessar)
OkSupervisionado, // Liberado manualmente OkSupervisionado, // Liberado manualmente MESMO com parâmetro não OK (override humano)
Critico // Nem a bateria garante atravessar com segurança 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 public class AutonomiaCorredorModel
{ {
// ==========================
// Configurações
// ==========================
private readonly TimeSpan _tempoParaTravar = TimeSpan.FromSeconds(10); // X segundos em crítico -> trava
// ==========================
// Estado público
// ==========================
[JsonProperty] [JsonProperty]
public bool Iniciado { get; private set; } public bool Iniciado { get; private set; }
[JsonProperty] [JsonProperty]
public AutonomiaCorredorStatus Status { get; private set; } public AutonomiaCorredorStatus Status { get; private set; } = AutonomiaCorredorStatus.Desconhecido;
private bool _liberado =>
Status == AutonomiaCorredorStatus.OkCompleto ||
Status == AutonomiaCorredorStatus.OkSomenteTransito ||
Status == AutonomiaCorredorStatus.OkSupervisionado;
[JsonProperty] [JsonProperty]
public bool Liberado { get; private set; } public bool Liberado { get; private set; } = true;
[JsonProperty] [JsonProperty]
public double DistanciaCorredor_m { get; private set; } public double DistanciaCorredor_m { get; private set; }
[JsonProperty] [JsonProperty]
public double DistanciaSeguraBateria_m { get; private set; } public double DistanciaSeguraBateria_m { get; private set; }
[JsonProperty] [JsonProperty]
public double DistanciaSeguraHerbicida_m { get; private set; } public double DistanciaSeguraHerbicida_m { get; private set; }
private Dictionary<int, AutonomiaCorredorModel> CorredoresLiberados { get; set; } = new Dictionary<int, AutonomiaCorredorModel>();
// Flags de conveniência
public bool BateriaSuficiente => DistanciaSeguraBateria_m >= DistanciaCorredor_m;
public bool HerbicidaSuficiente => DistanciaSeguraHerbicida_m >= DistanciaCorredor_m;
// Texto pra UI/log
[JsonProperty] [JsonProperty]
public string Motivo { get; private set; } = ""; 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<int, AutonomiaCorredorModel> CorredoresLiberados { get; set; } = new Dictionary<int, AutonomiaCorredorModel>();
// ==========================
// Trava (latch) por corredor
// ==========================
private readonly HashSet<int> _corredoresTravados = new HashSet<int>();
private readonly Dictionary<int, DateTime> _inicioCriticoUtc = new Dictionary<int, DateTime>();
// ==========================
// Ciclo de vida
// ==========================
public void Iniciar() public void Iniciar()
{ {
Iniciado = true; Iniciado = true;
@ -2538,10 +2562,18 @@ namespace AgroBase.Models
public void Parar() public void Parar()
{ {
Iniciado = false; 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) public void AtualizarDados(int idxCorredor, bool bateria_liberada, bool reservatorio_liberado)
{ {
// Se não iniciado, não bloqueia nada
if (!Iniciado) if (!Iniciado)
{ {
Liberado = true; Liberado = true;
@ -2550,12 +2582,18 @@ namespace AgroBase.Models
return; return;
} }
// Valida índice do corredor
if (Variaveis.OperacaoEmAndamento.Trajetoria?._Corredores?.Count <= idxCorredor) if (Variaveis.OperacaoEmAndamento.Trajetoria?._Corredores?.Count <= idxCorredor)
return; return;
// Se já registrou esse corredor como liberado e travado/feito, não recalcula
if (CorredoresLiberados.ContainsKey(idxCorredor) && CorredoresLiberados[idxCorredor].Liberado) if (CorredoresLiberados.ContainsKey(idxCorredor) && CorredoresLiberados[idxCorredor].Liberado)
return; 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; double distanciaCorredor = Variaveis.OperacaoEmAndamento.Trajetoria._Corredores[idxCorredor]?.DistanciaTotal ?? 0;
var _Bateria = Variaveis.OperacaoEmAndamento.DispSen?.Dados?.ContadorCarga; var _Bateria = Variaveis.OperacaoEmAndamento.DispSen?.Dados?.ContadorCarga;
@ -2567,41 +2605,157 @@ namespace AgroBase.Models
bool bateriaOk = BateriaSuficiente && (_Bateria?.Iniciado ?? false); bool bateriaOk = BateriaSuficiente && (_Bateria?.Iniciado ?? false);
bool herbOk = HerbicidaSuficiente && (_Reservatorio?.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<AutonomiaCorredorStatus> status_mods = new List<AutonomiaCorredorStatus>(); bool exigirPulverizacaoTotal =
string _motivo = ""; Variaveis.OperacaoEmAndamento.Parametros?.ModulosMandatorios?
string _motivo_bat = "Bateria " + (bateriaOk ? "suficiente" : "insuficiente") + $". Distância segura: {DistanciaSeguraBateria_m:F1} m, corredor: {distanciaCorredor:F1} m"; .Any(x => x.Dispositivo == T_Code.Atu && x.Mandatorio && x.Utilizar) ?? false;
string _motivo_herb = "Herbicida " + (herbOk ? "suficiente" : "insuficiente") + $" para pulverizar o corredor. Distância segura pulverizando: {DistanciaSeguraHerbicida_m:F1} m, corredor: {distanciaCorredor:F1} m";
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); statusBase = AutonomiaCorredorStatus.Critico;
if (bateria_liberada) _motivo_bat = $"Liberação humana: {_motivo_bat}";
_motivo += _motivo_bat; if (criticoBateria)
motivoBase += motivoBat;
if (!herbOk)
motivoBase += (string.IsNullOrEmpty(motivoBase) ? "" : "; ") + motivoHerb;
} }
else
if (!herbOk)
{ {
if (exigirPulverizacaoTotal) // Não crítico
status_mods.Add(!reservatorio_liberado ? AutonomiaCorredorStatus.Critico : AutonomiaCorredorStatus.OkSupervisionado); if (!herbOk && !exigirPulverizacaoTotal)
{
// Pode atravessar, mas não pulveriza tudo
statusBase = AutonomiaCorredorStatus.OkSomenteTransito;
motivoBase = motivoHerb;
}
else else
status_mods.Add(AutonomiaCorredorStatus.OkSomenteTransito); {
if (reservatorio_liberado) _motivo_herb = $"Liberação humana: {_motivo_herb}"; // Tudo ok
_motivo += (string.IsNullOrEmpty(_motivo) ? "" : "; ") + _motivo_herb; statusBase = AutonomiaCorredorStatus.OkCompleto;
}
} }
Status = // ==========================
status_mods.Contains(AutonomiaCorredorStatus.Critico) ? AutonomiaCorredorStatus.Critico : // 2) Atualiza timers de crítico e ativa trava se necessário
status_mods.Contains(AutonomiaCorredorStatus.OkSupervisionado) ? AutonomiaCorredorStatus.OkSupervisionado : // ==========================
status_mods.Contains(AutonomiaCorredorStatus.OkSomenteTransito) ? AutonomiaCorredorStatus.OkSomenteTransito : bool estaCriticoAgora = (statusBase == AutonomiaCorredorStatus.Critico);
AutonomiaCorredorStatus.OkCompleto;
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() var res = new AutonomiaCorredorModel()
{ {
Liberado = _liberado, Iniciado = this.Iniciado,
Motivo = $"Corredor {idxCorredor + 1}: " + (string.IsNullOrEmpty(_motivo) ? "Liberado" : $"{_motivo}"), Liberado = liberadoFinal,
Status = Status, Motivo = motivoFinal,
Status = statusFinal,
DistanciaCorredor_m = DistanciaCorredor_m, DistanciaCorredor_m = DistanciaCorredor_m,
DistanciaSeguraBateria_m = DistanciaSeguraBateria_m, DistanciaSeguraBateria_m = DistanciaSeguraBateria_m,
DistanciaSeguraHerbicida_m = DistanciaSeguraHerbicida_m, DistanciaSeguraHerbicida_m = DistanciaSeguraHerbicida_m,
@ -2611,11 +2765,17 @@ namespace AgroBase.Models
if (CorredoresLiberados.ContainsKey(idxCorredor)) 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; CorredoresLiberados[idxCorredor] = res;
addLog = true; addLog = true;
} }
else
{
// Mantém o mais recente motivo/status para UI (opcional)
CorredoresLiberados[idxCorredor] = res;
}
} }
else else
{ {
@ -2625,18 +2785,28 @@ namespace AgroBase.Models
if (addLog) 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; // Atualiza estado exposto
Motivo = CorredoresLiberados[idxCorredor].Motivo; Liberado = res.Liberado;
Motivo = res.Motivo;
} }
// ==========================
// Clone (mantido)
// ==========================
public AutonomiaCorredorModel Clone() public AutonomiaCorredorModel Clone()
{ {
return new AutonomiaCorredorModel() return new AutonomiaCorredorModel()
{ {
Iniciado = Iniciado,
Liberado = Liberado, Liberado = Liberado,
Motivo = Motivo, Motivo = Motivo,
Status = Status, Status = Status,

View File

@ -49,7 +49,7 @@ def main():
try: try:
# Atualizar SAÚDE DOS MÓDULOS a cada 2.5s # 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: try:
for t_code, modulo in modulos.items(): for t_code, modulo in modulos.items():
modulo.atualizar_saude() modulo.atualizar_saude()

View File

@ -1,9 +1,11 @@
import re
import time import time
import threading import threading
import statistics import statistics
import socket import socket
import psutil import psutil
import subprocess import subprocess
from concurrent.futures import ThreadPoolExecutor
from collections import deque from collections import deque
import paho.mqtt.client as mqtt import paho.mqtt.client as mqtt
from shared.contexto_global_redis import ContextoGlobalRedis from shared.contexto_global_redis import ContextoGlobalRedis
@ -19,7 +21,7 @@ class ModuloIPBribge(ModuloDiagnosticoBase):
self._thread_saude = None self._thread_saude = None
self._running = False self._running = False
self.nic_name = None self.nic_name = self.descobrir_nic_para_base()
self.window_size = window_size self.window_size = window_size
self.rtts = deque(maxlen=window_size) # ping válidos (ms) 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.last_heartbeat_ts = time.time() # atualize isso de fora
self.rover_id = None self.rover_id = None
self._mqtt_conectado = False
self.sub = 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: 300500
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): def atualizar_saude(self):
@ -72,38 +153,107 @@ class ModuloIPBribge(ModuloDiagnosticoBase):
self._mqtt.connect(base_ip, 1883, 60) self._mqtt.connect(base_ip, 1883, 60)
self._mqtt.loop_start() self._mqtt.loop_start()
except: except:
pass #pass
#print("Erro ao se conectar no broker mqtt") 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): def _on_connect(self, client, userdata, flags, rc):
if rc == 0: with self._mqtt_lock:
print("[HEARTBEAT] MQTT conectado") self._mqtt_conectado = (rc == 0)
self._mqtt_conectado = True self._mqtt_connecting = False
# reset backoff quando conecta
self._reset_janelas() if self._mqtt_conectado:
self._mqtt_backoff_s = 1.0
self.rover_id = ContextoGlobalRedis.get_equipamento().get("serial_number") self._reset_janelas()
self.rover_id = ContextoGlobalRedis.get_equipamento().get("serial_number")
if self.sub == False: if self.sub == False:
topic = f"agrobot/v1/rover/{self.rover_id}/heartbeat" topic = f"agrobot/v1/rover/{self.rover_id}/heartbeat"
client.subscribe(topic) client.subscribe(topic)
print("[HEARTBEAT] Subscribado em:", topic) print("[HEARTBEAT] Subscribado em:", topic)
self.sub = True self.sub = True
else: else:
print("Erro ao conectar MQTT:", rc) print("Erro ao conectar MQTT:", rc)
def _on_disconnect(self, client, userdata, rc): def _on_disconnect(self, client, userdata, rc):
print("[HEARTBEAT] MQTT desconectado!", rc) with self._mqtt_lock:
self._mqtt_conectado = False self._mqtt_conectado = False
self.sub = False self._mqtt_connecting = False
self.sub = False
print("[HEARTBEAT] MQTT desconectado!", rc)
# se rc != 0 significa desconexão inesperada # se rc != 0 significa desconexão inesperada
if rc != 0: if rc != 0:
print("[HEARTBEAT] Desconexão inesperada — pode ter perdido o link 900 MHz") print("[HEARTBEAT] Desconexão inesperada — pode ter perdido o link 900 MHz")
# Marca como "sem heartbeat" instantâneo (opcional) # Marca como "sem heartbeat" instantâneo (opcional)
# Aqui podemos forçar atraso grande # Aqui podemos forçar atraso grande
self.last_heartbeat_ts = 0 self.last_heartbeat_ts = 0
def _on_message(self, client, userdata, msg): def _on_message(self, client, userdata, msg):
topic = msg.topic topic = msg.topic
@ -185,7 +335,7 @@ class ModuloIPBribge(ModuloDiagnosticoBase):
self.last_bw_counters = counters self.last_bw_counters = counters
self.last_bw_ts = now 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. Retorna RTT em ms se ok, ou None se timeout.
Aqui exemplo pra Windows usando 'ping -n 1'. Aqui exemplo pra Windows usando 'ping -n 1'.
@ -198,14 +348,14 @@ class ModuloIPBribge(ModuloDiagnosticoBase):
# timeout de 1000 ms # timeout de 1000 ms
proc = subprocess.run( proc = subprocess.run(
["ping", "-n", "5", "-w", "500", base_ip], ["ping", "-n", f"{pings}", "-w", f"{timeout}", base_ip],
capture_output=True, text=True capture_output=True, text=True
) )
if proc.returncode != 0: if proc.returncode != 0:
#print("timeout") print("timeout")
return None, 100 return None, 100
#print(proc.stdout) print(proc.stdout)
perda = None perda = None
media_tempo = None media_tempo = None
@ -224,6 +374,51 @@ class ModuloIPBribge(ModuloDiagnosticoBase):
print("erro timeout") print("erro timeout")
return None, 100 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): def _update_nic_errors(self):
if self.nic_name is None: if self.nic_name is None:
self.nic_name = self.descobrir_nic_para_base() self.nic_name = self.descobrir_nic_para_base()
@ -306,7 +501,7 @@ class ModuloIPBribge(ModuloDiagnosticoBase):
loss_score = 0 loss_score = 0
# JITTER # JITTER
if len(self.rtts) > 1: if len(self.rtts) > 2:
jitter = statistics.pstdev(self.rtts) jitter = statistics.pstdev(self.rtts)
if jitter <= 10: if jitter <= 10:
jitter_score = 100 jitter_score = 100
@ -319,7 +514,7 @@ class ModuloIPBribge(ModuloDiagnosticoBase):
else: else:
jitter = 0.0 jitter = 0.0
# Sem dados suficientes, melhor considerar "desconhecido" => score neutro/baixo # Sem dados suficientes, melhor considerar "desconhecido" => score neutro/baixo
jitter_score = 0 jitter_score = 1
# NIC ERRORS # NIC ERRORS
if self.nic_error_rates: if self.nic_error_rates:
@ -337,14 +532,11 @@ class ModuloIPBribge(ModuloDiagnosticoBase):
nic_score = 0 nic_score = 0
# HEARTBEAT (score usando "atraso") # HEARTBEAT (score usando "atraso")
if atraso <= 1.5: t0, t1 = 1.0, 3.0 # 1s perfeito, 3s zerou
hb_score = 100 x = (atraso - t0) / (t1 - t0) # 0..1
elif atraso <= 2: x = max(0.0, min(1.0, x))
hb_score = 80 gamma = 1.6 # >1 deixa cair mais rápido perto do final
elif atraso <= 5: hb_score = int(round(100 * (1.0 - (x ** gamma))))
hb_score = 30
else:
hb_score = 0
# retorna scores + brutos # retorna scores + brutos
return ( return (
@ -353,8 +545,9 @@ class ModuloIPBribge(ModuloDiagnosticoBase):
) )
def atualizar_saude_interno(self): def atualizar_saude_interno(self):
#print("atualizando saude IPB...")
try: try:
self._start_mqtt_heartbeat() self._start_mqtt_heartbeat_async()
SAUDE_MIN_ALERTA = 80 SAUDE_MIN_ALERTA = 80
@ -363,14 +556,27 @@ class ModuloIPBribge(ModuloDiagnosticoBase):
saude_individual = [] saude_individual = []
conectado = self.get_base_ip() is not None and self.nic_name is not None and self._mqtt_conectado conectado = self.get_base_ip() is not None and self.nic_name is not None and self._mqtt_conectado
(rtt, loss) = self._ping_once() # 1) consome resultado pronto (não bloqueia)
#print(f"rtt: {rtt}, loss: {loss}") 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) self.timeouts.append(loss == 100)
if rtt is not None: if ok and rtt is not None:
self.rtts.append(rtt) self.rtts.append(rtt)
if loss is not None: if loss is not None:
self.loses.append(loss) 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_errors()
self._update_nic_bandwidth() self._update_nic_bandwidth()

View File

@ -50,10 +50,15 @@ def definir_comando(parada_necessaria: bool, erro: bool, filtro_vel: FiltroVeloc
pontos_min = 3 pontos_min = 3
vel_min = 15.0 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_com_ervas = _operacao.get("Mov", {}).get("percent_vel_min", 0)
vel_sem_ervas = _operacao.get("Mov", {}).get("percent_vel_max", 100) vel_sem_ervas = _operacao.get("Mov", {}).get("percent_vel_max", 100)
ang_max = _operacao.get("Dir", {}).get("angulo_max", 30) 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 status_carro = StatusCarroMapa.Parado
reduzir_para_pulverizar = False reduzir_para_pulverizar = False
@ -70,7 +75,9 @@ def definir_comando(parada_necessaria: bool, erro: bool, filtro_vel: FiltroVeloc
velocidade_sp = vel_com_ervas 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 velocidade_sp = 0
elif status_carro in [StatusCarroMapa.Direcionando, StatusCarroMapa.RetornandoBase]: 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) velocidade_sp = calcular_velocidade_relativa(vel_com_ervas, vel_sem_ervas, ang_max, erro_orientacao, k=0.85, curva=0.8, dead=10.0)

View File

@ -1262,6 +1262,12 @@ namespace OperationControl.Windows
txtTrajetoria_HerbOk.Text = (dados?.HerbicidaSuficiente ?? false) ? "Sim" : "Não"; txtTrajetoria_HerbOk.Text = (dados?.HerbicidaSuficiente ?? false) ? "Sim" : "Não";
btnTrajetoria_Confirmar.IsEnabled = !(dados?.Liberado ?? false); btnTrajetoria_Confirmar.IsEnabled = !(dados?.Liberado ?? false);
if (inicial)
{
chbTrajetoria_BatLiberada.IsChecked = dados?.BateriaSuficiente ?? false;
chbTrajetoria_HerbLiberado.IsChecked = dados?.HerbicidaSuficiente ?? false;
}
break; break;
} }
case AgroBase.Models.Enums.T_Code.Npc: case AgroBase.Models.Enums.T_Code.Npc:

View File

@ -13,6 +13,7 @@ Uso:
python _12_check_percent_class_labelmap.py python _12_check_percent_class_labelmap.py
""" """
import argparse
import os, json import os, json
import numpy as np import numpy as np
from PIL import Image from PIL import Image
@ -28,8 +29,6 @@ ROI_TAMANHO = config["roi_tamanho"]
pasta_base = os.path.join(MODELO, "dataset") pasta_base = os.path.join(MODELO, "dataset")
labelmap_path = os.path.join(pasta_base, "labelmap.txt") 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. # Limite de amostras por grupo (para rodar rápido). Ajuste se quiser.
MAX_SAMPLES_PER_GROUP = 1000 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 y_fim, y_ini = max(0, h - int(ROI_TAMANHO * h)), h
return slice(y_fim, y_ini) return slice(y_fim, y_ini)
# -------------- Labelmap -------------- def main(args):
_, colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path) # -------------- Labelmap --------------
ignore_id = infer_ignore_id(ignore_rgb, default_id=255) _, 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 # 'classes' esperado como dict: id -> nome
# Ordena por id para imprimir de forma estável # Ordena por id para imprimir de forma estável
class_ids_sorted = sorted(classes.keys()) class_ids_sorted = sorted(classes.keys())
class_names_sorted = [classes[cid] for cid in class_ids_sorted] class_names_sorted = [classes[cid] for cid in class_ids_sorted]
# -------------- Coleta -------------- root = os.path.join(pasta_base, "split", args.split, "group")
if not os.path.isdir(root): #root = os.path.join(pasta_base, "576x320", "group")
raise SystemExit(f"Nenhum diretório encontrado em {root}")
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): grupos = [g for g in os.listdir(root) if os.path.isdir(os.path.join(root, g))]
mdir = os.path.join(root, g, "masks")
if not os.path.isdir(mdir):
continue
totals = {cid: 0 for cid in class_ids_sorted} for g in sorted(grupos):
n = 0 mdir = os.path.join(root, g, "masks")
if not os.path.isdir(mdir):
for fname in os.listdir(mdir):
if not fname.lower().endswith(".png"):
continue 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: 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 print(f"{g:16s} " + " ".join(parts) + f" (amostras={n})")
if n >= MAX_SAMPLES_PER_GROUP:
break
s = sum(totals.values()) if __name__ == "__main__":
s = s if s > 0 else 1 # evita div/0 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).")
# Monta string dinâmica "nome=xx.xx%" args = ap.parse_args()
parts = [] main(args)
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})")

View File

@ -199,10 +199,11 @@ def localizar_raw_correspondente(pasta_new_raws, nome_img: str) -> str | None:
return candidate return candidate
return None return None
def processar_novas_imagens(classe, cor_classe_rgb, fazer_copia_final=True, manifesto_csv=None): def processar_novas_imagens(cana, horario, grupo, cor_classe_rgb, fazer_copia_final=True, manifesto_csv=None, orignais=False):
pasta_new_previews = os.path.join(MODELO, "dataset", "brutas", "group", classe, "previews") source = os.path.join(MODELO, "dataset", "brutas", f"cana_{cana}", horario) if not orignais else os.path.join(MODELO, "dataset", "original")
pasta_new_masks = os.path.join(MODELO, "dataset", "brutas", "group", classe, "masks") pasta_new_previews = os.path.join(source, "group", grupo, "previews")
pasta_new_raws = os.path.join(MODELO, "dataset", "brutas", "group", classe, "raws") 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_previews)
garantir_pasta(pasta_new_masks) garantir_pasta(pasta_new_masks)
@ -281,8 +282,12 @@ def build_cli():
ap = argparse.ArgumentParser( ap = argparse.ArgumentParser(
description="Gera máscaras sólidas para novas imagens de UMA classe (via labelmap) e copia para dataset final com dedup." 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("--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("--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).") ap.add_argument("--manifest", default="", help="Caminho do CSV de manifesto a gerar (ou vazio para não gerar).")
return ap return ap
@ -303,8 +308,11 @@ if __name__ == "__main__":
manifesto_csv = args.manifest if args.manifest else None manifesto_csv = args.manifest if args.manifest else None
processar_novas_imagens( processar_novas_imagens(
classe=args.classe, cana=args.cana,
horario=args.horario,
grupo=args.grupo,
cor_classe_rgb=cor_rgb, cor_classe_rgb=cor_rgb,
fazer_copia_final=fazer_copia, fazer_copia_final=fazer_copia,
manifesto_csv=manifesto_csv, manifesto_csv=manifesto_csv,
orignais=args.from_originals
) )

View File

@ -12,7 +12,7 @@ with open("config.json", "r") as f:
MODELO = config["camera"] MODELO = config["camera"]
# Raiz das brutas agrupadas # 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 # Onde as máscaras novas (rotuladas externamente) são colocadas
PASTA_NEW_MASKS = os.path.join(MODELO, "dataset", "new_masks") 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." "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("--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).") ap.add_argument("--manifest", default="", help="Caminho para CSV de manifesto (opcional).")
return ap return ap
@ -191,7 +193,9 @@ if __name__ == "__main__":
manifesto = args.manifest if args.manifest.strip() else None manifesto = args.manifest if args.manifest.strip() else None
mover = not args.copy 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] Brutas (group root): {PASTA_BRUTAS_GROUP_ROOT}")
print(f"[INFO] New masks : {PASTA_NEW_MASKS}") print(f"[INFO] New masks : {PASTA_NEW_MASKS}")
print(f"[INFO] Mover máscaras? : {mover}") print(f"[INFO] Mover máscaras? : {mover}")

View File

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

View File

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

View File

@ -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 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": "<GRUPO>",
"group_dir": caminho_da_pasta_do_grupo
}
Estrutura esperada (relativa a root_brutas):
cana_alta/meio_dia/group/<GRUPO>/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/<GRUPO>/raws
group_dir = os.path.dirname(dirpath) # .../cana_x/horario/group/<GRUPO>
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] = <GRUPO>
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)

View File

@ -25,7 +25,7 @@ import shutil
import argparse import argparse
import numpy as np import numpy as np
from utils import carregar_labelmap_completo from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids
# ====================== Configurações base ====================== # ====================== Configurações base ======================
@ -33,19 +33,22 @@ def carregar_config_e_paths():
with open("config.json", "r", encoding="utf-8") as f: with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
MODELO = config.get("camera") MODELO = config.get("camera")
USE_MASKS2 = config.get("dual_head", False)
pasta_base = os.path.join(MODELO, "dataset") pasta_base = os.path.join(MODELO, "dataset")
labelmap_path = os.path.join(pasta_base, "labelmap.txt") labelmap_path = os.path.join(pasta_base, "labelmap.txt")
# Pastas origem/destino # Pastas origem/destino
PASTA_NEW_IMAGES = os.path.join(pasta_base, "original", "images") PASTA_NEW_IMAGES = os.path.join(pasta_base, "original", "images")
PASTA_NEW_MASKS = os.path.join(pasta_base, "original", "masks") 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") 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 # Extensões aceitas
EXT_IMAGENS = (".jpg", ".jpeg", ".png") EXT_IMAGENS = (".jpg", ".jpeg", ".png")
EXT_MASKS = (".png", ".jpg", ".jpeg") # prioridade será .png quando houver EXT_MASKS = (".png", ".jpg", ".jpeg") # prioridade será .png quando houver
EXT_MASKS2 = (".png", ".jpg", ".jpeg") # idem
# Manifesto padrão # Manifesto padrão
MANIFESTO_DEFAULT = "manifest.csv" MANIFESTO_DEFAULT = "manifest.csv"
@ -84,6 +87,25 @@ def mapear_masks_por_base(pasta_masks):
mapa[base] = cam mapa[base] = cam
return mapa 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): def localizar_imagem_por_base(pasta_imgs, base):
"""Retorna caminho da imagem correspondente ao base se existir.""" """Retorna caminho da imagem correspondente ao base se existir."""
for ext in EXT_IMAGENS: 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): def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
""" """
Extrai IDs de classes presentes na máscara. Versão corrigida e otimizada:
- 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), - Se máscara for 1 canal: np.unique direto -> IDs.
e mapeia usando cor_para_id[(R,G,B)] -> id. - Se for 3 canais:
* se assume_rgb=True: labelmap está em RGB,
mas OpenCV 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) Retorna: set(ids_presentes)
""" """
m = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED) m = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
if m is None: if m is None:
raise RuntimeError(f"Falha ao abrir máscara: {mask_path}") 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): 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) return set(int(v) for v in vals)
# 3 canais (OpenCV lê BGR) # ---------------------------
h, w, c = m.shape # CASO 2: máscara RGB
flat = m.reshape(-1, 3) # ---------------------------
uniq_bgr = np.unique(flat, axis=0) # Se o labelmap está em RGB (assume_rgb=True),
ids = set() # convertemos a imagem BGR->RGB para casar com as chaves.
for b, g, r in uniq_bgr: if assume_rgb:
if assume_rgb: img = cv2.cvtColor(m, cv2.COLOR_BGR2RGB)
key = (int(r), int(g), int(b)) # converte para RGB else:
else: # labelmap já está em BGR; OpenCV entrega BGR; deixa como está
key = (int(b), int(g), int(r)) # já em BGR no labelmap img = m
id_ = cor_para_id.get(key)
if id_ is not None: # Aqui as chaves de cor_para_id estão no MESMO espaço de cor da imagem.
try: mapa_rgb = cor_para_id
ids.add(int(id_)) max_classes = len(cor_para_id)
except Exception:
pass # ---------- 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 return ids
def montar_nome_grupo(ids_presentes, id_para_nome): def montar_nome_grupo(ids_presentes, id_para_nome):
""" """
Constrói o nome do grupo a partir dos nomes das classes dos IDs presentes. Constrói o nome do grupo respeitando a ordem natural dos IDs do labelmap.
Preferência de ordenação: chao < erva < cana; demais nomes em ordem alfabética. Ex: {0,1,2} -> chao_cana_obstaculo
""" """
nomes = [] #print(f"ids_presentes: {ids_presentes}")
for cid in sorted(ids_presentes): if not ids_presentes:
nome = id_para_nome.get(cid, str(cid)) return "sem_classe"
nomes.append(nome) nomes = [id_para_nome.get(cid, str(cid)) for cid in sorted(ids_presentes)]
return "_".join(nomes)
# aplicar ordenação preferida quando disponíveis def copiar_ou_mover(img_src, mask_src, dest_img_dir, dest_mask_dir, mover=False, mask2_src=None, dest_mask2_dir=None):
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):
garantir_pasta(dest_img_dir) garantir_pasta(dest_img_dir)
garantir_pasta(dest_mask_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) dst_img = nome_disponivel(dest_img_dir, base_img, img_ext)
new_base = os.path.splitext(os.path.basename(dst_img))[0] new_base = os.path.splitext(os.path.basename(dst_img))[0]
dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext) dst_mask = os.path.join(dest_mask_dir, new_base + mask_ext)
dst_mask2 = None
if os.path.exists(dst_mask): if os.path.exists(dst_mask):
# evita colisão invertendo a ordem do "único" para a máscara # 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): if os.path.exists(dst_img):
dst_img = nome_disponivel(dest_img_dir, new_base, img_ext) 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: if mover:
shutil.move(img_src, dst_img) shutil.move(img_src, dst_img)
shutil.move(mask_src, dst_mask) shutil.move(mask_src, dst_mask)
if mask2_src and dst_mask2:
shutil.move(mask2_src, dst_mask2)
else: else:
shutil.copy2(img_src, dst_img) shutil.copy2(img_src, dst_img)
shutil.copy2(mask_src, dst_mask) 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 ====================== # ====================== Pipeline principal ======================
def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT, def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT,
validar_dim=True, estrito=False, labelmap_bgr=False): validar_dim=True, estrito=False, labelmap_bgr=False):
# carrega config/paths # 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 # carrega labelmap completo
cor_para_id, _colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path) 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) 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 # garante pastas
garantir_pasta(PASTA_NEW_IMAGES) garantir_pasta(PASTA_NEW_IMAGES)
garantir_pasta(PASTA_NEW_MASKS) garantir_pasta(PASTA_NEW_MASKS)
garantir_pasta(PASTA_FINAL) 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 # indexa máscaras
mapa_masks = mapear_masks_por_base(PASTA_NEW_MASKS) mapa_masks = mapear_masks_por_base(PASTA_NEW_MASKS)
mapa_masks2 = mapear_masks2_por_base(PASTA_NEW_MASKS2) if usar_masks2 else {}
registros = [] registros = []
totais = {"total_masks":0, "processados":0, "pulados":0, "sem_imagem":0, totais = {"total_masks":0, "processados":0, "pulados":0, "sem_imagem":0, "sem_mask2":0, "dim_mismatch":0, "erros":0}
"dim_mismatch":0, "erros":0}
por_grupo = {} por_grupo = {}
for base, mask_path in sorted(mapa_masks.items()): for base, mask_path in sorted(mapa_masks.items()):
@ -228,6 +283,11 @@ def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT,
continue continue
try: 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: if validar_dim:
try: try:
img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
@ -245,6 +305,24 @@ def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT,
continue continue
else: else:
print(msg + " → copiando mesmo assim.") 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: except Exception as e_dim:
print(f"[AVISO] Falha ao validar dimensões: {e_dim} → copiando mesmo assim.") 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_base = os.path.join(PASTA_FINAL, grupo)
dest_img_dir = os.path.join(dest_base, "images") dest_img_dir = os.path.join(dest_base, "images")
dest_mask_dir = os.path.join(dest_base, "masks") 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 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 por_grupo[grupo] = por_grupo.get(grupo, 0) + 1
print(f"[OK] {os.path.basename(dst_img)} → grupo: {grupo}") 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: except Exception as e:
totais["erros"] += 1 totais["erros"] += 1
@ -277,7 +358,7 @@ def processar(modelo_cli=None, mover=False, manifesto=MANIFESTO_DEFAULT,
if manifesto and registros: if manifesto and registros:
with open(manifesto, "w", newline="", encoding="utf-8") as f: with open(manifesto, "w", newline="", encoding="utf-8") as f:
w = csv.writer(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) w.writerows(registros)
print(f"[MANIFESTO] {manifesto} salvo ({len(registros)} entradas).") print(f"[MANIFESTO] {manifesto} salvo ({len(registros)} entradas).")

View File

@ -49,11 +49,16 @@ MANIFESTO_DEFAULT = "manifest.csv"
def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True): def extrair_ids_presentes(mask_path, cor_para_id, assume_rgb=True):
""" """
Extração ultra-rápida de IDs presentes na máscara. Extração ultra-rápida de IDs presentes na máscara.
Estratégia: Estratégia:
1) Tenta máscara 1 canal np.unique (instrantâneo) 1) Máscara 1 canal np.unique (instantâneo).
2) Amostragem em grid converter_mask_rgb_para_ids unique 2) Máscara 3 canais:
3) Early exit quando todas classes forem encontradas - se assume_rgb=True: labelmap está em RGB,
4) Fallback full-scan apenas se necessário (raríssimo) 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) m = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
if m is None: 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 # CASO 1: máscara indexada (1 canal) — instantâneo
# ------------------------------------------------------- # -------------------------------------------------------
if len(m.shape) == 2 or (len(m.shape) == 3 and m.shape[2] == 1): 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) 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: 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: 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) max_classes = len(cor_para_id)
# ---------- AMOSTRAGEM ---------- # ---------- AMOSTRAGEM ----------
step = 8 step = 8 # pode virar 4 se quiser mais precisão, 16 se quiser mais velocidade
amostra = m[::step, ::step] # reduz drasticamente o custo amostra = img[::step, ::step]
amostra_ids = converter_mask_rgb_para_ids(amostra, mapa_rgb, ignore_id=255) 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) 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: if len(ids) >= max_classes:
return ids return ids
# ---------- FULL-SCAN OTIMIZADO (último caso) ---------- # ---------- 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) ids = set(int(x) for x in np.unique(full_ids) if x != 255)
return ids return ids
@ -146,10 +146,14 @@ def inferir_ignore_id(ignore_rgb, cor_para_id):
return None return None
def montar_nome_grupo(ids_presentes, id_para_nome): 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)] nomes = [id_para_nome.get(cid, str(cid)) for cid in sorted(ids_presentes)]
prefer = {"chao": 0, "erva": 1, "cana": 2} return "_".join(nomes)
nomes = sorted(nomes, key=lambda n: (prefer.get(n, 99), n))
return "_".join(nomes) if nomes else "sem_classe"
# ============================================================ # ============================================================

View File

@ -32,6 +32,7 @@ import argparse
with open("config.json", "r", encoding="utf-8") as f: with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
MODELO = config.get("camera", ".") MODELO = config.get("camera", ".")
USE_MASKS2 = config.get("dual_head", False)
# Pastas base # Pastas base
DATASET_BASE = os.path.join(MODELO, "dataset") 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) # Fallback (modo antigo, sem grupos)
ORIG_OLD_IMG = os.path.join(DATASET_BASE, "original", "images") ORIG_OLD_IMG = os.path.join(DATASET_BASE, "original", "images")
ORIG_OLD_MSK = os.path.join(DATASET_BASE, "original", "masks") 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_IMG = os.path.join(DATASET_BASE, "augmented", "images")
AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks") AUG_OLD_MSK = os.path.join(DATASET_BASE, "augmented", "masks")
AUG_OLD_MSK2 = os.path.join(DATASET_BASE, "augmented", "masks2")
# Extensões aceitas # Extensões aceitas
IMG_EXTS = (".jpg", ".jpeg", ".png") IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXTS = (".png", ".jpg", ".jpeg") # manter prioridade PNG quando possível MSK_EXTS = (".png", ".jpg", ".jpeg") # manter prioridade PNG quando possível
MSK2_EXTS = (".png", ".jpg", ".jpeg")
def garantir_dir(p): def garantir_dir(p):
os.makedirs(p, exist_ok=True) os.makedirs(p, exist_ok=True)
@ -86,7 +90,9 @@ train_tf = A.Compose([
A.RandomSunFlare(p=0.10), A.RandomSunFlare(p=0.10),
A.ChannelShuffle(p=0.05), A.ChannelShuffle(p=0.05),
A.CoarseDropout(max_holes=6, max_height=16, max_width=16, p=0.10), 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): def load_rgb(path):
# cv2 lê BGR → converte para RGB # cv2 lê BGR → converte para RGB
@ -131,31 +137,60 @@ def map_masks_by_base(msk_dir):
by_base[base] = cand by_base[base] = cand
return by_base return by_base
def ensure_aug_dirs(group_name=None): def map_masks2_by_base(msk2_dir):
"""Cria diretórios de saída para o grupo ou modo antigo.""" """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: if group_name:
img_out = os.path.join(AUG_GROUP_ROOT, group_name, "images") img_out = os.path.join(AUG_GROUP_ROOT, group_name, "images")
msk_out = os.path.join(AUG_GROUP_ROOT, group_name, "masks") 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: else:
img_out = AUG_OLD_IMG img_out = AUG_OLD_IMG
msk_out = AUG_OLD_MSK msk_out = AUG_OLD_MSK
msk2_out = AUG_OLD_MSK2 if use_masks2 else None
garantir_dir(img_out) garantir_dir(img_out)
garantir_dir(msk_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_img, img_ext = os.path.splitext(os.path.basename(img_path))
base_msk, msk_ext = os.path.splitext(os.path.basename(msk_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 # padroniza pelo base da imagem
base = base_img base = base_img
img = load_rgb(img_path) img = load_rgb(img_path)
msk = load_rgb(msk_path) msk = load_rgb(msk_path)
msk2 = load_rgb(msk2_path) if msk2_path else None
gen = 0 gen = 0
for i in range(copies): 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"] img_aug = aug["image"]
msk_aug = aug["mask"] 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}") out_msk = os.path.join(msk_out_dir, f"{base}_aug_{i:02d}{msk_ext}")
save_rgb(out_img, img_aug) save_rgb(out_img, img_aug)
save_rgb(out_msk, msk_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 gen += 1
return gen return gen
@ -170,13 +211,16 @@ def process_group(group_name, copies):
"""Processa um grupo único (images/masks dentro de ORIG_GROUP_ROOT/<group_name>/).""" """Processa um grupo único (images/masks dentro de ORIG_GROUP_ROOT/<group_name>/)."""
img_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "images") img_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "images")
msk_dir = os.path.join(ORIG_GROUP_ROOT, group_name, "masks") 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)): 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.") print(f"[WARN] Grupo '{group_name}' inválido (sem images/masks). Pulando.")
return 0 return 0
imgs = [f for f in os.listdir(img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS] 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) 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 count = 0
for img_file in sorted(imgs): for img_file in sorted(imgs):
@ -185,13 +229,18 @@ def process_group(group_name, copies):
if not msk_file: if not msk_file:
print(f"[WARN] [{group_name}] Máscara não encontrada para {img_file}, pulando.") print(f"[WARN] [{group_name}] Máscara não encontrada para {img_file}, pulando.")
continue 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: try:
count += augment_pair( count += augment_pair(
os.path.join(img_dir, img_file), os.path.join(img_dir, img_file),
msk_file, msk_file,
img_out_dir, img_out_dir,
msk_out_dir, msk_out_dir,
copies=copies copies=copies,
msk2_path=msk2_file,
msk2_out_dir=msk2_out_dir
) )
except Exception as e: except Exception as e:
print(f"[ERRO] [{group_name}] {img_file}: {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] 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) 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 count = 0
for img_file in sorted(imgs): for img_file in sorted(imgs):
@ -215,13 +266,18 @@ def process_legacy(copies):
if not msk_file: if not msk_file:
print(f"[WARN] (legacy) Máscara não encontrada para {img_file}, pulando.") print(f"[WARN] (legacy) Máscara não encontrada para {img_file}, pulando.")
continue 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: try:
count += augment_pair( count += augment_pair(
os.path.join(ORIG_OLD_IMG, img_file), os.path.join(ORIG_OLD_IMG, img_file),
msk_file, msk_file,
img_out_dir, img_out_dir,
msk_out_dir, msk_out_dir,
copies=copies copies=copies,
msk2_path=msk2_file,
msk2_out_dir=msk2_out_dir
) )
except Exception as e: except Exception as e:
print(f"[ERRO] (legacy) {img_file}: {e}") print(f"[ERRO] (legacy) {img_file}: {e}")

View File

@ -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: with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
MODELO = config["camera"] MODELO = config["camera"]
USE_MASKS2 = config["dual_head"]
RESOLUCAO = tuple(config["resolucao"]) # [W, H] ou [width, height] RESOLUCAO = tuple(config["resolucao"]) # [W, H] ou [width, height]
pasta_base = os.path.join(MODELO, "dataset") pasta_base = os.path.join(MODELO, "dataset")
labelmap_path = os.path.join(pasta_base, "labelmap.txt") labelmap_path = os.path.join(pasta_base, "labelmap.txt")
@ -43,6 +44,7 @@ FONTES = ["original", "augmented"]
# Extensões aceitas # Extensões aceitas
IMG_EXTS = (".jpg", ".jpeg", ".png") IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXTS = (".png", ".jpg", ".jpeg") # preferir .png MSK_EXTS = (".png", ".jpg", ".jpeg") # preferir .png
MSK2_EXTS = (".png", ".jpg", ".jpeg") # idem
def infer_ignore_id(ignore_rgb, default_id=255): 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 by_base[base] = cand
return by_base 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, 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 = ""): 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).""" """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 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 é 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): 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 .../<fonte>/group/ agrupando por cada subpasta de grupo.""" """Processa uma raiz do tipo .../<fonte>/group/ agrupando por cada subpasta de grupo."""
total = 0 total = 0
@ -155,13 +213,17 @@ def process_group_root(fonte_root: str, fonte_nome: str, cor_para_id, ignore_id:
continue continue
in_img_dir = os.path.join(fonte_root, grupo, "images") in_img_dir = os.path.join(fonte_root, grupo, "images")
in_msk_dir = os.path.join(fonte_root, grupo, "masks") 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)): 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}") print(f"[WARN] Grupo inválido (sem images/masks): {grupo}")
continue continue
out_img_dir = os.path.join(out_root, grupo, "images") out_img_dir = os.path.join(out_root, grupo, "images")
out_msk_dir = os.path.join(out_root, grupo, "masks") 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) 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] imgs = [f for f in os.listdir(in_img_dir) if os.path.splitext(f.lower())[1] in IMG_EXTS]
n = len(imgs) 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) base, _ = os.path.splitext(fname)
caminho_rgb = os.path.join(in_img_dir, fname) caminho_rgb = os.path.join(in_img_dir, fname)
caminho_mask = msk_map.get(base) caminho_mask = msk_map.get(base)
caminho_mask2 = msk2_map.get(base) if usar_masks2 else None
ok = normalize_pair( ok = normalize_pair(
caminho_rgb, caminho_mask, cor_para_id, ignore_id, caminho_rgb, caminho_mask, cor_para_id, ignore_id,
out_img_dir, out_msk_dir, dim, prefix=f"{fonte_nome}_" 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: if ok:
total += 1 total += 1
print(f"[{fonte_nome} | {grupo} | {nome_res}] {i}/{n}{fname}") 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(): for nome_res, dim in RESOLUCOES.items():
out_img_dir = os.path.join(pasta_base, nome_res, "images") out_img_dir = os.path.join(pasta_base, nome_res, "images")
out_msk_dir = os.path.join(pasta_base, nome_res, "masks") 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) 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] imgs = [f for f in os.listdir(legacy_img) if os.path.splitext(f.lower())[1] in IMG_EXTS]
n = len(imgs) n = len(imgs)
for i, fname in enumerate(sorted(imgs), 1): for i, fname in enumerate(sorted(imgs), 1):
base, _ = os.path.splitext(fname) base, _ = os.path.splitext(fname)
caminho_rgb = os.path.join(legacy_img, fname) caminho_rgb = os.path.join(legacy_img, fname)
caminho_mask = msk_map.get(base) caminho_mask = msk_map.get(base)
caminho_mask2 = msk2_map.get(base) if usar_masks2 else None
ok = normalize_pair( ok = normalize_pair(
caminho_rgb, caminho_mask, cor_para_id, ignore_id, caminho_rgb, caminho_mask, cor_para_id, ignore_id,
out_img_dir, out_msk_dir, dim, prefix=f"{fonte_nome}_" 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: if ok:
total += 1 total += 1
print(f"[{fonte_nome} | legacy | {nome_res}] {i}/{n}{fname}") 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) # 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) 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) ignore_id = infer_ignore_id(ignore_rgb, default_id=255)
print(cor_para_id, _colormap_rgb, _id_para_nome)
total_geral = 0 total_geral = 0
# === ORIGINAL === # === ORIGINAL ===

View File

@ -36,6 +36,7 @@ import argparse
with open("config.json", "r", encoding="utf-8") as f: with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
MODELO = config.get("camera") MODELO = config.get("camera")
USE_MASKS2 = config.get("dual_head", False)
RESOLUCAO = tuple(config.get("resolucao")) RESOLUCAO = tuple(config.get("resolucao"))
# Pastas # Pastas
@ -77,6 +78,10 @@ def mask_from_image_name(img_name):
base, _ = os.path.splitext(img_name) base, _ = os.path.splitext(img_name)
return base + MSK_EXT 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): def classify_source_and_family(filename_no_ext):
""" """
Retorna (source, family_key) 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 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) 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 moved = 0
for nome in nomes: for nome in nomes:
mask_name = mask_from_image_name(nome) 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 continue
shutil.copy2(src_img, os.path.join(dst_img_dir, nome)) shutil.copy2(src_img, os.path.join(dst_img_dir, nome))
shutil.copy2(src_msk, os.path.join(dst_msk_dir, mask_name)) 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 moved += 1
return moved return moved
def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None): 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_img_dir = os.path.join(pasta_origem, group_name, "images")
src_msk_dir = os.path.join(pasta_origem, group_name, "masks") 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) familias = build_family_index(src_img_dir, src_msk_dir)
# apenas famílias que têm ORIGINAL para participar de val/test # 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_img = os.path.join(pasta_destino, "test", "group", group_name, "images")
dest_test_msk = os.path.join(pasta_destino, "test", "group", group_name, "masks") 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) dest_train_msk2 = os.path.join(pasta_destino, "train", "group", group_name, "masks2") if use_msk2 else None
m_val = copiar(nomes_val, src_img_dir, src_msk_dir, dest_val_img, dest_val_msk) dest_val_msk2 = os.path.join(pasta_destino, "val", "group", group_name, "masks2") if use_msk2 else None
m_test = copiar(nomes_test, src_img_dir, src_msk_dir, dest_test_img, dest_test_msk) 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}") 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} return {"train": m_train, "val": m_val, "test": m_test, "familias": total_familias}

View File

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

View File

@ -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": <tensor/ndarray>, "mask": <tensor/ndarray>, ...}
# 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/<filename>)
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()

View File

@ -213,7 +213,8 @@ def main():
cv2.destroyAllWindows() cv2.destroyAllWindows()
else: else:
# === Modo imagens (agrupado + fallback) === # === 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) image_paths, mask_paths, groups_idx = collect_pairs_grouped(test_root, want_groups=args.groups)
if not image_paths: if not image_paths:

View File

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

View File

@ -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):
"""
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/<qualquer>/images/*.png
split/test/group/<qualquer>/masks/*.png
split/test/group/<qualquer>/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/<grupo_ou_subgrupo>
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()

View File

@ -2,6 +2,7 @@
"camera": "gal5000", "camera": "gal5000",
"modelo": "segformer_b0", "modelo": "segformer_b0",
"model_name": "ndvi_big", "model_name": "ndvi_big",
"dual_head": false,
"main_class_name": "erva", "main_class_name": "erva",
"es_classes": "", "es_classes": "",
"model_to_use": "geral", "model_to_use": "geral",

View File

@ -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"
}

View File

@ -0,0 +1,4 @@
# label:color_rgb:parts:actions
naopulverizar:128,0,0::
pulverizar:0,128,0::
ignore:255,255,255::

View File

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

File diff suppressed because it is too large Load Diff

View File

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