454 lines
16 KiB
Python
454 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Split estratificado por GRUPO com **val/test só do ORIGINAL** e
|
|
garantia de NÃO VAZAMENTO entre splits, agora para estrutura RAW:
|
|
|
|
Lê de:
|
|
MODELO/dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}
|
|
|
|
Escreve em:
|
|
MODELO/dataset/split/<split>/group/<grupo>/{previews,raws,(masks),(masks2)}
|
|
|
|
Definições:
|
|
- "Família" = todas as variações da MESMA base original:
|
|
original_<base>.* e augmented_<base>_aug_XX.*
|
|
- Val/Test: só **original_<base>** (sem augmented)
|
|
- Train: original_<base> **e** todos augmented_<base>_aug_XX
|
|
|
|
Baseado no _7_split.py original (versão images/masks). :contentReference[oaicite:1]{index=1}
|
|
"""
|
|
import os
|
|
import re
|
|
import json
|
|
import shutil
|
|
import random
|
|
import argparse
|
|
|
|
# ⚙️ Configurações
|
|
with open("config.json", "r", encoding="utf-8") as f:
|
|
config = json.load(f)
|
|
MODELO = config.get("camera")
|
|
USE_MASKS2 = config.get("dual_head", False)
|
|
RESOLUCAO = tuple(config.get("resolucao"))
|
|
|
|
# Pastas (ajustadas para PREVIEWS/RAWS)
|
|
pasta_origem = os.path.join(MODELO, "dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
|
|
pasta_destino = os.path.join(MODELO, "dataset", "split")
|
|
|
|
IMG_EXTS = (".jpg", ".jpeg", ".png")
|
|
MSK_EXT = ".png" # máscaras normalizadas em PNG
|
|
RAW_EXT = ".raw"
|
|
|
|
# Regex para identificar famílias
|
|
RE_ORIGINAL_PREFIX = re.compile(r'^original_(.+)$', re.IGNORECASE)
|
|
RE_AUGMENTED_FAMILY = re.compile(r'^augmented_(.+?)(?:_aug[a-zA-Z0-9]*_\d+)?$', re.IGNORECASE)
|
|
RE_AUG_SUFFIX = re.compile(r'_aug[a-zA-Z0-9]*_\d+$', re.IGNORECASE)
|
|
|
|
|
|
def garantir(p):
|
|
os.makedirs(p, exist_ok=True)
|
|
|
|
|
|
def lista_grupos(root):
|
|
"""Lista grupos que têm previews + masks (raws opcional)."""
|
|
if not os.path.isdir(root):
|
|
return []
|
|
out = []
|
|
for g in sorted(os.listdir(root)):
|
|
gdir = os.path.join(root, g)
|
|
if not os.path.isdir(gdir):
|
|
continue
|
|
if os.path.isdir(os.path.join(gdir, "previews")) and os.path.isdir(os.path.join(gdir, "masks")):
|
|
out.append(g)
|
|
return out
|
|
|
|
|
|
def listar_previews(prev_dir):
|
|
if not os.path.isdir(prev_dir):
|
|
return []
|
|
fs = []
|
|
for f in os.listdir(prev_dir):
|
|
ext = os.path.splitext(f.lower())[1]
|
|
if ext in IMG_EXTS:
|
|
fs.append(f)
|
|
return sorted(fs)
|
|
|
|
|
|
def mask_from_image_name(img_name):
|
|
base, _ = os.path.splitext(img_name)
|
|
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
|
|
|
|
|
|
def raw_from_image_name(img_name):
|
|
base, _ = os.path.splitext(img_name)
|
|
return base + RAW_EXT
|
|
|
|
|
|
def classify_source_and_family(filename_no_ext):
|
|
"""
|
|
Retorna (source, family_key)
|
|
source ∈ {"original", "augmented", "unknown"}
|
|
family_key = base associada ao original (sem prefixo/sufixos), ex: "foo_001"
|
|
"""
|
|
m = RE_ORIGINAL_PREFIX.match(filename_no_ext)
|
|
if m:
|
|
return "original", m.group(1)
|
|
|
|
m = RE_AUGMENTED_FAMILY.match(filename_no_ext)
|
|
if m:
|
|
return "augmented", m.group(1)
|
|
|
|
# legado: tenta deduzir se é augmented por sufixo, e família é o próprio nome sem sufixo
|
|
if RE_AUG_SUFFIX.search(filename_no_ext):
|
|
fam = RE_AUG_SUFFIX.sub("", filename_no_ext)
|
|
return "augmented", fam
|
|
|
|
return "unknown", filename_no_ext
|
|
|
|
|
|
def build_family_index(prev_dir, msk_dir):
|
|
"""
|
|
Constroi índice de famílias a partir de previews/masks.
|
|
Retorna: dict family -> {"original": str|None, "augmented": [str], "all": [str]}
|
|
(strings são NOMES DE ARQUIVO, não paths completos; assumem que a máscara existe)
|
|
"""
|
|
familias = {}
|
|
imgs = listar_previews(prev_dir)
|
|
for img_name in imgs:
|
|
base_no_ext, ext = os.path.splitext(img_name)
|
|
mask_name = mask_from_image_name(img_name)
|
|
if not os.path.exists(os.path.join(msk_dir, mask_name)):
|
|
continue # garante pareamento preview/mask
|
|
|
|
source, fam = classify_source_and_family(base_no_ext)
|
|
d = familias.setdefault(fam, {"original": None, "augmented": [], "all": []})
|
|
d["all"].append(img_name)
|
|
if source == "original":
|
|
d["original"] = img_name
|
|
elif source == "augmented":
|
|
d["augmented"].append(img_name)
|
|
else:
|
|
# trata como original desconhecido para não perder dado
|
|
if d["original"] is None:
|
|
d["original"] = img_name
|
|
else:
|
|
d["augmented"].append(img_name)
|
|
return familias
|
|
|
|
|
|
def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
|
|
"""
|
|
Mesmo esquema do script original: calcula quantas FAMÍLIAS vão para train/val/test.
|
|
"""
|
|
n_train = int(round(n * p_train))
|
|
n_val = int(round(n * p_val))
|
|
n_test = n - n_train - n_val
|
|
|
|
if n_test < 0:
|
|
excesso = -n_test
|
|
take_train = min(excesso, max(0, n_train))
|
|
n_train -= take_train
|
|
excesso -= take_train
|
|
if excesso > 0:
|
|
take_val = min(excesso, max(0, n_val))
|
|
n_val -= take_val
|
|
excesso -= take_val
|
|
n_test = 0
|
|
|
|
min_sum = min_train + min_val + min_test
|
|
if n >= min_sum:
|
|
n_train = max(n_train, min_train)
|
|
n_val = max(n_val, min_val)
|
|
n_test = max(n_test, min_test)
|
|
|
|
total = n_train + n_val + n_test
|
|
while total > n:
|
|
if n_test > min_test:
|
|
n_test -= 1
|
|
elif n_val > min_val:
|
|
n_val -= 1
|
|
elif n_train > min_train:
|
|
n_train -= 1
|
|
else:
|
|
break
|
|
total = n_train + n_val + n_test
|
|
while total < n:
|
|
if n_train - min_train <= n_val - min_val:
|
|
n_train += 1
|
|
else:
|
|
n_val += 1
|
|
total = n_train + n_val + n_test
|
|
else:
|
|
n_train = min(n, max(1, min_train))
|
|
resto = n - n_train
|
|
n_val = max(0, min(resto, min_val))
|
|
n_test = max(0, resto - n_val)
|
|
|
|
# ajuste final
|
|
diff = n - (n_train + n_val + n_test)
|
|
if diff != 0:
|
|
if diff > 0:
|
|
# adiciona em train, depois val
|
|
take = min(diff, n - n_train)
|
|
n_train += take
|
|
diff -= take
|
|
if diff > 0:
|
|
n_val += diff
|
|
else:
|
|
diff = -diff
|
|
# tira de test, depois val
|
|
take = min(diff, n_test)
|
|
n_test -= take
|
|
diff -= take
|
|
if diff > 0:
|
|
n_val -= diff
|
|
|
|
return n_train, n_val, n_test
|
|
|
|
|
|
def copiar(
|
|
nomes,
|
|
src_prev_dir,
|
|
src_msk_dir,
|
|
dst_prev_dir,
|
|
dst_msk_dir,
|
|
src_msk2_dir=None,
|
|
dst_msk2_dir=None,
|
|
src_raw_dir=None,
|
|
dst_raw_dir=None,
|
|
):
|
|
"""
|
|
Copia preview + mask (+ mask2) (+ raw) para o split.
|
|
"""
|
|
garantir(dst_prev_dir)
|
|
garantir(dst_msk_dir)
|
|
|
|
use_msk2 = bool(src_msk2_dir and dst_msk2_dir and os.path.isdir(src_msk2_dir))
|
|
use_raw = bool(src_raw_dir and dst_raw_dir and os.path.isdir(src_raw_dir))
|
|
|
|
if use_msk2:
|
|
garantir(dst_msk2_dir)
|
|
if use_raw:
|
|
garantir(dst_raw_dir)
|
|
|
|
moved = 0
|
|
for nome in nomes:
|
|
mask_name = mask_from_image_name(nome)
|
|
src_prev = os.path.join(src_prev_dir, nome)
|
|
src_msk = os.path.join(src_msk_dir, mask_name)
|
|
|
|
if not (os.path.exists(src_prev) and os.path.exists(src_msk)):
|
|
continue
|
|
|
|
shutil.copy2(src_prev, os.path.join(dst_prev_dir, nome))
|
|
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))
|
|
|
|
if use_raw:
|
|
raw_name = raw_from_image_name(nome)
|
|
src_raw = os.path.join(src_raw_dir, raw_name)
|
|
if os.path.exists(src_raw):
|
|
shutil.copy2(src_raw, os.path.join(dst_raw_dir, raw_name))
|
|
|
|
moved += 1
|
|
return moved
|
|
|
|
|
|
def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
|
|
src_prev_dir = os.path.join(pasta_origem, group_name, "previews")
|
|
src_msk_dir = os.path.join(pasta_origem, group_name, "masks")
|
|
src_msk2_dir = os.path.join(pasta_origem, group_name, "masks2")
|
|
src_raw_dir = os.path.join(pasta_origem, group_name, "raws")
|
|
|
|
use_msk2 = USE_MASKS2 and os.path.isdir(src_msk2_dir)
|
|
use_raw = os.path.isdir(src_raw_dir)
|
|
|
|
familias = build_family_index(src_prev_dir, src_msk_dir)
|
|
|
|
# apenas famílias que têm ORIGINAL para participar de val/test
|
|
familias_originais = [fam for fam, d in familias.items() if d["original"] is not None]
|
|
total_familias = len(familias_originais)
|
|
if total_familias == 0:
|
|
print(f"[{group_name}] 0 famílias com original, pulando.")
|
|
return {"train": 0, "val": 0, "test": 0, "familias": 0}
|
|
|
|
rng = random.Random(seed)
|
|
rng.shuffle(familias_originais)
|
|
|
|
n_tr, n_va, n_te = allocate_counts(
|
|
total_familias, p_train, p_val, p_test,
|
|
mins["train"], mins["val"], mins["test"]
|
|
)
|
|
|
|
fam_train = set(familias_originais[:n_tr])
|
|
fam_val = set(familias_originais[n_tr:n_tr+n_va])
|
|
fam_test = set(familias_originais[n_tr+n_va: n_tr+n_va+n_te])
|
|
|
|
# CAP por grupo (apenas no TRAIN)
|
|
if caps_map and group_name in caps_map:
|
|
cap = caps_map[group_name]
|
|
if len(fam_train) > cap:
|
|
fam_list = list(fam_train)
|
|
rng.shuffle(fam_list)
|
|
kept = set(fam_list[:cap])
|
|
dropped = set(fam_list[cap:])
|
|
fam_train = kept
|
|
print(f"[{group_name}] cap-train-families={cap} → mantidas {len(kept)} famílias, descartadas {len(dropped)} do TRAIN")
|
|
|
|
nomes_train, nomes_val, nomes_test = [], [], []
|
|
|
|
for fam, d in familias.items():
|
|
if fam in fam_train:
|
|
if d["original"]:
|
|
nomes_train.append(d["original"])
|
|
if d["augmented"]:
|
|
nomes_train.extend(d["augmented"])
|
|
elif fam in fam_val:
|
|
if d["original"]:
|
|
nomes_val.append(d["original"])
|
|
elif fam in fam_test:
|
|
if d["original"]:
|
|
nomes_test.append(d["original"])
|
|
else:
|
|
pass
|
|
|
|
# dest dirs
|
|
dest_train_prev = os.path.join(pasta_destino, "train", "group", group_name, "previews")
|
|
dest_train_msk = os.path.join(pasta_destino, "train", "group", group_name, "masks")
|
|
dest_val_prev = os.path.join(pasta_destino, "val", "group", group_name, "previews")
|
|
dest_val_msk = os.path.join(pasta_destino, "val", "group", group_name, "masks")
|
|
dest_test_prev = os.path.join(pasta_destino, "test", "group", group_name, "previews")
|
|
dest_test_msk = os.path.join(pasta_destino, "test", "group", group_name, "masks")
|
|
|
|
dest_train_msk2 = os.path.join(pasta_destino, "train", "group", group_name, "masks2") if use_msk2 else None
|
|
dest_val_msk2 = os.path.join(pasta_destino, "val", "group", group_name, "masks2") if use_msk2 else None
|
|
dest_test_msk2 = os.path.join(pasta_destino, "test", "group", group_name, "masks2") if use_msk2 else None
|
|
|
|
dest_train_raw = os.path.join(pasta_destino, "train", "group", group_name, "raws") if use_raw else None
|
|
dest_val_raw = os.path.join(pasta_destino, "val", "group", group_name, "raws") if use_raw else None
|
|
dest_test_raw = os.path.join(pasta_destino, "test", "group", group_name, "raws") if use_raw else None
|
|
|
|
m_train = copiar(
|
|
nomes_train,
|
|
src_prev_dir, src_msk_dir,
|
|
dest_train_prev, dest_train_msk,
|
|
src_msk2_dir, dest_train_msk2,
|
|
src_raw_dir, dest_train_raw
|
|
)
|
|
m_val = copiar(
|
|
nomes_val,
|
|
src_prev_dir, src_msk_dir,
|
|
dest_val_prev, dest_val_msk,
|
|
src_msk2_dir, dest_val_msk2,
|
|
src_raw_dir, dest_val_raw
|
|
)
|
|
m_test = copiar(
|
|
nomes_test,
|
|
src_prev_dir, src_msk_dir,
|
|
dest_test_prev, dest_test_msk,
|
|
src_msk2_dir, dest_test_msk2,
|
|
src_raw_dir, dest_test_raw
|
|
)
|
|
|
|
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}
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Split estratificado por grupo SEM vazamento (RAW: previews/raws/masks).")
|
|
ap.add_argument("--train", type=float, default=0.70, help="Proporção de treino (default=0.70).")
|
|
ap.add_argument("--val", type=float, default=0.29, help="Proporção de validação (default=0.29).")
|
|
ap.add_argument("--test", type=float, default=0.01, help="Proporção de teste (default=0.01).")
|
|
ap.add_argument("--seed", type=int, default=42, help="Seed do embaralhamento (default=42).")
|
|
|
|
ap.add_argument("--min-train", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em train (default=1).")
|
|
ap.add_argument("--min-val", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em val (default=1).")
|
|
ap.add_argument("--min-test", type=int, default=0, help="Mínimo de FAMÍLIAS por grupo em test (default=0).")
|
|
|
|
ap.add_argument("--modelo", type=str, default=None, help="Sobrescreve MODELO do config.json.")
|
|
ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH (ex: 960x544).")
|
|
|
|
ap.add_argument("--cap-train-families", type=str, default="",
|
|
help="Mapa 'grupo:cap,...' p/ limitar número de FAMÍLIAS no TRAIN. Ex.: 'chao:350'")
|
|
|
|
args = ap.parse_args()
|
|
|
|
modelo = args.modelo or MODELO
|
|
if args.resolucao:
|
|
try:
|
|
w, h = args.resolucao.lower().split("x")
|
|
resolucao = (int(w), int(h))
|
|
except Exception:
|
|
resolucao = RESOLUCAO
|
|
else:
|
|
resolucao = RESOLUCAO
|
|
|
|
def parse_cap_map(s):
|
|
caps = {}
|
|
if not s:
|
|
return caps
|
|
for item in s.split(","):
|
|
k, v = item.strip().split(":")
|
|
caps[k.strip()] = int(v)
|
|
return caps
|
|
|
|
caps_map = parse_cap_map(args.cap_train_families)
|
|
|
|
global pasta_origem, pasta_destino
|
|
pasta_origem = os.path.join(modelo, "dataset", f"{resolucao[0]}x{resolucao[1]}", "group")
|
|
pasta_destino = os.path.join(modelo, "dataset", "split")
|
|
|
|
soma = args.train + args.val + args.test
|
|
if soma <= 0:
|
|
raise ValueError("Soma de proporções deve ser > 0.")
|
|
p_train = args.train / soma
|
|
p_val = args.val / soma
|
|
p_test = args.test / soma
|
|
|
|
mins = {
|
|
"train": max(0, args.min_train),
|
|
"val": max(0, args.min_val),
|
|
"test": max(0, args.min_test),
|
|
}
|
|
|
|
garantir(pasta_destino)
|
|
|
|
grupos = lista_grupos(pasta_origem)
|
|
if not grupos:
|
|
print(f"[WARN] Nenhum grupo encontrado em: {pasta_origem}")
|
|
return
|
|
|
|
random.seed(args.seed)
|
|
|
|
total_global = {"train": 0, "val": 0, "test": 0, "familias": 0}
|
|
print(f"Grupos: {', '.join(grupos)}")
|
|
print(f"Proporções normalizadas: train={p_train:.3f}, val={p_val:.3f}, test={p_test:.3f}")
|
|
print(f"Mínimos por grupo (famílias): train={mins['train']} val={mins['val']} test={mins['test']}")
|
|
|
|
for g in grupos:
|
|
res = split_group(g, p_train, p_val, p_test, args.seed, mins, caps_map=caps_map)
|
|
for k in total_global.keys():
|
|
total_global[k] += res.get(k, 0)
|
|
|
|
print("\nResumo global (imagens copiadas):")
|
|
print(f" train: {total_global['train']}")
|
|
print(f" val: {total_global['val']}")
|
|
print(f" test: {total_global['test']}")
|
|
print(f" famílias (total): {total_global['familias']}")
|
|
print("\n✅ Split RAW sem vazamento concluído!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|