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//previews/*.png|jpg brutas/group//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)