206 lines
7.0 KiB
Python
206 lines
7.0 KiB
Python
import json
|
|
import os
|
|
import shutil
|
|
import argparse
|
|
import csv
|
|
|
|
# ================= CONFIG =================
|
|
|
|
with open("config.json", "r") as f:
|
|
config = json.load(f)
|
|
|
|
MODELO = config["camera"]
|
|
|
|
# Raiz das brutas agrupadas
|
|
PASTA_BRUTAS_GROUP_ROOT = os.path.join(MODELO, "dataset", "brutas")
|
|
|
|
# Onde as máscaras novas (rotuladas externamente) são colocadas
|
|
PASTA_NEW_MASKS = os.path.join(MODELO, "dataset", "new_masks")
|
|
|
|
# Destino final do dataset consolidado
|
|
PASTA_FINAL_PREVIEWS = os.path.join(MODELO, "dataset", "original", "previews")
|
|
PASTA_FINAL_RAWS = os.path.join(MODELO, "dataset", "original", "raws")
|
|
PASTA_FINAL_MASKS = os.path.join(MODELO, "dataset", "original", "masks")
|
|
|
|
EXT_PREVIEWS = (".png", ".jpg", ".jpeg")
|
|
EXT_MASKS = (".png",)
|
|
EXT_RAWS = (".raw",)
|
|
|
|
# =========================================
|
|
|
|
def garantir_pasta(p: str):
|
|
os.makedirs(p, exist_ok=True)
|
|
|
|
def indexar_brutas_group(root_group: str):
|
|
"""
|
|
Varrre dataset/brutas/group recursivamente e monta dois dicionários:
|
|
base -> caminho_preview
|
|
base -> caminho_raw
|
|
|
|
Considera que a estrutura é:
|
|
brutas/group/<grupo>/previews/*.png|jpg
|
|
brutas/group/<grupo>/raws/*.raw
|
|
"""
|
|
previews_map = {}
|
|
raws_map = {}
|
|
|
|
if not os.path.isdir(root_group):
|
|
print(f"[AVISO] Pasta de grupos não existe: {root_group}")
|
|
return previews_map, raws_map
|
|
|
|
for dirpath, dirnames, filenames in os.walk(root_group):
|
|
base_dir = os.path.basename(dirpath).lower()
|
|
|
|
if base_dir == "previews":
|
|
for nome in filenames:
|
|
if not nome.lower().endswith(EXT_PREVIEWS):
|
|
continue
|
|
base = os.path.splitext(nome)[0]
|
|
# Se já existir, mantemos o primeiro encontrado (pode logar conflito se quiser)
|
|
if base not in previews_map:
|
|
previews_map[base] = os.path.join(dirpath, nome)
|
|
|
|
elif base_dir == "raws":
|
|
for nome in filenames:
|
|
if not nome.lower().endswith(EXT_RAWS):
|
|
continue
|
|
base = os.path.splitext(nome)[0]
|
|
if base not in raws_map:
|
|
raws_map[base] = os.path.join(dirpath, nome)
|
|
|
|
print(f"[INDEX] Previews indexados: {len(previews_map)}")
|
|
print(f"[INDEX] Raws indexados : {len(raws_map)}")
|
|
return previews_map, raws_map
|
|
|
|
def nome_disponivel(dest_dir: str, base: str, ext: str) -> str:
|
|
"""
|
|
Devolve um caminho disponível em dest_dir com base + ext.
|
|
Se já existir, adiciona _001, _002, ...
|
|
"""
|
|
p = os.path.join(dest_dir, base + ext)
|
|
if not os.path.exists(p):
|
|
return p
|
|
|
|
i = 1
|
|
while True:
|
|
p = os.path.join(dest_dir, f"{base}_{i:03d}{ext}")
|
|
if not os.path.exists(p):
|
|
return p
|
|
i += 1
|
|
|
|
def processar_new_masks(mover: bool = True, manifesto_csv: str | None = None):
|
|
garantir_pasta(PASTA_FINAL_PREVIEWS)
|
|
garantir_pasta(PASTA_FINAL_RAWS)
|
|
garantir_pasta(PASTA_FINAL_MASKS)
|
|
|
|
if not os.path.isdir(PASTA_NEW_MASKS):
|
|
raise SystemExit(f"[ERRO] Pasta new_masks não existe: {PASTA_NEW_MASKS}")
|
|
|
|
# 1) Indexa tudo que existe em brutas/group
|
|
previews_map, raws_map = indexar_brutas_group(PASTA_BRUTAS_GROUP_ROOT)
|
|
|
|
registros = []
|
|
total, copiados, ignorados = 0, 0, 0
|
|
|
|
for nome in os.listdir(PASTA_NEW_MASKS):
|
|
if not nome.lower().endswith(EXT_MASKS):
|
|
continue
|
|
|
|
total += 1
|
|
base = os.path.splitext(nome)[0]
|
|
caminho_mask_src = os.path.join(PASTA_NEW_MASKS, nome)
|
|
|
|
preview_src = previews_map.get(base)
|
|
raw_src = raws_map.get(base)
|
|
|
|
if preview_src is None or raw_src is None:
|
|
ignorados += 1
|
|
motivo = []
|
|
if preview_src is None:
|
|
motivo.append("preview")
|
|
if raw_src is None:
|
|
motivo.append("raw")
|
|
print(f"[SKIP] {base} -> faltando: {', '.join(motivo)}")
|
|
# NÃO move a máscara, ela fica em new_masks pra você analisar depois
|
|
continue
|
|
|
|
prev_ext = os.path.splitext(preview_src)[1].lower()
|
|
mask_ext = os.path.splitext(nome)[1].lower()
|
|
raw_ext = os.path.splitext(raw_src)[1].lower()
|
|
|
|
# Gera um base final único usando o preview como referência
|
|
dst_preview = nome_disponivel(PASTA_FINAL_PREVIEWS, base, prev_ext)
|
|
new_base = os.path.splitext(os.path.basename(dst_preview))[0]
|
|
|
|
dst_raw = os.path.join(PASTA_FINAL_RAWS, new_base + raw_ext)
|
|
dst_mask = os.path.join(PASTA_FINAL_MASKS, new_base + mask_ext)
|
|
|
|
# Copia preview + raw, move ou copia a máscara
|
|
shutil.copy2(preview_src, dst_preview)
|
|
shutil.copy2(raw_src, dst_raw)
|
|
if mover:
|
|
shutil.move(caminho_mask_src, dst_mask)
|
|
else:
|
|
shutil.copy2(caminho_mask_src, dst_mask)
|
|
|
|
copiados += 1
|
|
registros.append([
|
|
preview_src,
|
|
raw_src,
|
|
caminho_mask_src,
|
|
dst_preview,
|
|
dst_raw,
|
|
dst_mask,
|
|
])
|
|
|
|
print(f"[OK] {new_base}")
|
|
|
|
if manifesto_csv and registros:
|
|
with open(manifesto_csv, "w", newline="", encoding="utf-8") as f:
|
|
w = csv.writer(f)
|
|
w.writerow([
|
|
"src_preview",
|
|
"src_raw",
|
|
"src_mask",
|
|
"dst_preview",
|
|
"dst_raw",
|
|
"dst_mask",
|
|
])
|
|
w.writerows(registros)
|
|
print(f"[MANIFESTO] {manifesto_csv} salvo ({len(registros)} entradas).")
|
|
|
|
print(f"\nResumo: total_masks={total} | ingestas={copiados} | ignoradas={ignorados}")
|
|
|
|
# ================= CLI =================
|
|
|
|
def build_cli():
|
|
ap = argparse.ArgumentParser(
|
|
description=(
|
|
"Ingere máscaras novas (dataset/new_masks) procurando previews e raws "
|
|
"nas subpastas de dataset/brutas/group/** e copiando tudo para dataset/original."
|
|
)
|
|
)
|
|
ap.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
|
|
ap.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
|
|
ap.add_argument("--copy", action="store_true", help="Copia as máscaras em vez de mover (por padrão, move e esvazia new_masks conforme ingere).")
|
|
ap.add_argument("--manifest", default="", help="Caminho para CSV de manifesto (opcional).")
|
|
return ap
|
|
|
|
if __name__ == "__main__":
|
|
ap = build_cli()
|
|
args = ap.parse_args()
|
|
|
|
manifesto = args.manifest if args.manifest.strip() else None
|
|
mover = not args.copy
|
|
|
|
PASTA_BRUTAS_GROUP_ROOT = os.path.join(PASTA_BRUTAS_GROUP_ROOT, f"cana_{args.cana}", args.horario, "group")
|
|
|
|
print(f"[INFO] MODELO : {MODELO}")
|
|
print(f"[INFO] Brutas (group root): {PASTA_BRUTAS_GROUP_ROOT}")
|
|
print(f"[INFO] New masks : {PASTA_NEW_MASKS}")
|
|
print(f"[INFO] Mover máscaras? : {mover}")
|
|
print(f"[INFO] Manifesto : {manifesto or '(nenhum)'}")
|
|
print()
|
|
|
|
processar_new_masks(mover=mover, manifesto_csv=manifesto)
|