permitido navegar pelos arquivos do treinamento

Python/OAK/datasets/oak-fcc-3/utils/check_saved_files.py

@@ -663,21 +663,236 @@ def put_label(img: np.ndarray, title: str, subtitle: str = "") -> np.ndarray:
     return out
 
 
+def find_dataset_layout_root(path: Path) -> Path | None:
+    """
+    Detecta layout de dataset:
+      root/metas
+      root/previews
+      root/bins
+      root/masks opcional
+
+    Pode receber:
+      - root do dataset
+      - root/metas
+      - root/previews
+      - root/bins
+      - arquivo dentro dessas pastas
+    """
+    p = path.resolve()
+
+    candidates = []
+
+    if p.is_file():
+        candidates.append(p.parent)
+        candidates.append(p.parent.parent)
+    else:
+        candidates.append(p)
+        candidates.append(p.parent)
+
+    for c in candidates:
+        if not c:
+            continue
+
+        # Se o usuário apontou diretamente para metas/previews/bins
+        if c.name.lower() in ("metas", "metadata", "jsons", "previews", "bins", "masks"):
+            root = c.parent
+        else:
+            root = c
+
+        metas = root / "metas"
+        previews = root / "previews"
+        bins = root / "bins"
+
+        if metas.is_dir() and previews.is_dir() and bins.is_dir():
+            return root
+
+    return None
+
+
+def is_dataset_layout(path: Path) -> bool:
+    return find_dataset_layout_root(path) is not None
+
+
+def resolve_dataset_sibling_file(root: Path, subdir: str, stem: str, exts: tuple[str, ...]) -> Path | None:
+    folder = root / subdir
+    if not folder.is_dir():
+        return None
+
+    for ext in exts:
+        p = folder / f"{stem}{ext}"
+        if p.exists():
+            return p
+
+    return None
+
+
+def resolve_dataset_json_from_any_path(input_path: Path, dataset_root: Path) -> Path:
+    """
+    Dado um caminho qualquer dentro do dataset, descobre o JSON correspondente em metas/.
+    """
+    input_path = input_path.resolve()
+
+    if input_path.is_file():
+        stem = input_path.stem
+
+        # Se vier algo tipo xxx_CAM_A.bin ou xxx_cam0.bin, remove sufixo de câmera.
+        for token in ("_CAM_A", "_CAM_B", "_CAM_C", "_cam0", "_cam1", "_cam2", "_rgb", "_re", "_nir"):
+            if token in stem:
+                stem = stem.split(token)[0]
+                break
+
+        json_path = dataset_root / "metas" / f"{stem}.json"
+        if json_path.exists():
+            return json_path
+
+        if input_path.suffix.lower() == ".json" and input_path.parent.name.lower() == "metas":
+            return input_path
+
+    raise FileNotFoundError(f"Não consegui resolver JSON do dataset para: {input_path}")
+
+
 def resolve_capture_group(input_path: Path):
     """
     Resolve todos os arquivos relacionados a uma captura.
 
+    Suporta dois layouts:
+
+    1) Layout antigo, tudo na mesma pasta:
+       captura.json
+       captura.png
+       captura.raw/.bin
+       captura_CAM_A.bin ...
+
+    2) Layout dataset:
+       root/metas/captura.json
+       root/previews/captura.png
+       root/bins/captura_CAM_A.bin ...
+       root/masks/captura.png opcional
+
     Retorna:
     {
         "json": Path,
         "png": Path,
+        "mask": Path | None,
         "final_raw": Path | None,
-        "cameras": { "cam0": Path, ... }
+        "cameras": { "CAM_A": Path, ... },
+        "dataset_root": Path | None,
     }
     """
     input_path = input_path.resolve()
-    folder = input_path.parent
 
+    dataset_root = find_dataset_layout_root(input_path)
+
+    # =========================================================
+    # Caso novo: dataset/metas, previews, bins, masks
+    # =========================================================
+    if dataset_root is not None:
+        if input_path.is_dir():
+            # Se vier a raiz do dataset ou subpasta, pega o primeiro json depois pela navegação.
+            # Aqui só usamos fallback defensivo.
+            metas = sorted((dataset_root / "metas").glob("*.json"))
+            if not metas:
+                raise RuntimeError(f"Nenhum JSON encontrado em: {dataset_root / 'metas'}")
+            json_path = metas[0]
+        elif input_path.suffix.lower() == ".json":
+            json_path = input_path
+        else:
+            json_path = resolve_dataset_json_from_any_path(input_path, dataset_root)
+
+        meta = load_json(json_path)
+        base_name = json_path.stem
+
+        png_path = resolve_dataset_sibling_file(
+            dataset_root,
+            "previews",
+            base_name,
+            (".png", ".jpg", ".jpeg"),
+        )
+
+        if png_path is None:
+            raise FileNotFoundError(
+                f"Preview não encontrado para {base_name} em {dataset_root / 'previews'}"
+            )
+
+        mask_path = resolve_dataset_sibling_file(
+            dataset_root,
+            "masks",
+            base_name,
+            (".png", ".tif", ".tiff", ".npy"),
+        )
+
+        group = {
+            "json": json_path,
+            "png": png_path,
+            "mask": mask_path,
+            "final_raw": None,
+            "cameras": {},
+            "dataset_root": dataset_root,
+        }
+
+        bins_dir = dataset_root / "bins"
+
+        # Payload multi por câmera
+        if "saved_payload_paths" in meta:
+            for cam_id, fname in meta["saved_payload_paths"].items():
+                fname_path = Path(fname)
+
+                candidates = []
+
+                # 1) caminho exatamente como veio no JSON, relativo ao bins/
+                candidates.append(bins_dir / fname_path.name)
+
+                # 2) relativo à pasta do JSON
+                candidates.append(json_path.parent / fname)
+
+                # 3) relativo à raiz do dataset
+                candidates.append(dataset_root / fname)
+
+                # 4) fallback por padrões comuns
+                candidates.extend([
+                    bins_dir / f"{base_name}_{cam_id}.bin",
+                    bins_dir / f"{base_name}_{cam_id}.raw",
+                    bins_dir / f"{base_name}_{cam_id.lower()}.bin",
+                    bins_dir / f"{base_name}_{cam_id.lower()}.raw",
+                ])
+
+                found = None
+                for c in candidates:
+                    if c.exists():
+                        found = c
+                        break
+
+                if found is not None:
+                    group["cameras"][cam_id] = found
+                else:
+                    print(f"[WARN] bin não encontrado para {cam_id}: {fname}")
+
+        # Payload final único
+        elif "saved_payload_path" in meta:
+            fname = Path(meta["saved_payload_path"])
+
+            candidates = [
+                bins_dir / fname.name,
+                json_path.parent / fname,
+                dataset_root / fname,
+                bins_dir / f"{base_name}.raw",
+                bins_dir / f"{base_name}.bin",
+            ]
+
+            for c in candidates:
+                if c.exists():
+                    group["final_raw"] = c
+                    break
+
+            if group["final_raw"] is None:
+                print(f"[WARN] payload final não encontrado para {base_name}: {fname}")
+
+        return group
+
+    # =========================================================
+    # Caso antigo: tudo na mesma pasta
+    # =========================================================
+    folder = input_path.parent
     name = input_path.stem
 
     # remove sufixo _camX se existir
@@ -699,22 +914,20 @@ def resolve_capture_group(input_path: Path):
     group = {
         "json": json_path,
         "png": png_path,
+        "mask": None,
         "final_raw": None,
-        "cameras": {}
+        "cameras": {},
+        "dataset_root": None,
     }
 
-    # =========================
     # Caso MULTI payload
-    # =========================
     if "saved_payload_paths" in meta:
         for cam_id, fname in meta["saved_payload_paths"].items():
             path = folder / fname
             if path.exists():
                 group["cameras"][cam_id] = path
 
-    # =========================
     # Caso payload único (.raw)
-    # =========================
     elif "saved_payload_path" in meta:
         path = folder / meta["saved_payload_path"]
         if path.exists():
@@ -754,12 +967,29 @@ def resolve_sensor_dims_for_raw10_packed(arr: np.ndarray, cam_meta: dict, meta:
 
 def list_capture_groups_from_dir(folder: Path) -> list[Path]:
     """
-    Lista todos os JSONs de captura do diretório, ordenados por nome.
+    Lista todos os JSONs de captura.
-    Cada JSON representa uma captura.
+
+    Suporta:
+      - pasta antiga com *.json direto
+      - dataset root com metas/*.json
+      - dataset/metas
+      - dataset/previews
+      - dataset/bins
     """
+    folder = folder.resolve()
+
     if not folder.exists() or not folder.is_dir():
         raise FileNotFoundError(f"Diretório não encontrado: {folder}")
 
+    dataset_root = find_dataset_layout_root(folder)
+
+    if dataset_root is not None:
+        metas_dir = dataset_root / "metas"
+        items = sorted(metas_dir.glob("*.json"))
+        if not items:
+            raise RuntimeError(f"Nenhum arquivo .json encontrado em: {metas_dir}")
+        return items
+
     items = sorted(folder.glob("*.json"))
     if not items:
         raise RuntimeError(f"Nenhum arquivo .json encontrado em: {folder}")
@@ -772,22 +1002,53 @@ def resolve_navigation_inputs(input_path: Path) -> tuple[list[Path], int]:
     Retorna:
       entries: lista de JSONs de captura
       start_index: índice inicial baseado no input fornecido
+
+    Aceita:
+      - arquivo .json/.png/.bin/.raw do layout antigo
+      - pasta antiga
+      - root dataset
+      - dataset/metas
+      - dataset/previews
+      - dataset/bins
+      - arquivo dentro de metas/previews/bins
     """
     input_path = input_path.resolve()
 
-    # Caso 1: usuário passou uma pasta
+    # =========================================================
+    # Dataset layout
+    # =========================================================
+    dataset_root = find_dataset_layout_root(input_path)
+    if dataset_root is not None:
+        entries = list_capture_groups_from_dir(dataset_root)
+
+        if input_path.is_dir():
+            return entries, 0
+
+        if input_path.suffix.lower() == ".json":
+            target_json = input_path.resolve()
+        else:
+            target_json = resolve_dataset_json_from_any_path(input_path, dataset_root).resolve()
+
+        try:
+            idx = entries.index(target_json)
+        except ValueError:
+            idx = 0
+
+        return entries, idx
+
+    # =========================================================
+    # Layout antigo
+    # =========================================================
     if input_path.is_dir():
         entries = list_capture_groups_from_dir(input_path)
         return entries, 0
 
-    # Caso 2: usuário passou arquivo
     if not input_path.exists():
         raise FileNotFoundError(f"Arquivo não encontrado: {input_path}")
 
     folder = input_path.parent
     entries = list_capture_groups_from_dir(folder)
 
-    # Tenta descobrir qual JSON corresponde ao input
     if input_path.suffix.lower() == ".json":
         target_json = input_path.resolve()
     else:
@@ -814,6 +1075,8 @@ def render_group_to_canvas(json_path: Path, max_width: int):
         "group": group,
         "json": group["json"],
         "png": group["png"],
+        "mask": group.get("mask"),
+        "dataset_root": group.get("dataset_root"),
         "final_raw": group["final_raw"],
         "cameras": group["cameras"],
         "meta": meta,
@@ -854,6 +1117,8 @@ def main():
         print(f"[{current_idx + 1}/{len(entries)}]")
         print("Entrada JSON   :", info["json"])
         print("PNG            :", info["png"])
+        print("Dataset root   :", info.get("dataset_root"))
+        print("Mask           :", info.get("mask"))
         print("Final RAW      :", info["final_raw"])
         print("Câmeras        :", {k: str(v) for k, v in info["cameras"].items()})
         print("saved_payload_type :", meta.get("saved_payload_type"))