1203 lines
35 KiB
Python
1203 lines
35 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
_7_split_multi_source.py
|
|
|
|
Split estratificado por grupo/família para OAK-FCC-3, com suporte a múltiplas
|
|
raízes de dataset e proteção contra vazamento de dados sintéticos/augmentados.
|
|
|
|
Por que este script existe?
|
|
---------------------------
|
|
No fluxo atual temos, por exemplo:
|
|
|
|
dataset/1024x640/group -> dados reais normalizados
|
|
dataset/copypaste/group -> dados sintéticos copy/paste pós-normalização
|
|
|
|
A validação precisa continuar 100% real. Então este script:
|
|
- usa dados reais para decidir o split por família
|
|
- manda aug/copy-paste/sintéticos somente para TRAIN
|
|
- opcionalmente só inclui sintéticos cuja família real caiu no TRAIN
|
|
- copia tensores, máscaras, masks auxiliares, metas, previews e visuals
|
|
|
|
Estrutura esperada:
|
|
<root>/<grupo>/tensors/*.npy
|
|
<root>/<grupo>/masks/*.npy
|
|
<root>/<grupo>/masks_vegetation/*.npy
|
|
<root>/<grupo>/masks_cana/*.npy
|
|
<root>/<grupo>/metas/*.json
|
|
<root>/<grupo>/previews/*.png
|
|
<root>/<grupo>/visuals/*.png
|
|
|
|
Uso recomendado para real + copy/paste:
|
|
---------------------------------------
|
|
python _7_split_multi_source.py ^
|
|
--src-roots dataset/1024x640/group,dataset/copypaste/group ^
|
|
--train-only-roots dataset/copypaste/group ^
|
|
--train 0.7 --val 0.3 --test 0.0 ^
|
|
--synthetic-train-only ^
|
|
--synthetic-respect-family-split ^
|
|
--clear-dst
|
|
|
|
Com isso:
|
|
- dados reais vão para train/val conforme split
|
|
- copy/paste vai apenas para train
|
|
- copy/paste derivado de família que caiu em val/test é ignorado por padrão
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import csv
|
|
import json
|
|
import shutil
|
|
import random
|
|
import argparse
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple, Any
|
|
|
|
|
|
# ============================================================
|
|
# Config
|
|
# ============================================================
|
|
|
|
try:
|
|
with open("config.json", "r", encoding="utf-8") as f:
|
|
CONFIG = json.load(f)
|
|
except Exception:
|
|
CONFIG = {}
|
|
|
|
RESOLUCAO = tuple(CONFIG.get("resolucao", [1024, 640]))
|
|
MULTI_HEAD = bool(CONFIG.get("multi_head", False))
|
|
|
|
TENSOR_EXT = ".npy"
|
|
MASK_NPY_SUFFIX = ".npy"
|
|
|
|
AUX_MASK_DIRS = [
|
|
"masks_vegetation",
|
|
"masks_cana",
|
|
]
|
|
|
|
COPY_OPTIONAL_DIRS = [
|
|
"metas",
|
|
"previews",
|
|
"visuals",
|
|
]
|
|
|
|
IGNORE_FILENAMES = {
|
|
"normalize_manifest.csv",
|
|
"copy_paste_manifest.csv",
|
|
"copy_paste_rejected.csv",
|
|
"copy_paste_summary.json",
|
|
"audit_summary.json",
|
|
"audit_health.json",
|
|
}
|
|
|
|
|
|
# ============================================================
|
|
# Regex de família/origem
|
|
# ============================================================
|
|
|
|
RE_ORIGINAL_PREFIX = re.compile(r"^original_(.+)$", re.IGNORECASE)
|
|
|
|
# augmented_abc_augx_00 -> family abc
|
|
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)
|
|
|
|
# copy-paste: abc_cp_00 -> family abc
|
|
RE_COPYPASTE_SUFFIX = re.compile(r"(.+)_cp_\d+$", re.IGNORECASE)
|
|
|
|
# grupos copy/paste: chao_cana_copypaste -> split_group chao_cana
|
|
RE_GROUP_COPYPASTE_SUFFIX = re.compile(r"(.+)_copypaste$", re.IGNORECASE)
|
|
|
|
|
|
# ============================================================
|
|
# Utilidades
|
|
# ============================================================
|
|
|
|
def garantir(p: str | Path) -> None:
|
|
os.makedirs(str(p), exist_ok=True)
|
|
|
|
|
|
def limpar_dir(p: str | Path) -> None:
|
|
if os.path.isdir(str(p)):
|
|
shutil.rmtree(str(p))
|
|
garantir(p)
|
|
|
|
|
|
def norm_path(p: str | Path) -> str:
|
|
return os.path.normcase(os.path.abspath(str(p)))
|
|
|
|
|
|
def parse_csv_list(s: Optional[str]) -> List[str]:
|
|
if not s:
|
|
return []
|
|
return [x.strip() for x in str(s).split(",") if x.strip()]
|
|
|
|
|
|
def parse_map(s: str, value_type=int) -> Dict[str, Any]:
|
|
out: Dict[str, Any] = {}
|
|
if not s:
|
|
return out
|
|
for item in s.split(","):
|
|
item = item.strip()
|
|
if not item:
|
|
continue
|
|
if ":" not in item:
|
|
raise ValueError(f"Item de mapa inválido: {item}. Use grupo:valor")
|
|
k, v = item.split(":", 1)
|
|
out[k.strip()] = value_type(v.strip())
|
|
return out
|
|
|
|
|
|
def load_json_safe(path: str | Path) -> Dict[str, Any]:
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def save_json(path: str | Path, data: Dict[str, Any]) -> None:
|
|
garantir(os.path.dirname(str(path)))
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
def lista_grupos(root: str | Path) -> List[str]:
|
|
root = str(root)
|
|
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, "tensors"))
|
|
and os.path.isdir(os.path.join(gdir, "masks"))
|
|
):
|
|
out.append(g)
|
|
return out
|
|
|
|
|
|
def listar_tensors(tensor_dir: str | Path) -> List[str]:
|
|
if not os.path.isdir(str(tensor_dir)):
|
|
return []
|
|
return sorted([
|
|
f for f in os.listdir(str(tensor_dir))
|
|
if f.lower().endswith(TENSOR_EXT)
|
|
])
|
|
|
|
|
|
def mask_npy_from_tensor_name(tensor_name: str) -> str:
|
|
base, _ = os.path.splitext(tensor_name)
|
|
return base + MASK_NPY_SUFFIX
|
|
|
|
|
|
def base_no_ext(filename: str) -> str:
|
|
return os.path.splitext(filename)[0]
|
|
|
|
|
|
def normalize_group_for_split(group_name: str) -> str:
|
|
"""
|
|
chao_cana_copypaste -> chao_cana
|
|
outros ficam iguais.
|
|
"""
|
|
m = RE_GROUP_COPYPASTE_SUFFIX.match(group_name)
|
|
if m:
|
|
return m.group(1)
|
|
return group_name
|
|
|
|
|
|
def output_group_name(source_group: str, split_group: str, is_synthetic: bool, mode: str) -> str:
|
|
"""
|
|
mode:
|
|
source: mantém nome original do grupo da fonte
|
|
base: usa grupo normalizado do split
|
|
suffix: sintético vai para <base>_synthetic, real fica base
|
|
"""
|
|
if mode == "source":
|
|
return source_group
|
|
if mode == "base":
|
|
return split_group
|
|
if mode == "suffix":
|
|
if is_synthetic:
|
|
return f"{split_group}_synthetic"
|
|
return split_group
|
|
return source_group
|
|
|
|
|
|
def classify_source_and_family(filename_no_ext: str, meta: Optional[Dict[str, Any]] = None) -> Tuple[str, str, bool]:
|
|
"""
|
|
Retorna:
|
|
source: original | augmented | copypaste | synthetic | unknown
|
|
family: id da família para impedir vazamento
|
|
is_synthetic: True para aug/copy/synthetic
|
|
|
|
Regras:
|
|
original_abc -> original, abc
|
|
augmented_abc_aug_00 -> augmented, abc
|
|
abc_aug_00 -> augmented, abc
|
|
abc_cp_00 -> copypaste, abc
|
|
meta.synthetic == true -> synthetic/copypaste, tenta family pelo receiver/base
|
|
abc -> unknown, abc, tratado como original real
|
|
"""
|
|
meta = meta or {}
|
|
|
|
m = RE_ORIGINAL_PREFIX.match(filename_no_ext)
|
|
if m:
|
|
return "original", m.group(1), False
|
|
|
|
m = RE_AUGMENTED_FAMILY.match(filename_no_ext)
|
|
if m:
|
|
return "augmented", m.group(1), True
|
|
|
|
if RE_AUG_SUFFIX.search(filename_no_ext):
|
|
fam = RE_AUG_SUFFIX.sub("", filename_no_ext)
|
|
return "augmented", fam, True
|
|
|
|
m = RE_COPYPASTE_SUFFIX.match(filename_no_ext)
|
|
if m:
|
|
return "copypaste", m.group(1), True
|
|
|
|
# Meta do copy/paste v1/v2.
|
|
cp = meta.get("copy_paste_augmentation")
|
|
if isinstance(cp, dict):
|
|
recv = cp.get("receiver")
|
|
if isinstance(recv, dict):
|
|
fam = str(recv.get("base") or filename_no_ext)
|
|
return "copypaste", fam, True
|
|
return "copypaste", filename_no_ext, True
|
|
|
|
if bool(meta.get("synthetic", False)):
|
|
return "synthetic", filename_no_ext, True
|
|
|
|
return "unknown", filename_no_ext, False
|
|
|
|
|
|
@dataclass
|
|
class Sample:
|
|
src_root: str
|
|
source_root_label: str
|
|
source_root_train_only: bool
|
|
|
|
source_group: str
|
|
split_group: str
|
|
output_group_default: str
|
|
|
|
tensor_path: str
|
|
tensor_name: str
|
|
base: str
|
|
|
|
mask_path: str
|
|
source: str
|
|
family: str
|
|
is_synthetic: bool
|
|
|
|
meta_path: Optional[str]
|
|
preview_path: Optional[str]
|
|
visual_path: Optional[str]
|
|
|
|
|
|
# ============================================================
|
|
# Coleta de amostras
|
|
# ============================================================
|
|
|
|
def collect_samples_from_root(
|
|
src_root: str | Path,
|
|
train_only_roots_norm: set[str],
|
|
label: str = "",
|
|
groups_filter: Optional[set[str]] = None,
|
|
) -> List[Sample]:
|
|
src_root = str(src_root)
|
|
src_root_norm = norm_path(src_root)
|
|
root_train_only = src_root_norm in train_only_roots_norm
|
|
root_label = label or os.path.basename(os.path.normpath(src_root)) or src_root
|
|
|
|
samples: List[Sample] = []
|
|
|
|
for group_name in lista_grupos(src_root):
|
|
if groups_filter and group_name not in groups_filter and normalize_group_for_split(group_name) not in groups_filter:
|
|
continue
|
|
|
|
group_dir = os.path.join(src_root, group_name)
|
|
tensor_dir = os.path.join(group_dir, "tensors")
|
|
mask_dir = os.path.join(group_dir, "masks")
|
|
meta_dir = os.path.join(group_dir, "metas")
|
|
preview_dir = os.path.join(group_dir, "previews")
|
|
visual_dir = os.path.join(group_dir, "visuals")
|
|
|
|
split_group = normalize_group_for_split(group_name)
|
|
|
|
for tensor_name in listar_tensors(tensor_dir):
|
|
base = base_no_ext(tensor_name)
|
|
mask_name = mask_npy_from_tensor_name(tensor_name)
|
|
tensor_path = os.path.join(tensor_dir, tensor_name)
|
|
mask_path = os.path.join(mask_dir, mask_name)
|
|
|
|
if not os.path.exists(mask_path):
|
|
continue
|
|
|
|
meta_path = os.path.join(meta_dir, base + ".json")
|
|
if not os.path.exists(meta_path):
|
|
meta_path = None
|
|
|
|
meta = load_json_safe(meta_path) if meta_path else {}
|
|
source, family, is_synthetic = classify_source_and_family(base, meta)
|
|
|
|
preview_path = None
|
|
for ext in (".png", ".jpg", ".jpeg"):
|
|
cand = os.path.join(preview_dir, base + ext)
|
|
if os.path.exists(cand):
|
|
preview_path = cand
|
|
break
|
|
|
|
visual_path = None
|
|
for ext in (".png", ".jpg", ".jpeg"):
|
|
cand = os.path.join(visual_dir, base + "_debug" + ext)
|
|
if os.path.exists(cand):
|
|
visual_path = cand
|
|
break
|
|
cand2 = os.path.join(visual_dir, base + ext)
|
|
if os.path.exists(cand2):
|
|
visual_path = cand2
|
|
break
|
|
|
|
samples.append(Sample(
|
|
src_root=src_root,
|
|
source_root_label=root_label,
|
|
source_root_train_only=root_train_only,
|
|
|
|
source_group=group_name,
|
|
split_group=split_group,
|
|
output_group_default=group_name,
|
|
|
|
tensor_path=tensor_path,
|
|
tensor_name=tensor_name,
|
|
base=base,
|
|
|
|
mask_path=mask_path,
|
|
source=source,
|
|
family=family,
|
|
is_synthetic=is_synthetic,
|
|
|
|
meta_path=meta_path,
|
|
preview_path=preview_path,
|
|
visual_path=visual_path,
|
|
))
|
|
|
|
return samples
|
|
|
|
|
|
def collect_all_samples(
|
|
src_roots: List[str],
|
|
train_only_roots: List[str],
|
|
groups: Optional[List[str]] = None,
|
|
) -> List[Sample]:
|
|
train_only_norm = {norm_path(p) for p in train_only_roots}
|
|
groups_filter = set(groups) if groups else None
|
|
|
|
all_samples: List[Sample] = []
|
|
for i, root in enumerate(src_roots):
|
|
label = f"root{i}"
|
|
samples = collect_samples_from_root(
|
|
src_root=root,
|
|
train_only_roots_norm=train_only_norm,
|
|
label=label,
|
|
groups_filter=groups_filter,
|
|
)
|
|
all_samples.extend(samples)
|
|
|
|
return all_samples
|
|
|
|
|
|
# ============================================================
|
|
# Split por família real
|
|
# ============================================================
|
|
|
|
def allocate_counts(n: int, p_train: float, p_val: float, p_test: float, min_train: int, min_val: int, min_test: int) -> Tuple[int, int, int]:
|
|
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)
|
|
|
|
diff = n - (n_train + n_val + n_test)
|
|
if diff != 0:
|
|
if diff > 0:
|
|
take = min(diff, n - n_train)
|
|
n_train += take
|
|
diff -= take
|
|
if diff > 0:
|
|
n_val += diff
|
|
else:
|
|
diff = -diff
|
|
take = min(diff, n_test)
|
|
n_test -= take
|
|
diff -= take
|
|
if diff > 0:
|
|
n_val -= diff
|
|
|
|
return n_train, n_val, n_test
|
|
|
|
|
|
def build_family_split(
|
|
samples: List[Sample],
|
|
p_train: float,
|
|
p_val: float,
|
|
p_test: float,
|
|
seed: int,
|
|
mins: Dict[str, int],
|
|
caps_map: Optional[Dict[str, int]] = None,
|
|
) -> Tuple[Dict[str, Dict[str, str]], Dict[str, Dict[str, int]]]:
|
|
"""
|
|
Usa apenas amostras reais/originais para decidir a partição das famílias.
|
|
|
|
Retorna:
|
|
family_split[split_group][family] = train|val|test
|
|
group_summary[split_group] = contagens de famílias
|
|
"""
|
|
caps_map = caps_map or {}
|
|
|
|
families_by_group: Dict[str, set[str]] = {}
|
|
|
|
for s in samples:
|
|
if s.source_root_train_only:
|
|
continue
|
|
if s.is_synthetic:
|
|
continue
|
|
if s.source in ("augmented", "copypaste", "synthetic"):
|
|
continue
|
|
|
|
families_by_group.setdefault(s.split_group, set()).add(s.family)
|
|
|
|
family_split: Dict[str, Dict[str, str]] = {}
|
|
group_summary: Dict[str, Dict[str, int]] = {}
|
|
|
|
for group_name in sorted(families_by_group.keys()):
|
|
fams = sorted(families_by_group[group_name])
|
|
rng = random.Random(seed)
|
|
rng.shuffle(fams)
|
|
|
|
total_familias = len(fams)
|
|
if total_familias == 0:
|
|
group_summary[group_name] = {"familias": 0, "train_families": 0, "val_families": 0, "test_families": 0}
|
|
continue
|
|
|
|
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(fams[:n_tr])
|
|
fam_val = set(fams[n_tr:n_tr + n_va])
|
|
fam_test = set(fams[n_tr + n_va:n_tr + n_va + n_te])
|
|
|
|
if group_name in caps_map:
|
|
cap = int(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} → "
|
|
f"mantidas {len(kept)} famílias, descartadas {len(dropped)} do TRAIN"
|
|
)
|
|
|
|
group_map: Dict[str, str] = {}
|
|
for f in fam_train:
|
|
group_map[f] = "train"
|
|
for f in fam_val:
|
|
group_map[f] = "val"
|
|
for f in fam_test:
|
|
group_map[f] = "test"
|
|
|
|
family_split[group_name] = group_map
|
|
group_summary[group_name] = {
|
|
"familias": total_familias,
|
|
"train_families": len(fam_train),
|
|
"val_families": len(fam_val),
|
|
"test_families": len(fam_test),
|
|
}
|
|
|
|
return family_split, group_summary
|
|
|
|
|
|
# ============================================================
|
|
# Cópia
|
|
# ============================================================
|
|
|
|
def copy_optional_file(src_path: Optional[str], dst_dir: str, dst_base: str, suffix: str = "") -> Optional[str]:
|
|
if not src_path or not os.path.exists(src_path):
|
|
return None
|
|
|
|
garantir(dst_dir)
|
|
ext = os.path.splitext(src_path)[1]
|
|
dst = os.path.join(dst_dir, dst_base + suffix + ext)
|
|
shutil.copy2(src_path, dst)
|
|
return dst
|
|
|
|
|
|
def copy_named_optional(src_dir: str, dst_dir: str, base: str, ext: str) -> Optional[str]:
|
|
src = os.path.join(src_dir, base + ext)
|
|
if not os.path.exists(src):
|
|
return None
|
|
garantir(dst_dir)
|
|
dst = os.path.join(dst_dir, base + ext)
|
|
shutil.copy2(src, dst)
|
|
return dst
|
|
|
|
|
|
def copiar_sample(
|
|
s: Sample,
|
|
dst_root: str,
|
|
split_name: str,
|
|
out_group_name: str,
|
|
copy_meta_preview: bool = True,
|
|
copy_visuals: bool = True,
|
|
) -> Optional[Dict[str, Any]]:
|
|
dst_group_dir = os.path.join(dst_root, split_name, "group", out_group_name)
|
|
|
|
dst_tensor_dir = os.path.join(dst_group_dir, "tensors")
|
|
dst_mask_dir = os.path.join(dst_group_dir, "masks")
|
|
dst_meta_dir = os.path.join(dst_group_dir, "metas")
|
|
dst_preview_dir = os.path.join(dst_group_dir, "previews")
|
|
dst_visual_dir = os.path.join(dst_group_dir, "visuals")
|
|
|
|
garantir(dst_tensor_dir)
|
|
garantir(dst_mask_dir)
|
|
|
|
tensor_dst = os.path.join(dst_tensor_dir, s.tensor_name)
|
|
mask_dst = os.path.join(dst_mask_dir, os.path.basename(s.mask_path))
|
|
|
|
if not (os.path.exists(s.tensor_path) and os.path.exists(s.mask_path)):
|
|
return None
|
|
|
|
shutil.copy2(s.tensor_path, tensor_dst)
|
|
shutil.copy2(s.mask_path, mask_dst)
|
|
|
|
base = s.base
|
|
|
|
# mask debug PNG opcional
|
|
mask_png_src = os.path.join(os.path.dirname(s.mask_path), base + ".png")
|
|
mask_png_dst = None
|
|
if os.path.exists(mask_png_src):
|
|
mask_png_dst = os.path.join(dst_mask_dir, base + ".png")
|
|
shutil.copy2(mask_png_src, mask_png_dst)
|
|
|
|
# aux masks
|
|
aux_masks: Dict[str, Dict[str, Optional[str]]] = {}
|
|
src_group_dir = os.path.join(s.src_root, s.source_group)
|
|
|
|
for aux_dir in AUX_MASK_DIRS:
|
|
aux_src_dir = os.path.join(src_group_dir, aux_dir)
|
|
aux_dst_dir = os.path.join(dst_group_dir, aux_dir)
|
|
|
|
aux_npy_src = os.path.join(aux_src_dir, base + ".npy")
|
|
aux_png_src = os.path.join(aux_src_dir, base + ".png")
|
|
|
|
aux_npy_dst = None
|
|
aux_png_dst = None
|
|
|
|
if os.path.exists(aux_npy_src):
|
|
garantir(aux_dst_dir)
|
|
aux_npy_dst = os.path.join(aux_dst_dir, base + ".npy")
|
|
shutil.copy2(aux_npy_src, aux_npy_dst)
|
|
|
|
if os.path.exists(aux_png_src):
|
|
aux_png_dst = os.path.join(aux_dst_dir, base + ".png")
|
|
shutil.copy2(aux_png_src, aux_png_dst)
|
|
|
|
aux_masks[aux_dir] = {
|
|
"npy": aux_npy_dst,
|
|
"png": aux_png_dst,
|
|
}
|
|
|
|
if MULTI_HEAD and aux_npy_dst is None:
|
|
raise RuntimeError(
|
|
f"multi_head=true, mas máscara auxiliar ausente: "
|
|
f"{aux_dir}/{base}.npy em {src_group_dir}"
|
|
)
|
|
|
|
meta_dst = None
|
|
preview_dst = None
|
|
visual_dst = None
|
|
|
|
if copy_meta_preview:
|
|
if s.meta_path and os.path.exists(s.meta_path):
|
|
garantir(dst_meta_dir)
|
|
meta_dst = os.path.join(dst_meta_dir, base + ".json")
|
|
shutil.copy2(s.meta_path, meta_dst)
|
|
|
|
if s.preview_path and os.path.exists(s.preview_path):
|
|
garantir(dst_preview_dir)
|
|
preview_dst = os.path.join(dst_preview_dir, base + os.path.splitext(s.preview_path)[1])
|
|
shutil.copy2(s.preview_path, preview_dst)
|
|
|
|
if copy_visuals and s.visual_path and os.path.exists(s.visual_path):
|
|
garantir(dst_visual_dir)
|
|
visual_dst = os.path.join(dst_visual_dir, os.path.basename(s.visual_path))
|
|
shutil.copy2(s.visual_path, visual_dst)
|
|
|
|
return {
|
|
"split": split_name,
|
|
"group": out_group_name,
|
|
"source_group": s.source_group,
|
|
"split_group": s.split_group,
|
|
"base": base,
|
|
"family": s.family,
|
|
"source": s.source,
|
|
"is_synthetic": int(bool(s.is_synthetic)),
|
|
"source_root": s.src_root,
|
|
"source_root_label": s.source_root_label,
|
|
"source_root_train_only": int(bool(s.source_root_train_only)),
|
|
|
|
"tensor": tensor_dst,
|
|
"mask_npy": mask_dst,
|
|
"mask_png": mask_png_dst,
|
|
|
|
"mask_vegetation_npy": aux_masks.get("masks_vegetation", {}).get("npy"),
|
|
"mask_vegetation_png": aux_masks.get("masks_vegetation", {}).get("png"),
|
|
|
|
"mask_cana_npy": aux_masks.get("masks_cana", {}).get("npy"),
|
|
"mask_cana_png": aux_masks.get("masks_cana", {}).get("png"),
|
|
|
|
"meta": meta_dst,
|
|
"preview": preview_dst,
|
|
"visual_debug": visual_dst,
|
|
}
|
|
|
|
|
|
def decide_sample_split(
|
|
s: Sample,
|
|
family_split: Dict[str, Dict[str, str]],
|
|
args,
|
|
) -> Tuple[Optional[str], str]:
|
|
"""
|
|
Retorna (split, reason)
|
|
split None = ignorado
|
|
"""
|
|
group_map = family_split.get(s.split_group, {})
|
|
assigned = group_map.get(s.family)
|
|
|
|
is_train_only_candidate = (
|
|
s.source_root_train_only
|
|
or (bool(args.synthetic_train_only) and s.is_synthetic)
|
|
or s.source in ("augmented", "copypaste", "synthetic")
|
|
)
|
|
|
|
if is_train_only_candidate:
|
|
if bool(args.synthetic_respect_family_split):
|
|
if assigned is None:
|
|
if bool(args.allow_orphan_synthetic_train):
|
|
return "train", "train_only_orphan_allowed"
|
|
return None, "skip_train_only_orphan_no_real_family"
|
|
|
|
if assigned != "train":
|
|
return None, f"skip_train_only_family_assigned_{assigned}"
|
|
|
|
return "train", "train_only_family_train"
|
|
|
|
return "train", "train_only_forced"
|
|
|
|
# real/original
|
|
if assigned is None:
|
|
return None, "skip_real_no_family_assignment"
|
|
|
|
return assigned, f"real_assigned_{assigned}"
|
|
|
|
|
|
# ============================================================
|
|
# Manifestos
|
|
# ============================================================
|
|
|
|
def write_manifest(path: str, rows: List[Dict[str, Any]]) -> None:
|
|
garantir(os.path.dirname(path))
|
|
|
|
fieldnames = [
|
|
"split",
|
|
"group",
|
|
"source_group",
|
|
"split_group",
|
|
"base",
|
|
"family",
|
|
"source",
|
|
"is_synthetic",
|
|
"source_root",
|
|
"source_root_label",
|
|
"source_root_train_only",
|
|
|
|
"tensor",
|
|
"mask_npy",
|
|
"mask_png",
|
|
|
|
"mask_vegetation_npy",
|
|
"mask_vegetation_png",
|
|
|
|
"mask_cana_npy",
|
|
"mask_cana_png",
|
|
|
|
"meta",
|
|
"preview",
|
|
"visual_debug",
|
|
]
|
|
|
|
with open(path, "w", newline="", encoding="utf-8") as f:
|
|
w = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
|
|
w.writeheader()
|
|
w.writerows(rows)
|
|
|
|
|
|
def write_skipped(path: str, rows: List[Dict[str, Any]]) -> None:
|
|
garantir(os.path.dirname(path))
|
|
fieldnames = [
|
|
"source_root",
|
|
"source_group",
|
|
"split_group",
|
|
"base",
|
|
"family",
|
|
"source",
|
|
"is_synthetic",
|
|
"reason",
|
|
]
|
|
with open(path, "w", newline="", encoding="utf-8") as f:
|
|
w = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
|
|
w.writeheader()
|
|
w.writerows(rows)
|
|
|
|
|
|
# ============================================================
|
|
# Main
|
|
# ============================================================
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser(
|
|
description="Split multi-source OAK-FCC-3 sem vazamento, com sintéticos/aug apenas no train."
|
|
)
|
|
|
|
ap.add_argument("--train", type=float, default=0.70)
|
|
ap.add_argument("--val", type=float, default=0.29)
|
|
ap.add_argument("--test", type=float, default=0.01)
|
|
ap.add_argument("--seed", type=int, default=42)
|
|
|
|
ap.add_argument("--min-train", type=int, default=1)
|
|
ap.add_argument("--min-val", type=int, default=1)
|
|
ap.add_argument("--min-test", type=int, default=0)
|
|
|
|
ap.add_argument(
|
|
"--resolucao",
|
|
type=str,
|
|
default=None,
|
|
help="Sobrescreve resolução no formato WxH. Ex: 1024x640.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--src-root",
|
|
type=str,
|
|
default=None,
|
|
help="Compatibilidade: uma raiz normalizada. Default: dataset/<resolucao>/group",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--src-roots",
|
|
type=str,
|
|
default="",
|
|
help="Lista de raízes separadas por vírgula. Ex: dataset/1024x640/group,dataset/copypaste/group",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--train-only-roots",
|
|
type=str,
|
|
default="",
|
|
help="Raízes que devem entrar somente no TRAIN. Ex: dataset/copypaste/group",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--dst-root",
|
|
type=str,
|
|
default="dataset/split",
|
|
help="Raiz do split.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--groups",
|
|
type=str,
|
|
default=None,
|
|
help="Lista de grupos/split_groups separados por vírgula.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--cap-train-families",
|
|
type=str,
|
|
default="",
|
|
help="Mapa 'grupo:cap,...' para limitar famílias reais no TRAIN. Ex: 'chao:50'",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--clear-dst",
|
|
action="store_true",
|
|
help="Apaga dst-root antes de copiar.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--no-meta-preview",
|
|
action="store_true",
|
|
help="Não copia metas/previews para o split.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--no-visuals",
|
|
action="store_true",
|
|
help="Não copia pasta visuals/debug.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--synthetic-train-only",
|
|
action="store_true",
|
|
default=True,
|
|
help="Força samples synthetic/augmented/copypaste para TRAIN apenas.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--allow-synthetic-val",
|
|
dest="synthetic_train_only",
|
|
action="store_false",
|
|
help="Permite sintéticos no val/test. Não recomendado.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--synthetic-respect-family-split",
|
|
action="store_true",
|
|
default=True,
|
|
help="Só inclui sintético no TRAIN se a família real correspondente caiu no TRAIN.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--no-synthetic-respect-family-split",
|
|
dest="synthetic_respect_family_split",
|
|
action="store_false",
|
|
help="Inclui sintéticos no TRAIN mesmo sem checar a família real. Mais arriscado.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--allow-orphan-synthetic-train",
|
|
action="store_true",
|
|
default=False,
|
|
help="Permite sintético no TRAIN mesmo sem família real encontrada.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--output-group-mode",
|
|
choices=["source", "base", "suffix"],
|
|
default="source",
|
|
help="Nome do grupo de saída. source mantém chao_cana_copypaste; base junta em chao_cana; suffix cria <base>_synthetic.",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--manifest",
|
|
type=str,
|
|
default="",
|
|
help="CSV de manifesto. Default: <dst-root>/split_manifest.csv",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--summary",
|
|
type=str,
|
|
default="",
|
|
help="JSON de resumo. Default: <dst-root>/split_summary.json",
|
|
)
|
|
|
|
ap.add_argument(
|
|
"--skipped",
|
|
type=str,
|
|
default="",
|
|
help="CSV de samples ignorados. Default: <dst-root>/split_skipped.csv",
|
|
)
|
|
|
|
args = ap.parse_args()
|
|
|
|
if args.resolucao:
|
|
try:
|
|
w, h = args.resolucao.lower().split("x")
|
|
resolucao = (int(w), int(h))
|
|
except Exception:
|
|
resolucao = RESOLUCAO
|
|
else:
|
|
resolucao = RESOLUCAO
|
|
|
|
default_src = os.path.join("dataset", f"{resolucao[0]}x{resolucao[1]}", "group")
|
|
|
|
src_roots = parse_csv_list(args.src_roots)
|
|
if args.src_root:
|
|
src_roots.insert(0, args.src_root)
|
|
if not src_roots:
|
|
src_roots = [default_src]
|
|
|
|
# Remove duplicatas preservando ordem.
|
|
seen = set()
|
|
src_roots_unique = []
|
|
for r in src_roots:
|
|
nr = norm_path(r)
|
|
if nr not in seen:
|
|
seen.add(nr)
|
|
src_roots_unique.append(r)
|
|
src_roots = src_roots_unique
|
|
|
|
train_only_roots = parse_csv_list(args.train_only_roots)
|
|
|
|
for r in src_roots:
|
|
if not os.path.isdir(r):
|
|
raise SystemExit(f"[ERRO] src-root não encontrado: {r}")
|
|
|
|
dst_root = args.dst_root
|
|
|
|
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),
|
|
}
|
|
|
|
groups = parse_csv_list(args.groups) if args.groups else None
|
|
caps_map = parse_map(args.cap_train_families, int)
|
|
|
|
if args.clear_dst:
|
|
print(f"[INFO] Limpando destino: {dst_root}")
|
|
limpar_dir(dst_root)
|
|
else:
|
|
garantir(dst_root)
|
|
|
|
samples = collect_all_samples(
|
|
src_roots=src_roots,
|
|
train_only_roots=train_only_roots,
|
|
groups=groups,
|
|
)
|
|
|
|
if not samples:
|
|
print("[WARN] Nenhuma amostra encontrada.")
|
|
return
|
|
|
|
family_split, family_summary = build_family_split(
|
|
samples=samples,
|
|
p_train=p_train,
|
|
p_val=p_val,
|
|
p_test=p_test,
|
|
seed=args.seed,
|
|
mins=mins,
|
|
caps_map=caps_map,
|
|
)
|
|
|
|
print("==========================================")
|
|
print("Split OAK-FCC-3 Multi-source")
|
|
print("SRC_ROOTS:")
|
|
for r in src_roots:
|
|
marker = " [TRAIN_ONLY]" if norm_path(r) in {norm_path(x) for x in train_only_roots} else ""
|
|
print(f" - {r}{marker}")
|
|
print(f"DST : {dst_root}")
|
|
print(f"Split : train={p_train:.3f}, val={p_val:.3f}, test={p_test:.3f}")
|
|
print(f"Mínimos : train={mins['train']} val={mins['val']} test={mins['test']}")
|
|
print(f"Seed : {args.seed}")
|
|
print(f"Synthetic : train_only={args.synthetic_train_only} respect_family={args.synthetic_respect_family_split}")
|
|
print("==========================================")
|
|
|
|
print("\nFamílias reais por split_group:")
|
|
for g, info in family_summary.items():
|
|
print(
|
|
f"[{g}] famílias={info['familias']} → "
|
|
f"train={info['train_families']}, val={info['val_families']}, test={info['test_families']}"
|
|
)
|
|
|
|
all_rows: List[Dict[str, Any]] = []
|
|
skipped_rows: List[Dict[str, Any]] = []
|
|
|
|
totals = {
|
|
"train": 0,
|
|
"val": 0,
|
|
"test": 0,
|
|
"skipped": 0,
|
|
"synthetic_train": 0,
|
|
"real_train": 0,
|
|
"real_val": 0,
|
|
"real_test": 0,
|
|
}
|
|
|
|
by_group: Dict[str, Dict[str, int]] = {}
|
|
|
|
for s in samples:
|
|
split_name, reason = decide_sample_split(s, family_split, args)
|
|
|
|
if split_name is None:
|
|
totals["skipped"] += 1
|
|
skipped_rows.append({
|
|
"source_root": s.src_root,
|
|
"source_group": s.source_group,
|
|
"split_group": s.split_group,
|
|
"base": s.base,
|
|
"family": s.family,
|
|
"source": s.source,
|
|
"is_synthetic": int(s.is_synthetic),
|
|
"reason": reason,
|
|
})
|
|
continue
|
|
|
|
out_group = output_group_name(
|
|
source_group=s.source_group,
|
|
split_group=s.split_group,
|
|
is_synthetic=s.is_synthetic,
|
|
mode=args.output_group_mode,
|
|
)
|
|
|
|
row = copiar_sample(
|
|
s=s,
|
|
dst_root=dst_root,
|
|
split_name=split_name,
|
|
out_group_name=out_group,
|
|
copy_meta_preview=not args.no_meta_preview,
|
|
copy_visuals=not args.no_visuals,
|
|
)
|
|
|
|
if row is None:
|
|
totals["skipped"] += 1
|
|
skipped_rows.append({
|
|
"source_root": s.src_root,
|
|
"source_group": s.source_group,
|
|
"split_group": s.split_group,
|
|
"base": s.base,
|
|
"family": s.family,
|
|
"source": s.source,
|
|
"is_synthetic": int(s.is_synthetic),
|
|
"reason": "copy_failed_missing_tensor_or_mask",
|
|
})
|
|
continue
|
|
|
|
all_rows.append(row)
|
|
|
|
totals[split_name] += 1
|
|
if split_name == "train" and s.is_synthetic:
|
|
totals["synthetic_train"] += 1
|
|
elif split_name == "train":
|
|
totals["real_train"] += 1
|
|
elif split_name == "val":
|
|
totals["real_val"] += 1
|
|
elif split_name == "test":
|
|
totals["real_test"] += 1
|
|
|
|
gsum = by_group.setdefault(out_group, {"train": 0, "val": 0, "test": 0, "synthetic_train": 0, "skipped": 0})
|
|
gsum[split_name] = gsum.get(split_name, 0) + 1
|
|
if split_name == "train" and s.is_synthetic:
|
|
gsum["synthetic_train"] += 1
|
|
|
|
manifest_path = args.manifest or os.path.join(dst_root, "split_manifest.csv")
|
|
summary_path = args.summary or os.path.join(dst_root, "split_summary.json")
|
|
skipped_path = args.skipped or os.path.join(dst_root, "split_skipped.csv")
|
|
|
|
write_manifest(manifest_path, all_rows)
|
|
write_skipped(skipped_path, skipped_rows)
|
|
|
|
summary = {
|
|
"schema": "oak_fcc3_multi_source_split_v1",
|
|
"src_roots": src_roots,
|
|
"train_only_roots": train_only_roots,
|
|
"dst_root": dst_root,
|
|
"resolution": list(resolucao),
|
|
"proportions": {
|
|
"train": p_train,
|
|
"val": p_val,
|
|
"test": p_test,
|
|
},
|
|
"mins": mins,
|
|
"seed": args.seed,
|
|
"multi_head": MULTI_HEAD,
|
|
"options": {
|
|
"synthetic_train_only": args.synthetic_train_only,
|
|
"synthetic_respect_family_split": args.synthetic_respect_family_split,
|
|
"allow_orphan_synthetic_train": args.allow_orphan_synthetic_train,
|
|
"output_group_mode": args.output_group_mode,
|
|
},
|
|
"family_summary": family_summary,
|
|
"groups": by_group,
|
|
"total": totals,
|
|
"samples_seen": len(samples),
|
|
"manifest": manifest_path,
|
|
"skipped": skipped_path,
|
|
}
|
|
|
|
save_json(summary_path, summary)
|
|
|
|
print("\nResumo global:")
|
|
print(f" train total: {totals['train']}")
|
|
print(f" real train: {totals['real_train']}")
|
|
print(f" synthetic train: {totals['synthetic_train']}")
|
|
print(f" val real: {totals['real_val']}")
|
|
print(f" test real: {totals['real_test']}")
|
|
print(f" skipped: {totals['skipped']}")
|
|
print(f"\nManifest: {manifest_path}")
|
|
print(f"Skipped : {skipped_path}")
|
|
print(f"Summary : {summary_path}")
|
|
|
|
print("\n✅ Split multi-source sem vazamento concluído!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|