145 lines
5.1 KiB
Python
145 lines
5.1 KiB
Python
import cv2
|
|
import numpy as np
|
|
|
|
# ----------------------------
|
|
# Helpers LABELMAP
|
|
# ----------------------------
|
|
def carregar_labelmap_completo(caminho):
|
|
cor_para_id = {}
|
|
id_para_nome = {}
|
|
cores_rgb = []
|
|
|
|
with open(caminho, 'r') as arquivo:
|
|
idx = 0
|
|
for linha in arquivo:
|
|
if linha.startswith("#") or not linha.strip():
|
|
continue
|
|
partes = linha.strip().split(':')
|
|
if len(partes) >= 2:
|
|
nome_classe, cor_rgb_str = partes[0], partes[1]
|
|
r, g, b = map(int, cor_rgb_str.split(','))
|
|
cor_rgb = (r, g, b)
|
|
|
|
if nome_classe.lower() == "ignore":
|
|
ignore_rgb = cor_rgb
|
|
continue # NÃO adiciona ignore no LUT de classes
|
|
|
|
cor_para_id[cor_rgb] = idx
|
|
cores_rgb.append(cor_rgb)
|
|
id_para_nome[idx] = nome_classe
|
|
idx += 1
|
|
|
|
#print(f"Mapa: {cor_para_id}")
|
|
#print(f"Colormap RGB: {cores_rgb}")
|
|
#print(f"Classes: {id_para_nome}")
|
|
#print(f"Ignore RGB: {ignore_rgb}")
|
|
|
|
return cor_para_id, cores_rgb, id_para_nome, ignore_rgb
|
|
|
|
def converter_mask_rgb_para_ids(img_rgb, mapa_rgb, ignore_id):
|
|
# Cria um mapa 256^3 para IDs (usa int32 para indexar)
|
|
lut = np.full((256**3,), ignore_id, dtype=np.uint8)
|
|
|
|
for cor, classe_id in mapa_rgb.items():
|
|
r, g, b = cor
|
|
lut[(r << 16) + (g << 8) + b] = classe_id
|
|
|
|
# Converte RGB para índice único
|
|
flat_idx = (img_rgb[:,:,0].astype(np.int32) << 16) + \
|
|
(img_rgb[:,:,1].astype(np.int32) << 8) + \
|
|
img_rgb[:,:,2].astype(np.int32)
|
|
|
|
# Aplica LUT vetorizada
|
|
return lut[flat_idx]
|
|
|
|
def converter_mask_ids_para_rgb(mask_ids: np.ndarray, colormap_rgb: list, ignore_id: int = 255) -> np.ndarray:
|
|
# Criar lookup table (256 cores possíveis)
|
|
lut = np.zeros((256, 3), dtype=np.uint8)
|
|
for i, color in enumerate(colormap_rgb):
|
|
lut[i] = color
|
|
lut[ignore_id] = (255, 255, 255)
|
|
|
|
# Aplicar LUT direto (vetorizado)
|
|
return lut[mask_ids]
|
|
|
|
def converter_mask_ids_para_bgr(mask_ids: np.ndarray, colormap_rgb: list, ignore_id: int = 255) -> np.ndarray:
|
|
"""
|
|
Converte máscara de IDs para imagem BGR (uint8),
|
|
pronta para uso com OpenCV.
|
|
"""
|
|
lut = np.zeros((256, 3), dtype=np.uint8)
|
|
|
|
for i, (r, g, b) in enumerate(colormap_rgb):
|
|
lut[i] = (b, g, r) # RGB -> BGR
|
|
|
|
lut[ignore_id] = (255, 255, 255) # branco em BGR = RGB
|
|
|
|
return lut[mask_ids]
|
|
|
|
def desenhar_legenda_vertical(colormap_rgb, classes, largura=200):
|
|
"""
|
|
Retorna uma imagem com a legenda das classes (cor + nome)
|
|
"""
|
|
nomes_classes = [classes[i] for i in range(len(classes))]
|
|
altura_por_classe = 30
|
|
altura_total = altura_por_classe * len(colormap_rgb)
|
|
legenda = np.ones((altura_total, largura, 3), dtype=np.uint8) * 255
|
|
|
|
for idx, (rgb, nome) in enumerate(zip(colormap_rgb, nomes_classes)):
|
|
y = idx * altura_por_classe
|
|
color = tuple(int(c) for c in rgb)
|
|
cv2.rectangle(legenda, (10, y + 5), (30, y + 25), color, -1)
|
|
cv2.putText(legenda, nome, (40, y + 20),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1, cv2.LINE_AA)
|
|
|
|
return legenda
|
|
|
|
def desenhar_legenda_horizontal(colormap_rgb, classes, altura=30, largura_por_classe=120):
|
|
"""
|
|
Retorna uma imagem com a legenda das classes (cor + nome), em uma única linha horizontal
|
|
"""
|
|
nomes_classes = [classes[i] for i in range(len(classes))]
|
|
largura_total = largura_por_classe * len(colormap_rgb)
|
|
legenda = np.ones((altura, largura_total, 3), dtype=np.uint8) * 255 # faixa branca
|
|
|
|
for idx, (rgb, nome) in enumerate(zip(colormap_rgb, nomes_classes)):
|
|
x = idx * largura_por_classe
|
|
color = tuple(int(c) for c in rgb)
|
|
# Retângulo colorido
|
|
cv2.rectangle(legenda, (x + 10, 5), (x + 30, 25), color, -1)
|
|
# Texto da classe
|
|
cv2.putText(legenda, nome, (x + 35, 20),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1, cv2.LINE_AA)
|
|
|
|
return legenda
|
|
|
|
def _infer_ignore_id(ignore_rgb, default_id=255):
|
|
import numpy as _np
|
|
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
|
|
|
|
# ----------------------------
|
|
# Helpers ROI
|
|
# ----------------------------
|
|
def compute_roi_indices(H: int, zona_inicio: float, faixa_atuacao: float):
|
|
y_inicio = int((1.0 - zona_inicio) * H)
|
|
y_fim = int((1.0 - (zona_inicio + faixa_atuacao)) * H)
|
|
y_fim = max(0, min(H, y_fim))
|
|
y_inicio = max(0, min(H, y_inicio))
|
|
if y_fim >= y_inicio:
|
|
y_fim = max(0, y_inicio - 1)
|
|
return y_fim, y_inicio
|
|
|
|
def resize_keep_width(img: np.ndarray, new_w: int, min_h: int, interpolation: int) -> np.ndarray:
|
|
h, w = img.shape[:2]
|
|
new_h = int(round(new_w * (h / w)))
|
|
if min_h is not None and new_h < min_h:
|
|
new_h = min_h
|
|
return cv2.resize(img, (new_w, new_h), interpolation=interpolation)
|