Ajustado scripts de treinamento e health worker
This commit is contained in:
parent
4a1561fb16
commit
cb0ebd8b1d
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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: 300–500
|
||||||
|
self._ping_count = 2 # recomendado: 1 (janela já suaviza)
|
||||||
|
self._ping_min_interval = 0.6 # não precisa pingar a cada 100ms
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def _start_ping_async_if_needed(self):
|
||||||
|
"""Dispara ping em background se não houver um em andamento."""
|
||||||
|
now = time.time()
|
||||||
|
|
||||||
|
# respeita intervalo mínimo (evita spam de ping)
|
||||||
|
if (now - self._last_ping_ts) < self._ping_min_interval:
|
||||||
|
return
|
||||||
|
|
||||||
|
# já tem um ping rodando?
|
||||||
|
if self._ping_future is not None and not self._ping_future.done():
|
||||||
|
return
|
||||||
|
|
||||||
|
base_ip = self.get_base_ip()
|
||||||
|
if not base_ip:
|
||||||
|
# sem rota ainda -> considera "sem dados" (não bloqueia)
|
||||||
|
with self._ping_lock:
|
||||||
|
self._last_ping_ts = now
|
||||||
|
self._last_ping_ok = False
|
||||||
|
self._last_ping_rtt_ms = None
|
||||||
|
self._last_ping_loss_pct = 100.0
|
||||||
|
return
|
||||||
|
|
||||||
|
# dispara o ping em background
|
||||||
|
self._ping_future = self._ping_executor.submit(
|
||||||
|
self._ping_once, base_ip, self._ping_count, self._ping_timeout_ms
|
||||||
|
)
|
||||||
|
|
||||||
|
def _consume_ping_result_if_ready(self):
|
||||||
|
"""Se o ping terminou, atualiza o snapshot sem bloquear o loop."""
|
||||||
|
if self._ping_future is None or not self._ping_future.done():
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
# _ping_once deve retornar: (ok: bool, avg_rtt_ms: Optional[float], loss_pct: float)
|
||||||
|
ok, avg_rtt_ms, loss_pct = self._ping_future.result(timeout=0)
|
||||||
|
except Exception:
|
||||||
|
ok, avg_rtt_ms, loss_pct = False, None, 100.0
|
||||||
|
|
||||||
|
with self._ping_lock:
|
||||||
|
self._last_ping_ts = time.time()
|
||||||
|
self._last_ping_ok = ok
|
||||||
|
self._last_ping_rtt_ms = avg_rtt_ms
|
||||||
|
self._last_ping_loss_pct = loss_pct
|
||||||
|
|
||||||
|
# limpa
|
||||||
|
self._ping_future = None
|
||||||
|
|
||||||
|
def _get_ping_snapshot(self):
|
||||||
|
"""Lê o último resultado do ping (snapshot thread-safe)."""
|
||||||
|
with self._ping_lock:
|
||||||
|
return (self._last_ping_ok, self._last_ping_rtt_ms, self._last_ping_loss_pct, self._last_ping_ts)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def atualizar_saude(self):
|
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()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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})")
|
|
||||||
|
|
@ -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
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -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}")
|
||||||
|
|
|
||||||
|
|
@ -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))
|
|
||||||
|
|
@ -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()
|
||||||
|
|
@ -0,0 +1,328 @@
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import shutil
|
||||||
|
import argparse
|
||||||
|
import csv
|
||||||
|
|
||||||
|
# ================= CONFIG =================
|
||||||
|
|
||||||
|
with open("config.json", "r") as f:
|
||||||
|
config = json.load(f)
|
||||||
|
|
||||||
|
MODELO = config["camera"]
|
||||||
|
|
||||||
|
# Raiz das brutas (todas as canas/horários/grupos)
|
||||||
|
PASTA_BRUTAS_ROOT = os.path.join(MODELO, "dataset", "brutas")
|
||||||
|
|
||||||
|
# Onde você coloca os previews selecionados (tudo misturado)
|
||||||
|
PASTA_SELECTED_PREVIEWS = os.path.join(MODELO, "dataset", "selected_previews")
|
||||||
|
|
||||||
|
# Destino final, organizado por grupo:
|
||||||
|
# dataset/original/group/{GRUPO}/{previews,raws,metas,masks}
|
||||||
|
PASTA_ORIGINAL_GROUP_ROOT = os.path.join(MODELO, "dataset", "original", "group")
|
||||||
|
|
||||||
|
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
|
||||||
|
EXT_RAWS = (".raw",)
|
||||||
|
EXT_MASKS = (".png",)
|
||||||
|
# Ajusta se suas metas tiverem outra extensão
|
||||||
|
EXT_METAS = (".json", ".yml", ".yaml", ".txt", ".csv")
|
||||||
|
|
||||||
|
|
||||||
|
# ================= HELPERS =================
|
||||||
|
|
||||||
|
def garantir_pasta(p: str):
|
||||||
|
os.makedirs(p, exist_ok=True)
|
||||||
|
|
||||||
|
|
||||||
|
def nome_disponivel(dest_dir: str, base: str, ext: str) -> str:
|
||||||
|
"""
|
||||||
|
Devolve um caminho disponível em dest_dir com base + ext.
|
||||||
|
Se já existir, adiciona _001, _002, ...
|
||||||
|
"""
|
||||||
|
p = os.path.join(dest_dir, base + ext)
|
||||||
|
if not os.path.exists(p):
|
||||||
|
return p
|
||||||
|
|
||||||
|
i = 1
|
||||||
|
while True:
|
||||||
|
p = os.path.join(dest_dir, f"{base}_{i:03d}{ext}")
|
||||||
|
if not os.path.exists(p):
|
||||||
|
return p
|
||||||
|
i += 1
|
||||||
|
|
||||||
|
|
||||||
|
def buscar_por_base(pasta: str, base: str, exts: tuple[str, ...]) -> str | None:
|
||||||
|
"""
|
||||||
|
Procura um arquivo em 'pasta' com o mesmo 'base' e qualquer extensão em 'exts'.
|
||||||
|
Retorna o caminho completo ou None.
|
||||||
|
"""
|
||||||
|
if not os.path.isdir(pasta):
|
||||||
|
return None
|
||||||
|
|
||||||
|
for nome in os.listdir(pasta):
|
||||||
|
nome_lower = nome.lower()
|
||||||
|
root, ext = os.path.splitext(nome_lower)
|
||||||
|
if root == base.lower() and ext in exts:
|
||||||
|
return os.path.join(pasta, nome)
|
||||||
|
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def indexar_brutas_por_raw(root_brutas: str):
|
||||||
|
"""
|
||||||
|
Varre dataset/brutas recursivamente, olhando pastas 'raws' e montando índice:
|
||||||
|
|
||||||
|
base -> {
|
||||||
|
"raw": caminho_raw,
|
||||||
|
"cana": "cana_alta" | "cana_baixa" | ... (se conseguir inferir),
|
||||||
|
"horario": "meio_dia" | "cedo" | ... (se conseguir inferir),
|
||||||
|
"grupo": "<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)
|
||||||
|
|
@ -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 lê BGR -> convertemos BGR -> RGB.
|
||||||
|
* se assume_rgb=False: labelmap está em BGR,
|
||||||
|
mantemos BGR como está.
|
||||||
|
- Usa converter_mask_rgb_para_ids em amostragem + fallback full-scan.
|
||||||
|
|
||||||
Retorna: set(ids_presentes)
|
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).")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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"
|
|
||||||
|
|
||||||
|
|
||||||
# ============================================================
|
# ============================================================
|
||||||
|
|
|
||||||
|
|
@ -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}")
|
||||||
|
|
|
||||||
|
|
@ -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 já é uma máscara "pronta".
|
||||||
|
- Se for RGB/BGR (3 canais): converte para cinza e faz threshold (0/255) antes de redimensionar.
|
||||||
|
- Se for 1 canal: mantém, faz threshold (0/255).
|
||||||
|
- Redimensiona com INTER_NEAREST.
|
||||||
|
Saída sempre .png com o mesmo nome base/prefixo do arquivo de imagem.
|
||||||
|
"""
|
||||||
|
if not caminho_mask2 or not os.path.isfile(caminho_mask2):
|
||||||
|
return False
|
||||||
|
|
||||||
|
nome = os.path.basename(caminho_rgb)
|
||||||
|
nome_saida = f"{prefix}{nome}" if prefix else nome
|
||||||
|
for ext in (".jpg", ".jpeg", ".png"):
|
||||||
|
if nome_saida.lower().endswith(ext):
|
||||||
|
nome_saida = nome_saida[: -len(ext)] + ".png"
|
||||||
|
break
|
||||||
|
|
||||||
|
m2 = cv2.imread(caminho_mask2, cv2.IMREAD_UNCHANGED)
|
||||||
|
if m2 is None:
|
||||||
|
print(f"[!] Erro ao ler máscara2: {caminho_mask2}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
if len(m2.shape) == 3:
|
||||||
|
# BGR/RGB -> gray
|
||||||
|
m2g = cv2.cvtColor(m2, cv2.COLOR_BGR2GRAY)
|
||||||
|
else:
|
||||||
|
m2g = m2
|
||||||
|
|
||||||
|
# binariza para 0/255 (evita lixo de compressão)
|
||||||
|
_, m2bin = cv2.threshold(m2g, 127, 255, cv2.THRESH_BINARY)
|
||||||
|
m2res = cv2.resize(m2bin, dim, interpolation=cv2.INTER_NEAREST)
|
||||||
|
garantir_dir(out_msk2_dir)
|
||||||
|
cv2.imwrite(os.path.join(out_msk2_dir, nome_saida), m2res)
|
||||||
|
return True
|
||||||
|
|
||||||
def process_group_root(fonte_root: str, fonte_nome: str, cor_para_id, ignore_id: int, groups_except: str = None):
|
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 ===
|
||||||
|
|
|
||||||
|
|
@ -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}
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
@ -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()
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
@ -0,0 +1,570 @@
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import glob
|
||||||
|
import argparse
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import List, Optional
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
from transformers import SegformerForSemanticSegmentation
|
||||||
|
|
||||||
|
|
||||||
|
# -----------------------------
|
||||||
|
# Utils
|
||||||
|
# -----------------------------
|
||||||
|
def find_device():
|
||||||
|
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||||
|
|
||||||
|
def load_labelmap_lut_bgr(labelmap_path: str):
|
||||||
|
"""
|
||||||
|
Lê labelmap no formato:
|
||||||
|
classe:r,g,b
|
||||||
|
Ignora linhas vazias e comentários '#'
|
||||||
|
Se existir 'ignore', retorna ignore_id=255 e ignore_bgr.
|
||||||
|
|
||||||
|
Retorna:
|
||||||
|
lut_bgr: dict[int, tuple(b,g,r)]
|
||||||
|
id_to_name: dict[int, str]
|
||||||
|
ignore_id: int (default 255)
|
||||||
|
"""
|
||||||
|
lut_bgr = {}
|
||||||
|
id_to_name = {}
|
||||||
|
ignore_id = 255
|
||||||
|
ignore_bgr = (180, 0, 180) # fallback
|
||||||
|
|
||||||
|
idx = 0
|
||||||
|
with open(labelmap_path, "r", encoding="utf-8") as f:
|
||||||
|
for line in f:
|
||||||
|
line = line.strip()
|
||||||
|
if not line or line.startswith("#"):
|
||||||
|
continue
|
||||||
|
parts = line.split(":")
|
||||||
|
if len(parts) < 2:
|
||||||
|
continue
|
||||||
|
name = parts[0].strip()
|
||||||
|
rgb_str = parts[1].strip()
|
||||||
|
r, g, b = map(int, rgb_str.split(","))
|
||||||
|
bgr = (b, g, r)
|
||||||
|
|
||||||
|
if name.lower() == "ignore":
|
||||||
|
ignore_bgr = bgr
|
||||||
|
continue
|
||||||
|
|
||||||
|
lut_bgr[idx] = bgr
|
||||||
|
id_to_name[idx] = name
|
||||||
|
idx += 1
|
||||||
|
|
||||||
|
lut_bgr[ignore_id] = ignore_bgr
|
||||||
|
return lut_bgr, id_to_name, ignore_id
|
||||||
|
|
||||||
|
def imread_bgr(path: str) -> np.ndarray:
|
||||||
|
img = cv2.imread(path, cv2.IMREAD_COLOR)
|
||||||
|
if img is None:
|
||||||
|
raise RuntimeError(f"Falha ao ler imagem: {path}")
|
||||||
|
return img
|
||||||
|
|
||||||
|
|
||||||
|
def imread_gray(path: str) -> np.ndarray:
|
||||||
|
m = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
|
||||||
|
if m is None:
|
||||||
|
raise RuntimeError(f"Falha ao ler máscara: {path}")
|
||||||
|
return m
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_uint8(x: np.ndarray) -> np.ndarray:
|
||||||
|
if x.dtype == np.uint8:
|
||||||
|
return x
|
||||||
|
x = np.clip(x, 0, 255).astype(np.uint8)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
def normalize01(x: np.ndarray, eps: float = 1e-6) -> np.ndarray:
|
||||||
|
x = x.astype(np.float32)
|
||||||
|
mn, mx = float(x.min()), float(x.max())
|
||||||
|
return (x - mn) / (mx - mn + eps)
|
||||||
|
|
||||||
|
|
||||||
|
def colorize_seg(mask_ids: np.ndarray, lut_bgr: dict) -> np.ndarray:
|
||||||
|
h, w = mask_ids.shape[:2]
|
||||||
|
out = np.zeros((h, w, 3), dtype=np.uint8)
|
||||||
|
|
||||||
|
# pinta cada id presente
|
||||||
|
for cid, bgr in lut_bgr.items():
|
||||||
|
out[mask_ids == cid] = bgr
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def overlay_mask(img_bgr: np.ndarray, mask_bgr: np.ndarray, alpha: float = 0.45) -> np.ndarray:
|
||||||
|
return cv2.addWeighted(img_bgr, 1.0 - alpha, mask_bgr, alpha, 0.0)
|
||||||
|
|
||||||
|
|
||||||
|
def put_hud(img: np.ndarray, lines: List[str]) -> np.ndarray:
|
||||||
|
out = img.copy()
|
||||||
|
y = 22
|
||||||
|
for s in lines:
|
||||||
|
cv2.putText(out, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (15, 15, 15), 3, cv2.LINE_AA)
|
||||||
|
cv2.putText(out, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (245, 245, 245), 1, cv2.LINE_AA)
|
||||||
|
y += 22
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
# -----------------------------
|
||||||
|
# Corridor Head (igual ao treino)
|
||||||
|
# -----------------------------
|
||||||
|
class CorridorHead(nn.Module):
|
||||||
|
"""
|
||||||
|
Head simples:
|
||||||
|
feat (B,C,H,W) + logits_seg (B,K,H,W) -> concat -> convs -> 1 canal (logit corredor)
|
||||||
|
"""
|
||||||
|
def __init__(self, feat_ch: int, num_classes: int, hidden: int = 256, dropout: float = 0.1):
|
||||||
|
super().__init__()
|
||||||
|
in_ch = feat_ch + num_classes
|
||||||
|
self.in_ch = in_ch
|
||||||
|
self.net = nn.Sequential(
|
||||||
|
nn.Conv2d(in_ch, hidden, kernel_size=3, padding=1),
|
||||||
|
nn.BatchNorm2d(hidden),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
nn.Dropout2d(dropout),
|
||||||
|
nn.Conv2d(hidden, hidden, kernel_size=3, padding=1),
|
||||||
|
nn.BatchNorm2d(hidden),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
nn.Dropout2d(dropout),
|
||||||
|
nn.Conv2d(hidden, 1, kernel_size=1)
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, feat: torch.Tensor, logits_seg: torch.Tensor) -> torch.Tensor:
|
||||||
|
x = torch.cat([feat, logits_seg], dim=1)
|
||||||
|
return self.net(x)
|
||||||
|
|
||||||
|
|
||||||
|
# -----------------------------
|
||||||
|
# Dataset discovery (split/test/group/**)
|
||||||
|
# -----------------------------
|
||||||
|
@dataclass
|
||||||
|
class Sample:
|
||||||
|
img_path: str
|
||||||
|
mask_path: Optional[str]
|
||||||
|
mask2_path: Optional[str]
|
||||||
|
group_name: str
|
||||||
|
filename: str
|
||||||
|
|
||||||
|
|
||||||
|
IMG_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp")
|
||||||
|
|
||||||
|
|
||||||
|
def discover_samples(split_root: str) -> List[Sample]:
|
||||||
|
"""
|
||||||
|
Espera estrutura:
|
||||||
|
split/test/group/<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()
|
||||||
|
|
@ -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",
|
||||||
|
|
|
||||||
|
|
@ -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"
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,4 @@
|
||||||
|
# label:color_rgb:parts:actions
|
||||||
|
naopulverizar:128,0,0::
|
||||||
|
pulverizar:0,128,0::
|
||||||
|
ignore:255,255,255::
|
||||||
|
|
@ -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
|
|
@ -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()
|
||||||
Loading…
Reference in New Issue