345 lines
12 KiB
Python
345 lines
12 KiB
Python
import json
|
|
import os
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import cv2
|
|
from PIL import Image, ImageDraw
|
|
import matplotlib.pyplot as plt
|
|
from utils import carregar_labelmap_completo
|
|
|
|
|
|
# =========================
|
|
# CONFIGURAÇÃO DE CLASSES
|
|
# =========================
|
|
# Ajuste aqui conforme suas máscaras:
|
|
# - Se sua máscara for "indexada" (modo P) ou grayscale com IDs por pixel:
|
|
# class_ids = {0: 1, 1: 2} # exemplo: erva=1, cana=2
|
|
# - Se sua máscara for RGB com cores fixas:
|
|
# class_colors = {0: (0,255,0), 1: (0,0,255)} # exemplo
|
|
#
|
|
# Por padrão abaixo: tenta RGB primeiro; se a máscara vier indexada, usa IDs.
|
|
def build_maps_from_labelmap(alpha: int = 90, ignore_names=None):
|
|
"""
|
|
Lê labelmap e constrói maps dinâmicos:
|
|
- class_names: {new_id: name}
|
|
- class_colors_rgb: {new_id: (r,g,b)}
|
|
- class_ids: {new_id: new_id} (para máscaras indexed alinhadas com o new_id)
|
|
- overlay_rgba: {new_id: (r,g,b,alpha)}
|
|
- ignore_rgb: cor da classe "ignore" (se existir no labelmap original)
|
|
- id_old_to_new: {old_id: new_id} (útil se sua máscara indexed usa ids antigos)
|
|
- id_new_to_old: {new_id: old_id}
|
|
"""
|
|
if ignore_names is None:
|
|
ignore_names = []
|
|
|
|
ignore_set = {n.strip().lower() for n in ignore_names if n and n.strip()}
|
|
|
|
with open("config.json", "r", encoding="utf-8") as f:
|
|
config = json.load(f)
|
|
|
|
MODELO = config["camera"]
|
|
labelmap_path = os.path.join(MODELO, "dataset", "labelmap.txt")
|
|
|
|
# Lê tudo do labelmap (mantém seus ids originais)
|
|
cor_para_id, cores_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
|
|
|
# --- Filtra classes por nome ---
|
|
kept_old_ids = []
|
|
for old_id, name in id_para_nome.items():
|
|
if name.strip().lower() in ignore_set:
|
|
continue
|
|
kept_old_ids.append(old_id)
|
|
|
|
# Reindexa para ficar 0..N-1
|
|
kept_old_ids = sorted(kept_old_ids)
|
|
id_old_to_new = {old_id: new_id for new_id, old_id in enumerate(kept_old_ids)}
|
|
id_new_to_old = {new_id: old_id for old_id, new_id in id_old_to_new.items()}
|
|
|
|
# Constrói maps novos (compactos)
|
|
class_names = {}
|
|
class_colors_rgb = {}
|
|
overlay_rgba = {}
|
|
|
|
# cor_para_id: { (r,g,b): old_id }
|
|
# id_para_nome: { old_id: name }
|
|
for cor_rgb, old_id in cor_para_id.items():
|
|
if old_id not in id_old_to_new:
|
|
continue
|
|
new_id = id_old_to_new[old_id]
|
|
class_names[new_id] = id_para_nome[old_id]
|
|
class_colors_rgb[new_id] = cor_rgb
|
|
overlay_rgba[new_id] = (cor_rgb[0], cor_rgb[1], cor_rgb[2], alpha)
|
|
|
|
# Para máscara indexed:
|
|
# - Se sua máscara indexed já usa os IDs NOVOS (compactos), isso aqui está ok.
|
|
# - Se ela usa IDs ANTIGOS, você precisa mapear (old -> new) antes de extrair polígonos.
|
|
class_ids = {new_id: new_id for new_id in class_names.keys()}
|
|
|
|
return {
|
|
"class_names": class_names,
|
|
"class_colors_rgb": class_colors_rgb,
|
|
"class_ids": class_ids,
|
|
"overlay_rgba": overlay_rgba,
|
|
"ignore_rgb": ignore_rgb,
|
|
"id_old_to_new": id_old_to_new,
|
|
"id_new_to_old": id_new_to_old,
|
|
"labelmap_path": labelmap_path,
|
|
}
|
|
|
|
maps = build_maps_from_labelmap(ignore_names=["chao", "ignore"])
|
|
|
|
CLASS_NAMES = maps["class_names"]
|
|
CLASS_IDS = maps["class_ids"]
|
|
CLASS_COLORS_RGB = maps["class_colors_rgb"]
|
|
OVERLAY_RGBA = maps["overlay_rgba"]
|
|
IGNORE_RGB = maps["ignore_rgb"]
|
|
|
|
|
|
# =========================
|
|
# UTILITÁRIOS
|
|
# =========================
|
|
def imread_unicode(path: Path) -> np.ndarray:
|
|
"""Lê imagem com caminho unicode no Windows."""
|
|
data = np.fromfile(str(path), dtype=np.uint8)
|
|
img = cv2.imdecode(data, cv2.IMREAD_UNCHANGED)
|
|
return img
|
|
|
|
|
|
def load_mask(mask_path: Path):
|
|
"""
|
|
Retorna:
|
|
mask_type: 'indexed' ou 'rgb'
|
|
mask_data:
|
|
- indexed: np.ndarray (H,W) int
|
|
- rgb: np.ndarray (H,W,3) uint8 em RGB
|
|
"""
|
|
pil = Image.open(mask_path)
|
|
if pil.mode == "P":
|
|
arr = np.array(pil, dtype=np.int32)
|
|
return "indexed", arr
|
|
if pil.mode in ("L", "I;16"):
|
|
arr = np.array(pil, dtype=np.int32)
|
|
return "indexed", arr
|
|
|
|
# RGB/RGBA
|
|
pil = pil.convert("RGBA")
|
|
rgba = np.array(pil, dtype=np.uint8)
|
|
rgb = rgba[:, :, :3]
|
|
return "rgb", rgb
|
|
|
|
|
|
def class_binary_from_mask(mask_type, mask_data, cls, class_ids, class_colors, rgb_tol=10):
|
|
"""Gera máscara binária (uint8 0/255) para uma classe."""
|
|
if mask_type == "indexed":
|
|
target_id = class_ids[cls]
|
|
bin_mask = (mask_data == target_id).astype(np.uint8) * 255
|
|
return bin_mask
|
|
|
|
# rgb
|
|
target = np.array(class_colors[cls], dtype=np.int16)
|
|
img = mask_data.astype(np.int16)
|
|
diff = np.abs(img - target[None, None, :])
|
|
ok = (diff[:, :, 0] <= rgb_tol) & (diff[:, :, 1] <= rgb_tol) & (diff[:, :, 2] <= rgb_tol)
|
|
return ok.astype(np.uint8) * 255
|
|
|
|
|
|
def simplify_contour(cnt, epsilon_px=1.0, epsilon_rel=0.001):
|
|
peri = cv2.arcLength(cnt, True)
|
|
eps = max(epsilon_px, epsilon_rel * peri)
|
|
return cv2.approxPolyDP(cnt, eps, True)
|
|
|
|
|
|
def contours_to_polygons(bin_mask, min_area_px=50, epsilon_px=2.0):
|
|
"""
|
|
bin_mask: uint8 0/255
|
|
Retorna lista de polígonos, cada um como array (N,2) em pixels (float).
|
|
"""
|
|
# limpa ruído e fecha pequenos buracos
|
|
kernel = np.ones((3, 3), np.uint8)
|
|
m = cv2.morphologyEx(bin_mask, cv2.MORPH_OPEN, kernel, iterations=1)
|
|
m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, kernel, iterations=1)
|
|
|
|
contours, _hier = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
|
polys = []
|
|
for cnt in contours:
|
|
area = cv2.contourArea(cnt)
|
|
if area < min_area_px:
|
|
continue
|
|
approx = simplify_contour(cnt, epsilon_px=epsilon_px)
|
|
if len(approx) < 3:
|
|
continue
|
|
pts = approx.reshape(-1, 2).astype(np.float32)
|
|
polys.append(pts)
|
|
return polys
|
|
|
|
|
|
def polygon_px_to_yolo(poly_px, w, h):
|
|
"""(N,2) px -> lista [x1,y1,x2,y2,...] normalizada 0..1"""
|
|
xs = np.clip(poly_px[:, 0] / float(w), 0.0, 1.0)
|
|
ys = np.clip(poly_px[:, 1] / float(h), 0.0, 1.0)
|
|
coords = []
|
|
for x, y in zip(xs, ys):
|
|
coords.append(float(x))
|
|
coords.append(float(y))
|
|
return coords
|
|
|
|
|
|
def draw_polygons_on_preview(preview_path: Path, polygons_by_class, out_path: Path):
|
|
"""Cria overlay (PIL) com polígonos extraídos por cima do preview."""
|
|
img = Image.open(preview_path).convert("RGB")
|
|
w, h = img.size
|
|
draw = ImageDraw.Draw(img, "RGBA")
|
|
|
|
for cls, polys in polygons_by_class.items():
|
|
color = OVERLAY_RGBA.get(cls, (255, 255, 255, 90))
|
|
outline = color[:3] + (255,)
|
|
for poly in polys:
|
|
pts = [(float(x), float(y)) for x, y in poly]
|
|
if len(pts) >= 3:
|
|
draw.polygon(pts, fill=color, outline=outline)
|
|
|
|
img.save(out_path)
|
|
|
|
|
|
def make_triview(preview_path: Path, mask_path: Path, overlay_path: Path, out_path: Path, title: str = ""):
|
|
"""Salva uma imagem com 3 colunas: preview | mask | overlay."""
|
|
prev = Image.open(preview_path).convert("RGB")
|
|
msk = Image.open(mask_path).convert("RGB")
|
|
ovl = Image.open(overlay_path).convert("RGB")
|
|
|
|
fig = plt.figure(figsize=(16, 6))
|
|
fig.suptitle(title, fontsize=12)
|
|
|
|
ax1 = fig.add_subplot(1, 3, 1)
|
|
ax1.imshow(prev)
|
|
ax1.set_title("Preview")
|
|
ax1.axis("off")
|
|
|
|
ax2 = fig.add_subplot(1, 3, 2)
|
|
ax2.imshow(msk)
|
|
ax2.set_title("Mask (manual)")
|
|
ax2.axis("off")
|
|
|
|
ax3 = fig.add_subplot(1, 3, 3)
|
|
ax3.imshow(ovl)
|
|
ax3.set_title("Overlay (polígonos extraídos)")
|
|
ax3.axis("off")
|
|
|
|
plt.tight_layout()
|
|
fig.savefig(out_path, dpi=140)
|
|
plt.close(fig)
|
|
|
|
|
|
def find_matching_preview(previews_dir: Path, stem: str):
|
|
"""Procura preview com mesmo stem em extensões comuns."""
|
|
for ext in [".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"]:
|
|
p = previews_dir / f"{stem}{ext}"
|
|
if p.exists():
|
|
return p
|
|
return None
|
|
|
|
|
|
# =========================
|
|
# PIPELINE PRINCIPAL
|
|
# =========================
|
|
def process_dataset(root_dir: Path, out_labels_dir: Path, out_vis_dir: Path,
|
|
class_ids, class_colors, rgb_tol=10,
|
|
min_area_px=50, epsilon_px=2.0):
|
|
previews_dir = root_dir / "previews"
|
|
masks_dir = root_dir / "masks"
|
|
|
|
if not previews_dir.exists() or not masks_dir.exists():
|
|
raise FileNotFoundError(f"Esperado encontrar previews/ e masks/ dentro de {root_dir}")
|
|
|
|
out_labels_dir.mkdir(parents=True, exist_ok=True)
|
|
out_vis_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
mask_files = sorted(list(masks_dir.glob("*.png")) + list(masks_dir.glob("*.jpg")) + list(masks_dir.glob("*.jpeg")))
|
|
if not mask_files:
|
|
print(f"[WARN] Nenhuma máscara encontrada em: {masks_dir}")
|
|
return
|
|
|
|
total = 0
|
|
for mask_path in mask_files:
|
|
stem = mask_path.stem
|
|
preview_path = find_matching_preview(previews_dir, stem)
|
|
if preview_path is None:
|
|
print(f"[WARN] Sem preview para máscara: {mask_path.name}")
|
|
continue
|
|
|
|
# tamanhos
|
|
prev_img = Image.open(preview_path)
|
|
w, h = prev_img.size
|
|
|
|
mask_type, mask_data = load_mask(mask_path)
|
|
|
|
polygons_by_class = {}
|
|
yolo_lines = []
|
|
|
|
for cls in sorted(CLASS_NAMES.keys()):
|
|
if cls not in class_ids or cls not in class_colors:
|
|
continue
|
|
|
|
bin_mask = class_binary_from_mask(mask_type, mask_data, cls, class_ids, class_colors, rgb_tol=rgb_tol)
|
|
polys = contours_to_polygons(bin_mask, min_area_px=min_area_px, epsilon_px=epsilon_px)
|
|
if not polys:
|
|
continue
|
|
|
|
polygons_by_class[cls] = polys
|
|
|
|
for poly_px in polys:
|
|
coords = polygon_px_to_yolo(poly_px, w, h)
|
|
# YOLOv8-seg exige pelo menos 3 pontos (6 nums)
|
|
if len(coords) >= 6:
|
|
line = str(cls) + " " + " ".join(f"{v:.6f}" for v in coords)
|
|
yolo_lines.append(line)
|
|
|
|
# salva label
|
|
label_path = out_labels_dir / f"{stem}.txt"
|
|
label_path.write_text("\n".join(yolo_lines) + ("\n" if yolo_lines else ""), encoding="utf-8")
|
|
|
|
# gera overlay e triview
|
|
overlay_path = out_vis_dir / f"{stem}_overlay.png"
|
|
triview_path = out_vis_dir / f"{stem}_triview.png"
|
|
draw_polygons_on_preview(preview_path, polygons_by_class, overlay_path)
|
|
make_triview(preview_path, mask_path, overlay_path, triview_path, title=stem)
|
|
|
|
total += 1
|
|
print(f"[OK] {stem}: polys={sum(len(v) for v in polygons_by_class.values())} -> {label_path.name}")
|
|
|
|
print(f"\nFeito ✅ Processados: {total} arquivos")
|
|
print(f"Labels: {out_labels_dir}")
|
|
print(f"Vis: {out_vis_dir}")
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--root", type=str, required=True, help="Pasta raiz no formato antigo (contendo previews/ masks/ raws/ metas/)")
|
|
|
|
ap.add_argument("--rgb_tol", type=int, default=10, help="Tolerância p/ match de cor RGB na máscara")
|
|
ap.add_argument("--min_area", type=int, default=50, help="Área mínima (px) pra descartar sujeira")
|
|
ap.add_argument("--eps", type=float, default=2.0, help="Epsilon (px) pra simplificar polígonos")
|
|
|
|
args = ap.parse_args()
|
|
|
|
root_dir = Path(args.root)
|
|
out_labels_dir = Path(f"{args.root}/labels")
|
|
out_vis_dir = Path(f"{args.root}/vis")
|
|
|
|
process_dataset(
|
|
root_dir=root_dir,
|
|
out_labels_dir=out_labels_dir,
|
|
out_vis_dir=out_vis_dir,
|
|
class_ids=CLASS_IDS,
|
|
class_colors=CLASS_COLORS_RGB,
|
|
rgb_tol=args.rgb_tol,
|
|
min_area_px=args.min_area,
|
|
epsilon_px=args.eps
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|