agrobot_base/Python/OAK/datasets/utils.py

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

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