654 lines
22 KiB
Python
654 lines
22 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
corridor_state_labeler.py
|
|
-------------------------------------------------
|
|
Rotulador rápido de estado global do corredor para dataset de segmentação.
|
|
|
|
Objetivo
|
|
- Percorrer uma raiz de dataset no formato:
|
|
|
|
dataset_root/
|
|
├── grupo_1/
|
|
│ ├── images/
|
|
│ ├── masks/
|
|
│ └── labels/ # criado automaticamente
|
|
├── grupo_2/
|
|
│ ├── images/
|
|
│ ├── masks/
|
|
│ └── labels/
|
|
└── grupo_3/
|
|
├── images/
|
|
├── masks/
|
|
└── labels/
|
|
|
|
- Mostrar imagem original e máscara lado a lado.
|
|
- Permitir classificar cada frame com teclas numéricas.
|
|
- Salvar um JSON por imagem dentro de labels/.
|
|
- Suportar resume, skip, voltar, avançar e undo.
|
|
|
|
Exemplo de labels globais:
|
|
Direcionando,EntrandoRua,CaminhandoRua,SaindoRua
|
|
|
|
Dependências:
|
|
pip install pillow
|
|
|
|
Uso com GUI:
|
|
python corridor_state_labeler.py
|
|
|
|
Uso por linha de comando:
|
|
python corridor_state_labeler.py --dataset-root "C:/.../oak-d/dataset/original/group" --states Direcionando EntrandoRua CaminhandoRua SaindoRua --resume
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import tkinter as tk
|
|
from tkinter import filedialog, messagebox
|
|
|
|
try:
|
|
from PIL import Image, ImageTk, ImageOps
|
|
except Exception:
|
|
print("ERRO: Pillow não encontrado. Instale com: pip install pillow", file=sys.stderr)
|
|
raise
|
|
|
|
|
|
IMG_EXTS = [".png", ".jpg", ".jpeg", ".bmp", ".webp"]
|
|
MASK_EXTS = [".png", ".jpg", ".jpeg", ".bmp", ".webp"]
|
|
|
|
IMAGE_DIR_NAMES = ["images", "image", "imgs", "rgb"]
|
|
MASK_DIR_NAMES = ["masks", "mask", "segmentacao", "seg", "segmentations"]
|
|
LABEL_DIR_NAME = "labels"
|
|
|
|
|
|
@dataclass
|
|
class SampleItem:
|
|
group: str
|
|
image_path: Path
|
|
mask_path: Path
|
|
label_path: Path
|
|
base: str
|
|
|
|
|
|
def normalize_stem(stem: str) -> str:
|
|
"""
|
|
Remove sufixos comuns para tentar casar image e mask mesmo quando nomes diferem.
|
|
Exemplos:
|
|
001_rgb -> 001
|
|
001_image -> 001
|
|
001_mask -> 001
|
|
001_segmentacao -> 001
|
|
"""
|
|
suffixes = [
|
|
"_rgb", "_RGB", "_Rgb",
|
|
"_image", "_img", "_frame",
|
|
"_mask", "_masks",
|
|
"_seg", "_SEG", "_segment", "_segmentacao", "_Segmentacao",
|
|
]
|
|
out = stem
|
|
changed = True
|
|
while changed:
|
|
changed = False
|
|
for sfx in suffixes:
|
|
if out.endswith(sfx):
|
|
out = out[: -len(sfx)]
|
|
changed = True
|
|
break
|
|
return out
|
|
|
|
|
|
def natural_key(text: str):
|
|
"""Ordenação natural simples: 2 antes de 10."""
|
|
import re
|
|
parts = re.split(r"(\d+)", text)
|
|
return [int(p) if p.isdigit() else p.lower() for p in parts]
|
|
|
|
|
|
def find_first_existing_dir(group_dir: Path, candidates: List[str]) -> Optional[Path]:
|
|
for name in candidates:
|
|
p = group_dir / name
|
|
if p.is_dir():
|
|
return p
|
|
return None
|
|
|
|
|
|
def list_image_files(folder: Path, exts: List[str]) -> List[Path]:
|
|
files = []
|
|
for p in folder.iterdir():
|
|
if p.is_file() and p.suffix.lower() in exts:
|
|
files.append(p)
|
|
files.sort(key=lambda x: natural_key(x.name))
|
|
return files
|
|
|
|
|
|
def collect_samples(dataset_root: Path) -> List[SampleItem]:
|
|
"""
|
|
Procura grupos dentro de dataset_root.
|
|
Cada grupo precisa ter uma pasta images/ e masks/.
|
|
Cria labels/ automaticamente.
|
|
"""
|
|
if not dataset_root.exists() or not dataset_root.is_dir():
|
|
raise FileNotFoundError(f"Raiz inválida: {dataset_root}")
|
|
|
|
samples: List[SampleItem] = []
|
|
|
|
group_dirs = [p for p in dataset_root.iterdir() if p.is_dir()]
|
|
group_dirs.sort(key=lambda x: natural_key(x.name))
|
|
|
|
for group_dir in group_dirs:
|
|
image_dir = find_first_existing_dir(group_dir, IMAGE_DIR_NAMES)
|
|
mask_dir = find_first_existing_dir(group_dir, MASK_DIR_NAMES)
|
|
|
|
if image_dir is None or mask_dir is None:
|
|
continue
|
|
|
|
label_dir = group_dir / LABEL_DIR_NAME
|
|
label_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
image_files = list_image_files(image_dir, IMG_EXTS)
|
|
mask_files = list_image_files(mask_dir, MASK_EXTS)
|
|
|
|
mask_by_base: Dict[str, Path] = {}
|
|
for m in mask_files:
|
|
mask_by_base[normalize_stem(m.stem)] = m
|
|
|
|
for img in image_files:
|
|
base = normalize_stem(img.stem)
|
|
mask = mask_by_base.get(base)
|
|
if mask is None:
|
|
# fallback: tenta mesmo stem exato
|
|
exact_candidates = [m for m in mask_files if m.stem == img.stem]
|
|
mask = exact_candidates[0] if exact_candidates else None
|
|
|
|
if mask is None:
|
|
continue
|
|
|
|
label_path = label_dir / f"{base}.json"
|
|
samples.append(
|
|
SampleItem(
|
|
group=group_dir.name,
|
|
image_path=img,
|
|
mask_path=mask,
|
|
label_path=label_path,
|
|
base=base,
|
|
)
|
|
)
|
|
|
|
samples.sort(key=lambda s: (natural_key(s.group), natural_key(s.base)))
|
|
return samples
|
|
|
|
|
|
def read_existing_label(label_path: Path) -> Optional[Dict]:
|
|
if not label_path.exists():
|
|
return None
|
|
try:
|
|
with label_path.open("r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def safe_relative(path: Path, root: Path) -> str:
|
|
try:
|
|
return str(path.resolve().relative_to(root.resolve())).replace("\\", "/")
|
|
except Exception:
|
|
return str(path).replace("\\", "/")
|
|
|
|
|
|
class CorridorStateLabelerApp:
|
|
def __init__(
|
|
self,
|
|
dataset_root: Path,
|
|
samples: List[SampleItem],
|
|
states: List[str],
|
|
resume: bool = True,
|
|
display_height: int = 520,
|
|
show_existing: bool = False,
|
|
):
|
|
self.dataset_root = dataset_root
|
|
self.states = states
|
|
self.display_height = int(display_height)
|
|
self.show_existing = show_existing
|
|
|
|
if resume:
|
|
self.samples = [s for s in samples if not s.label_path.exists()]
|
|
else:
|
|
self.samples = samples
|
|
|
|
self.idx = 0
|
|
self.history: List[Dict] = []
|
|
|
|
self.root = tk.Tk()
|
|
self.root.title("Agrobot - Corridor State Labeler")
|
|
self.root.geometry("1600x880")
|
|
self.root.bind("<Key>", self.on_key)
|
|
|
|
self.top_frame = tk.Frame(self.root)
|
|
self.top_frame.pack(side=tk.TOP, fill=tk.X)
|
|
|
|
self.info_label = tk.Label(self.top_frame, text="", font=("Segoe UI", 11), anchor="w")
|
|
self.info_label.pack(side=tk.LEFT, padx=10, pady=6, fill=tk.X, expand=True)
|
|
|
|
self.legend_label = tk.Label(self.top_frame, text=self.build_legend_text(), font=("Segoe UI", 10), anchor="e")
|
|
self.legend_label.pack(side=tk.RIGHT, padx=10, pady=6)
|
|
|
|
self.image_frame = tk.Frame(self.root)
|
|
self.image_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
|
|
|
|
self.left_frame = tk.LabelFrame(self.image_frame, text="Imagem")
|
|
self.left_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=8, pady=8)
|
|
|
|
self.right_frame = tk.LabelFrame(self.image_frame, text="Máscara")
|
|
self.right_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=8, pady=8)
|
|
|
|
self.image_label = tk.Label(self.left_frame)
|
|
self.image_label.pack(fill=tk.BOTH, expand=True)
|
|
|
|
self.mask_label = tk.Label(self.right_frame)
|
|
self.mask_label.pack(fill=tk.BOTH, expand=True)
|
|
|
|
self.current_label_var = tk.StringVar(value="")
|
|
self.current_label = tk.Label(self.root, textvariable=self.current_label_var, font=("Segoe UI", 11), anchor="w")
|
|
self.current_label.pack(side=tk.BOTTOM, fill=tk.X, padx=8, pady=2)
|
|
|
|
self.status_var = tk.StringVar(value="Pronto.")
|
|
self.status_label = tk.Label(self.root, textvariable=self.status_var, font=("Segoe UI", 10), anchor="w")
|
|
self.status_label.pack(side=tk.BOTTOM, fill=tk.X, padx=8, pady=4)
|
|
|
|
self.footer = tk.Label(
|
|
self.root,
|
|
text="1..9/0=classificar | Espaço=saltar | →/n=próxima | ←/p=anterior | b=desfazer | r=recarregar | q/Esc=sair",
|
|
font=("Segoe UI", 10),
|
|
)
|
|
self.footer.pack(side=tk.BOTTOM, fill=tk.X, pady=2)
|
|
|
|
self.render()
|
|
|
|
def build_legend_text(self) -> str:
|
|
parts = []
|
|
for i, state in enumerate(self.states, start=1):
|
|
key = i if i <= 9 else 0
|
|
parts.append(f"[{key}] {state}")
|
|
return " | ".join(parts)
|
|
|
|
def load_and_resize(self, path: Path) -> ImageTk.PhotoImage:
|
|
img = Image.open(path).convert("RGB")
|
|
img = ImageOps.exif_transpose(img)
|
|
|
|
w, h = img.size
|
|
if h <= 0:
|
|
raise ValueError(f"Imagem inválida: {path}")
|
|
|
|
new_h = self.display_height
|
|
new_w = max(1, int(w * (new_h / h)))
|
|
img = img.resize((new_w, new_h), Image.BILINEAR)
|
|
return ImageTk.PhotoImage(img)
|
|
|
|
def render(self):
|
|
if not self.samples:
|
|
messagebox.showinfo("Fim", "Nenhum item para rotular. Talvez tudo já esteja rotulado com resume ligado.")
|
|
self.root.destroy()
|
|
return
|
|
|
|
self.idx = max(0, min(self.idx, len(self.samples) - 1))
|
|
item = self.samples[self.idx]
|
|
|
|
try:
|
|
self.tk_img = self.load_and_resize(item.image_path)
|
|
self.tk_mask = self.load_and_resize(item.mask_path)
|
|
self.image_label.configure(image=self.tk_img)
|
|
self.mask_label.configure(image=self.tk_mask)
|
|
except Exception as e:
|
|
self.status_var.set(f"ERRO ao abrir imagem/máscara: {e}")
|
|
return
|
|
|
|
existing = read_existing_label(item.label_path)
|
|
existing_txt = ""
|
|
if existing:
|
|
existing_txt = f" | Já rotulado: {existing.get('estado_corredor')}"
|
|
|
|
self.info_label.configure(
|
|
text=(
|
|
f"{self.idx + 1}/{len(self.samples)} | Grupo: {item.group} | Base: {item.base} "
|
|
f"| Imagem: {item.image_path.name} | Máscara: {item.mask_path.name}{existing_txt}"
|
|
)
|
|
)
|
|
|
|
self.current_label_var.set(
|
|
f"Arquivo de label: {safe_relative(item.label_path, self.dataset_root)}"
|
|
)
|
|
|
|
def save_label(self, state_index: int):
|
|
if state_index < 0 or state_index >= len(self.states):
|
|
return
|
|
|
|
item = self.samples[self.idx]
|
|
state = self.states[state_index]
|
|
|
|
previous_content = None
|
|
previous_existed = item.label_path.exists()
|
|
if previous_existed:
|
|
try:
|
|
previous_content = item.label_path.read_text(encoding="utf-8")
|
|
except Exception:
|
|
previous_content = None
|
|
|
|
rec = {
|
|
"estado_corredor": state,
|
|
"label_id": state_index,
|
|
"states": self.states,
|
|
"group": item.group,
|
|
"base": item.base,
|
|
"image": safe_relative(item.image_path, self.dataset_root),
|
|
"mask": safe_relative(item.mask_path, self.dataset_root),
|
|
"label": safe_relative(item.label_path, self.dataset_root),
|
|
"timestamp": datetime.now().isoformat(timespec="seconds"),
|
|
"source": "corridor_state_labeler",
|
|
}
|
|
|
|
try:
|
|
item.label_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with item.label_path.open("w", encoding="utf-8") as f:
|
|
json.dump(rec, f, ensure_ascii=False, indent=2)
|
|
|
|
self.history.append(
|
|
{
|
|
"action": "label",
|
|
"index": self.idx,
|
|
"label_path": item.label_path,
|
|
"previous_existed": previous_existed,
|
|
"previous_content": previous_content,
|
|
"new_state": state,
|
|
}
|
|
)
|
|
|
|
self.status_var.set(f"Salvo: {item.group}/{item.base} → {state}")
|
|
self.idx += 1
|
|
if self.idx >= len(self.samples):
|
|
messagebox.showinfo("Concluído", "Você chegou ao final da fila.")
|
|
self.root.destroy()
|
|
return
|
|
self.render()
|
|
|
|
except Exception as e:
|
|
messagebox.showerror("Erro", f"Falha ao salvar label: {e}")
|
|
self.status_var.set(f"ERRO ao salvar label: {e}")
|
|
|
|
def skip(self):
|
|
item = self.samples[self.idx]
|
|
self.history.append({"action": "skip", "index": self.idx})
|
|
self.status_var.set(f"Pulou: {item.group}/{item.base}")
|
|
self.idx += 1
|
|
if self.idx >= len(self.samples):
|
|
messagebox.showinfo("Concluído", "Você chegou ao final da fila.")
|
|
self.root.destroy()
|
|
return
|
|
self.render()
|
|
|
|
def undo(self):
|
|
if not self.history:
|
|
self.status_var.set("Nada para desfazer.")
|
|
return
|
|
|
|
last = self.history.pop()
|
|
action = last.get("action")
|
|
|
|
if action == "label":
|
|
label_path: Path = last["label_path"]
|
|
try:
|
|
if last.get("previous_existed"):
|
|
previous_content = last.get("previous_content")
|
|
if previous_content is not None:
|
|
label_path.write_text(previous_content, encoding="utf-8")
|
|
else:
|
|
if label_path.exists():
|
|
label_path.unlink()
|
|
|
|
self.idx = max(0, min(int(last.get("index", self.idx)), len(self.samples) - 1))
|
|
self.status_var.set(f"Desfeito label: {label_path.name}")
|
|
self.render()
|
|
except Exception as e:
|
|
messagebox.showerror("Erro", f"Falha ao desfazer: {e}")
|
|
self.status_var.set(f"ERRO no undo: {e}")
|
|
|
|
elif action == "skip":
|
|
self.idx = max(0, min(int(last.get("index", self.idx)), len(self.samples) - 1))
|
|
self.status_var.set("Desfeito pulo.")
|
|
self.render()
|
|
|
|
def next_item(self):
|
|
self.idx = min(len(self.samples) - 1, self.idx + 1)
|
|
self.status_var.set("Avançou uma imagem.")
|
|
self.render()
|
|
|
|
def prev_item(self):
|
|
self.idx = max(0, self.idx - 1)
|
|
self.status_var.set("Voltou uma imagem.")
|
|
self.render()
|
|
|
|
def on_key(self, event):
|
|
ch = event.keysym.lower()
|
|
|
|
if ch in [str(i) for i in range(1, 10)] or ch == "0":
|
|
idx = 9 if ch == "0" else int(ch) - 1
|
|
self.save_label(idx)
|
|
return
|
|
|
|
if ch in ("space", "s"):
|
|
self.skip()
|
|
return
|
|
|
|
if ch in ("right", "n"):
|
|
self.next_item()
|
|
return
|
|
|
|
if ch in ("left", "p"):
|
|
self.prev_item()
|
|
return
|
|
|
|
if ch == "b":
|
|
self.undo()
|
|
return
|
|
|
|
if ch == "r":
|
|
self.status_var.set("Recarregado.")
|
|
self.render()
|
|
return
|
|
|
|
if ch in ("q", "escape"):
|
|
self.root.destroy()
|
|
return
|
|
|
|
def run(self):
|
|
self.root.mainloop()
|
|
|
|
|
|
class SetupWindow:
|
|
def __init__(self):
|
|
self.root = tk.Tk()
|
|
self.root.title("Configurar - Corridor State Labeler")
|
|
self.root.geometry("760x430")
|
|
|
|
frm_root = tk.LabelFrame(self.root, text="Raiz do dataset")
|
|
frm_root.pack(fill=tk.X, padx=10, pady=8)
|
|
|
|
self.dataset_root_var = tk.StringVar(value="")
|
|
tk.Entry(frm_root, textvariable=self.dataset_root_var).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=6, pady=6)
|
|
tk.Button(frm_root, text="Escolher...", command=self.choose_dataset_root).pack(side=tk.RIGHT, padx=6, pady=6)
|
|
|
|
frm_states = tk.LabelFrame(self.root, text="Estados do corredor separados por vírgula")
|
|
frm_states.pack(fill=tk.X, padx=10, pady=8)
|
|
|
|
self.states_var = tk.StringVar(value="Parado,EntrandoRua,CaminhandoRua,SaindoRua,Manobrando,Direcionando,RetornandoBase,Indefinido")
|
|
tk.Entry(frm_states, textvariable=self.states_var).pack(fill=tk.X, padx=6, pady=6)
|
|
|
|
frm_opts = tk.LabelFrame(self.root, text="Opções")
|
|
frm_opts.pack(fill=tk.X, padx=10, pady=8)
|
|
|
|
self.resume_var = tk.BooleanVar(value=True)
|
|
self.height_var = tk.IntVar(value=520)
|
|
|
|
tk.Checkbutton(frm_opts, text="Retomar: pular imagens que já possuem labels/*.json", variable=self.resume_var).pack(anchor="w", padx=6, pady=4)
|
|
|
|
frm_height = tk.Frame(frm_opts)
|
|
frm_height.pack(fill=tk.X, padx=6, pady=4)
|
|
tk.Label(frm_height, text="Altura de exibição em px:").pack(side=tk.LEFT)
|
|
tk.Entry(frm_height, textvariable=self.height_var, width=8).pack(side=tk.LEFT, padx=6)
|
|
|
|
self.preview_var = tk.StringVar(value="")
|
|
tk.Label(self.root, textvariable=self.preview_var, fg="#555", anchor="w", justify="left").pack(fill=tk.X, padx=12, pady=4)
|
|
|
|
btn_frame = tk.Frame(self.root)
|
|
btn_frame.pack(fill=tk.X, padx=10, pady=10)
|
|
tk.Button(btn_frame, text="Verificar dataset", command=self.preview_dataset).pack(side=tk.LEFT, padx=4)
|
|
tk.Button(btn_frame, text="Iniciar classificação", command=self.start).pack(side=tk.RIGHT, padx=4)
|
|
|
|
tk.Label(
|
|
self.root,
|
|
text="Teclas: 1..9/0=classificar | Espaço=saltar | b=desfazer | ←/→ navegar | q/Esc=sair",
|
|
fg="#555",
|
|
).pack(pady=4)
|
|
|
|
self.result = None
|
|
|
|
def choose_dataset_root(self):
|
|
p = filedialog.askdirectory(title="Selecione a raiz do dataset")
|
|
if p:
|
|
self.dataset_root_var.set(p)
|
|
self.preview_dataset()
|
|
|
|
def parse_states(self) -> List[str]:
|
|
return [s.strip() for s in self.states_var.get().split(",") if s.strip()]
|
|
|
|
def preview_dataset(self):
|
|
root_raw = self.dataset_root_var.get().strip()
|
|
if not root_raw:
|
|
self.preview_var.set("Selecione uma raiz de dataset.")
|
|
return
|
|
root = Path(root_raw)
|
|
try:
|
|
samples = collect_samples(root)
|
|
groups = sorted(set(s.group for s in samples), key=natural_key)
|
|
labeled = sum(1 for s in samples if s.label_path.exists())
|
|
self.preview_var.set(
|
|
f"Amostras encontradas: {len(samples)} | Já rotuladas: {labeled} | Grupos: {', '.join(groups) if groups else '-'}"
|
|
)
|
|
except Exception as e:
|
|
self.preview_var.set(f"Erro ao verificar dataset: {e}")
|
|
|
|
def start(self):
|
|
root_raw = self.dataset_root_var.get().strip()
|
|
states = self.parse_states()
|
|
|
|
if not root_raw:
|
|
messagebox.showwarning("Faltando raiz", "Selecione a raiz do dataset.")
|
|
return
|
|
if not states:
|
|
messagebox.showwarning("Faltando estados", "Informe pelo menos um estado.")
|
|
return
|
|
if len(states) > 10:
|
|
messagebox.showwarning("Muitos estados", "Este rotulador suporta até 10 estados nas teclas 1..9 e 0.")
|
|
return
|
|
|
|
root = Path(root_raw)
|
|
try:
|
|
samples = collect_samples(root)
|
|
except Exception as e:
|
|
messagebox.showerror("Erro", f"Falha ao ler dataset: {e}")
|
|
return
|
|
|
|
if not samples:
|
|
messagebox.showinfo(
|
|
"Sem amostras",
|
|
"Nenhum par image/mask encontrado. Verifique se cada grupo possui images/ e masks/.",
|
|
)
|
|
return
|
|
|
|
self.result = {
|
|
"dataset_root": root,
|
|
"states": states,
|
|
"resume": self.resume_var.get(),
|
|
"display_height": self.height_var.get(),
|
|
}
|
|
self.root.destroy()
|
|
|
|
def run(self):
|
|
self.root.mainloop()
|
|
return self.result
|
|
|
|
|
|
def run_gui_setup():
|
|
setup = SetupWindow()
|
|
res = setup.run()
|
|
if not res:
|
|
return
|
|
|
|
dataset_root = res["dataset_root"]
|
|
samples = collect_samples(dataset_root)
|
|
|
|
app = CorridorStateLabelerApp(
|
|
dataset_root=dataset_root,
|
|
samples=samples,
|
|
states=res["states"],
|
|
resume=res["resume"],
|
|
display_height=res["display_height"],
|
|
)
|
|
app.run()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Rotulador de estado global do corredor para dataset de segmentação.")
|
|
parser.add_argument("--dataset-root", help="Raiz do dataset contendo grupos com images/ e masks/.")
|
|
parser.add_argument("--states", nargs="+", help="Estados do corredor. Ex: Direcionando EntrandoRua CaminhandoRua SaindoRua")
|
|
parser.add_argument("--resume", action="store_true", help="Pular itens que já possuem labels/*.json.")
|
|
parser.add_argument("--no-resume", action="store_true", help="Não pular itens já rotulados.")
|
|
parser.add_argument("--display-height", type=int, default=520, help="Altura de exibição das imagens em px.")
|
|
parser.add_argument("--no-gui-setup", action="store_true", help="Não abrir janela de configuração.")
|
|
|
|
args = parser.parse_args()
|
|
|
|
need_gui = not args.no_gui_setup and (not args.dataset_root or not args.states)
|
|
if need_gui:
|
|
run_gui_setup()
|
|
return
|
|
|
|
if not args.dataset_root or not args.states:
|
|
print("ERRO: informe --dataset-root e --states, ou rode sem argumentos para abrir a GUI.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
dataset_root = Path(args.dataset_root)
|
|
states = args.states
|
|
if len(states) > 10:
|
|
print("ERRO: máximo de 10 estados: teclas 1..9 e 0.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
samples = collect_samples(dataset_root)
|
|
if not samples:
|
|
print("ERRO: nenhum par image/mask encontrado.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
resume = True
|
|
if args.no_resume:
|
|
resume = False
|
|
elif args.resume:
|
|
resume = True
|
|
|
|
app = CorridorStateLabelerApp(
|
|
dataset_root=dataset_root,
|
|
samples=samples,
|
|
states=states,
|
|
resume=resume,
|
|
display_height=args.display_height,
|
|
)
|
|
app.run()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|