criacao de testes para modo stereo
This commit is contained in:
Diego Freitas
parent
c304401ef0
commit
b4397d0954
|
|
@ -69,6 +69,9 @@ Python/OAK/datasets/oak-fcc-3/backup/
|
|||
Python/OAK/datasets/oak-fcc-3/dataset/
|
||||
Python/OAK/datasets/oak-fcc-3/audit_multispec_out/
|
||||
Python/OAK/datasets/oak-fcc-3/depth_probe_out/
|
||||
Python/OAK/datasets/oak-fcc-3/calibration/multicam_charuco_calib_out/debug/
|
||||
Python/OAK/datasets/oak-fcc-3/calibration/stereo_charuco_calib_out/debug/
|
||||
Python/OAK/datasets/oak-fcc-3/calibration/stereo_dataset/
|
||||
Python/OAK/datasets/oak-fcc-3/.cache/
|
||||
Python/OAK/datasets/gal5000/dataset/
|
||||
Python/OAK/datasets/gal5000/backup/
|
||||
|
|
|
|||
Binary file not shown.
|
|
@ -0,0 +1,236 @@
|
|||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import pipeline
|
||||
|
||||
|
||||
IMG_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
|
||||
|
||||
|
||||
def robust_normalize(depth: np.ndarray, invert: bool = False) -> np.ndarray:
|
||||
d = depth.astype(np.float32)
|
||||
|
||||
p2 = np.percentile(d, 2)
|
||||
p98 = np.percentile(d, 98)
|
||||
|
||||
dn = (d - p2) / (p98 - p2 + 1e-6)
|
||||
dn = np.clip(dn, 0.0, 1.0)
|
||||
|
||||
if invert:
|
||||
dn = 1.0 - dn
|
||||
|
||||
return dn
|
||||
|
||||
|
||||
def depth_to_colormap(depth_norm: np.ndarray) -> np.ndarray:
|
||||
u8 = (depth_norm * 255).astype(np.uint8)
|
||||
return cv2.applyColorMap(u8, cv2.COLORMAP_TURBO)
|
||||
|
||||
|
||||
def make_bands(depth_norm: np.ndarray) -> np.ndarray:
|
||||
bands = np.zeros_like(depth_norm, dtype=np.uint8)
|
||||
bands[depth_norm >= 0.33] = 1
|
||||
bands[depth_norm >= 0.66] = 2
|
||||
|
||||
out = np.zeros((bands.shape[0], bands.shape[1], 3), dtype=np.uint8)
|
||||
|
||||
# BGR
|
||||
out[bands == 0] = (80, 80, 255) # faixa 0
|
||||
out[bands == 1] = (80, 255, 255) # faixa 1
|
||||
out[bands == 2] = (80, 255, 80) # faixa 2
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def run_depth(pipe, image_path: Path, invert: bool):
|
||||
img_pil = Image.open(image_path).convert("RGB")
|
||||
rgb = np.array(img_pil)
|
||||
preview_bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
|
||||
|
||||
result = pipe(img_pil)
|
||||
|
||||
if "predicted_depth" in result:
|
||||
depth = result["predicted_depth"]
|
||||
|
||||
if hasattr(depth, "detach"):
|
||||
depth = depth.detach().cpu().numpy()
|
||||
|
||||
depth = np.array(depth).squeeze().astype(np.float32)
|
||||
|
||||
else:
|
||||
depth_img = result["depth"]
|
||||
depth = np.array(depth_img).astype(np.float32)
|
||||
|
||||
if depth.ndim == 3:
|
||||
depth = cv2.cvtColor(
|
||||
depth.astype(np.uint8),
|
||||
cv2.COLOR_RGB2GRAY
|
||||
).astype(np.float32)
|
||||
|
||||
depth = cv2.resize(
|
||||
depth,
|
||||
(preview_bgr.shape[1], preview_bgr.shape[0]),
|
||||
interpolation=cv2.INTER_CUBIC
|
||||
)
|
||||
|
||||
depth_norm = robust_normalize(depth, invert=invert)
|
||||
depth_color = depth_to_colormap(depth_norm)
|
||||
depth_bands = make_bands(depth_norm)
|
||||
|
||||
return preview_bgr, depth_norm, depth_color, depth_bands
|
||||
|
||||
|
||||
def resize_to_height(img: np.ndarray, target_h: int) -> np.ndarray:
|
||||
h, w = img.shape[:2]
|
||||
if h == target_h:
|
||||
return img
|
||||
|
||||
scale = target_h / h
|
||||
new_w = int(w * scale)
|
||||
return cv2.resize(img, (new_w, target_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
|
||||
def compose_view(preview_bgr, depth_color, depth_bands, image_path, index, total, mode):
|
||||
target_h = 520
|
||||
|
||||
left = resize_to_height(preview_bgr, target_h)
|
||||
|
||||
if mode == "depth":
|
||||
right_img = depth_color
|
||||
right_title = "Depth Anything V2"
|
||||
else:
|
||||
right_img = depth_bands
|
||||
right_title = "Depth bands"
|
||||
|
||||
right = resize_to_height(right_img, target_h)
|
||||
|
||||
# Garante mesma altura
|
||||
h = min(left.shape[0], right.shape[0])
|
||||
left = left[:h]
|
||||
right = right[:h]
|
||||
|
||||
canvas = np.hstack([left, right])
|
||||
|
||||
text1 = f"{index + 1}/{total} - {image_path.name}"
|
||||
text2 = f"Modo: {right_title} | N/SPACE prox | A ant | M modo | I invert | S salvar | Q sair"
|
||||
|
||||
cv2.rectangle(canvas, (0, 0), (canvas.shape[1], 58), (0, 0, 0), -1)
|
||||
cv2.putText(canvas, text1, (12, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.58, (255, 255, 255), 1, cv2.LINE_AA)
|
||||
cv2.putText(canvas, text2, (12, 48), cv2.FONT_HERSHEY_SIMPLEX, 0.50, (220, 220, 220), 1, cv2.LINE_AA)
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input_dir", required=True, help="Pasta com imagens preview PNG/JPG")
|
||||
parser.add_argument("--model", default="depth-anything/Depth-Anything-V2-Small-hf")
|
||||
parser.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"])
|
||||
parser.add_argument("--invert", action="store_true")
|
||||
parser.add_argument("--save_dir", default="depth_viewer_saves")
|
||||
args = parser.parse_args()
|
||||
|
||||
input_dir = Path(args.input_dir)
|
||||
save_dir = Path(args.save_dir)
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
image_paths = sorted([
|
||||
p for p in input_dir.rglob("*")
|
||||
if p.suffix.lower() in IMG_EXTS
|
||||
])
|
||||
|
||||
if not image_paths:
|
||||
raise RuntimeError(f"Nenhuma imagem encontrada em: {input_dir}")
|
||||
|
||||
if args.device == "auto":
|
||||
device = 0 if torch.cuda.is_available() else -1
|
||||
elif args.device == "cuda":
|
||||
device = 0
|
||||
else:
|
||||
device = -1
|
||||
|
||||
print(f"[INFO] imagens: {len(image_paths)}")
|
||||
print(f"[INFO] modelo: {args.model}")
|
||||
print(f"[INFO] device: {'cuda' if device == 0 else 'cpu'}")
|
||||
print("[INFO] controles:")
|
||||
print(" N ou SPACE = próxima")
|
||||
print(" A = anterior")
|
||||
print(" M = alterna depth/faixas")
|
||||
print(" I = inverte depth")
|
||||
print(" S = salva visual atual")
|
||||
print(" Q ou ESC = sair")
|
||||
|
||||
pipe = pipeline(
|
||||
task="depth-estimation",
|
||||
model=args.model,
|
||||
device=device
|
||||
)
|
||||
|
||||
idx = 0
|
||||
invert = args.invert
|
||||
mode = "depth"
|
||||
|
||||
cached_path = None
|
||||
cached_data = None
|
||||
|
||||
cv2.namedWindow("Depth Anything V2 Viewer", cv2.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
image_path = image_paths[idx]
|
||||
|
||||
need_reprocess = cached_path != image_path or cached_data is None
|
||||
|
||||
if need_reprocess:
|
||||
print(f"[RUN] {idx + 1}/{len(image_paths)} - {image_path.name}")
|
||||
preview_bgr, depth_norm, depth_color, depth_bands = run_depth(pipe, image_path, invert=invert)
|
||||
cached_data = (preview_bgr, depth_norm, depth_color, depth_bands)
|
||||
cached_path = image_path
|
||||
else:
|
||||
preview_bgr, depth_norm, depth_color, depth_bands = cached_data
|
||||
|
||||
view = compose_view(
|
||||
preview_bgr=preview_bgr,
|
||||
depth_color=depth_color,
|
||||
depth_bands=depth_bands,
|
||||
image_path=image_path,
|
||||
index=idx,
|
||||
total=len(image_paths),
|
||||
mode=mode
|
||||
)
|
||||
|
||||
cv2.imshow("Depth Anything V2 Viewer", view)
|
||||
key = cv2.waitKey(0) & 0xFF
|
||||
|
||||
if key in [27, ord("q"), ord("Q")]:
|
||||
break
|
||||
|
||||
elif key in [ord("n"), ord("N"), 32]:
|
||||
idx = min(idx + 1, len(image_paths) - 1)
|
||||
cached_path = None
|
||||
|
||||
elif key in [ord("a"), ord("A")]:
|
||||
idx = max(idx - 1, 0)
|
||||
cached_path = None
|
||||
|
||||
elif key in [ord("m"), ord("M")]:
|
||||
mode = "bands" if mode == "depth" else "depth"
|
||||
|
||||
elif key in [ord("i"), ord("I")]:
|
||||
invert = not invert
|
||||
cached_path = None
|
||||
print(f"[INFO] invert={invert}")
|
||||
|
||||
elif key in [ord("s"), ord("S")]:
|
||||
out_path = save_dir / f"{image_path.stem}_viewer_{mode}.png"
|
||||
cv2.imwrite(str(out_path), view)
|
||||
print(f"[SAVE] {out_path}")
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,771 @@
|
|||
import argparse
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
# ============================================================
|
||||
# RAW10
|
||||
# ============================================================
|
||||
|
||||
def unpack_raw10_packed(raw: bytes, width: int, height: int) -> np.ndarray:
|
||||
"""
|
||||
RAW10 packed padrão:
|
||||
5 bytes = 4 pixels de 10 bits.
|
||||
Retorna uint16 HxW com valores 0..1023.
|
||||
"""
|
||||
arr = np.frombuffer(raw, dtype=np.uint8)
|
||||
|
||||
pixel_count = width * height
|
||||
expected_bytes = (pixel_count // 4) * 5
|
||||
|
||||
if pixel_count % 4 != 0:
|
||||
raise RuntimeError(f"width*height precisa ser múltiplo de 4. Recebido: {pixel_count}")
|
||||
|
||||
if arr.size < expected_bytes:
|
||||
raise RuntimeError(
|
||||
f"RAW10 menor que esperado. bytes={arr.size}, esperado={expected_bytes}, "
|
||||
f"width={width}, height={height}"
|
||||
)
|
||||
|
||||
arr = arr[:expected_bytes]
|
||||
groups = arr.reshape(-1, 5).astype(np.uint16)
|
||||
|
||||
p0 = (groups[:, 0] << 2) | ((groups[:, 4] >> 0) & 0x03)
|
||||
p1 = (groups[:, 1] << 2) | ((groups[:, 4] >> 2) & 0x03)
|
||||
p2 = (groups[:, 2] << 2) | ((groups[:, 4] >> 4) & 0x03)
|
||||
p3 = (groups[:, 3] << 2) | ((groups[:, 4] >> 6) & 0x03)
|
||||
|
||||
out = np.empty(groups.shape[0] * 4, dtype=np.uint16)
|
||||
out[0::4] = p0
|
||||
out[1::4] = p1
|
||||
out[2::4] = p2
|
||||
out[3::4] = p3
|
||||
|
||||
return out.reshape(height, width)
|
||||
|
||||
|
||||
def normalize_to_u8(img: np.ndarray, p_low=1.0, p_high=99.0) -> np.ndarray:
|
||||
arr = img.astype(np.float32)
|
||||
valid = np.isfinite(arr)
|
||||
|
||||
if np.count_nonzero(valid) < 20:
|
||||
return np.zeros(arr.shape[:2], dtype=np.uint8)
|
||||
|
||||
vals = arr[valid]
|
||||
lo = np.percentile(vals, p_low)
|
||||
hi = np.percentile(vals, p_high)
|
||||
|
||||
out = (arr - lo) / (hi - lo + 1e-6)
|
||||
out = np.clip(out, 0.0, 1.0)
|
||||
|
||||
return (out * 255).astype(np.uint8)
|
||||
|
||||
|
||||
def raw10_bin_to_gray(path: Path, width: int, height: int, use_clahe=True) -> np.ndarray:
|
||||
raw = path.read_bytes()
|
||||
mono10 = unpack_raw10_packed(raw, width, height)
|
||||
gray = normalize_to_u8(mono10)
|
||||
|
||||
if use_clahe:
|
||||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
gray = clahe.apply(gray)
|
||||
|
||||
return gray
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Metadata
|
||||
# ============================================================
|
||||
|
||||
def find_meta(root_dir: Path) -> Path | None:
|
||||
candidates = list(root_dir.rglob("meta.json")) + list(root_dir.rglob("metadata.json"))
|
||||
return candidates[0] if candidates else None
|
||||
|
||||
|
||||
def extract_camera_info_from_meta(meta: dict, cam_key: str):
|
||||
for root_key in ["camera_info", "cameras", "camera_meta", "payload_sources_info"]:
|
||||
root = meta.get(root_key)
|
||||
if isinstance(root, dict):
|
||||
info = root.get(cam_key)
|
||||
if isinstance(info, dict):
|
||||
return info
|
||||
|
||||
stack = [meta]
|
||||
while stack:
|
||||
obj = stack.pop()
|
||||
if isinstance(obj, dict):
|
||||
if cam_key in obj and isinstance(obj[cam_key], dict):
|
||||
return obj[cam_key]
|
||||
for v in obj.values():
|
||||
if isinstance(v, (dict, list)):
|
||||
stack.append(v)
|
||||
elif isinstance(obj, list):
|
||||
for v in obj:
|
||||
if isinstance(v, (dict, list)):
|
||||
stack.append(v)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def try_get_width_height_from_meta(root_dir: Path, cam_key: str):
|
||||
meta_path = find_meta(root_dir)
|
||||
if meta_path is None:
|
||||
return None, None
|
||||
|
||||
try:
|
||||
meta = json.loads(meta_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
meta = json.loads(meta_path.read_text(encoding="latin-1"))
|
||||
|
||||
info = extract_camera_info_from_meta(meta, cam_key)
|
||||
if not info:
|
||||
return None, None
|
||||
|
||||
width = (
|
||||
info.get("width")
|
||||
or info.get("w")
|
||||
or info.get("sensor_width")
|
||||
or info.get("frame_width")
|
||||
)
|
||||
|
||||
height = (
|
||||
info.get("height")
|
||||
or info.get("h")
|
||||
or info.get("sensor_height")
|
||||
or info.get("frame_height")
|
||||
)
|
||||
|
||||
if width is None or height is None:
|
||||
return None, None
|
||||
|
||||
return int(width), int(height)
|
||||
|
||||
|
||||
def resolve_width_height(args, cam_key: str):
|
||||
if args.width > 0 and args.height > 0:
|
||||
return args.width, args.height
|
||||
|
||||
w, h = try_get_width_height_from_meta(Path(args.root_dir), cam_key)
|
||||
if w and h:
|
||||
return w, h
|
||||
|
||||
raise RuntimeError(
|
||||
f"Não consegui descobrir width/height para {cam_key}. "
|
||||
f"Passe manualmente: --width 1280 --height 800"
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Pairing
|
||||
# ============================================================
|
||||
|
||||
def clean_stem_for_pair(path: Path, cam_key: str):
|
||||
stem = path.stem
|
||||
s = stem
|
||||
|
||||
variants = [
|
||||
cam_key,
|
||||
cam_key.lower(),
|
||||
cam_key.replace("_", ""),
|
||||
cam_key.replace("_", "").lower(),
|
||||
]
|
||||
|
||||
for v in variants:
|
||||
s = s.replace(v, "")
|
||||
|
||||
s = re.sub(r"[_\-\s]+", "_", s).strip("_").lower()
|
||||
return s
|
||||
|
||||
|
||||
def find_cam_bins(root_dir: Path, cam_key: str):
|
||||
return sorted([
|
||||
p for p in root_dir.rglob("*.bin")
|
||||
if cam_key.lower() in p.name.lower()
|
||||
])
|
||||
|
||||
|
||||
def pair_cam_bins(root_dir: Path, left_cam: str, right_cam: str):
|
||||
left_bins = find_cam_bins(root_dir, left_cam)
|
||||
right_bins = find_cam_bins(root_dir, right_cam)
|
||||
|
||||
right_by_folder_key = {}
|
||||
right_by_key = {}
|
||||
|
||||
for rp in right_bins:
|
||||
key = clean_stem_for_pair(rp, right_cam)
|
||||
right_by_folder_key[(rp.parent, key)] = rp
|
||||
right_by_key.setdefault(key, rp)
|
||||
|
||||
pairs = []
|
||||
|
||||
for lp in left_bins:
|
||||
key = clean_stem_for_pair(lp, left_cam)
|
||||
|
||||
rp = right_by_folder_key.get((lp.parent, key))
|
||||
if rp is None:
|
||||
rp = right_by_key.get(key)
|
||||
|
||||
if rp is None:
|
||||
same_folder = [r for r in right_bins if r.parent == lp.parent]
|
||||
if len(same_folder) == 1:
|
||||
rp = same_folder[0]
|
||||
|
||||
if rp is not None:
|
||||
pairs.append((lp, rp))
|
||||
|
||||
if pairs:
|
||||
return pairs, left_bins, right_bins, "name/key"
|
||||
|
||||
n = min(len(left_bins), len(right_bins))
|
||||
return list(zip(left_bins[:n], right_bins[:n])), left_bins, right_bins, "order"
|
||||
|
||||
|
||||
# ============================================================
|
||||
# ChArUco helpers
|
||||
# ============================================================
|
||||
|
||||
def get_aruco_dict(dict_name: str):
|
||||
aruco = cv2.aruco
|
||||
|
||||
mapping = {
|
||||
"4X4_50": aruco.DICT_4X4_50,
|
||||
"4X4_100": aruco.DICT_4X4_100,
|
||||
"4X4_250": aruco.DICT_4X4_250,
|
||||
"4X4_1000": aruco.DICT_4X4_1000,
|
||||
"5X5_50": aruco.DICT_5X5_50,
|
||||
"5X5_100": aruco.DICT_5X5_100,
|
||||
"5X5_250": aruco.DICT_5X5_250,
|
||||
"5X5_1000": aruco.DICT_5X5_1000,
|
||||
"6X6_50": aruco.DICT_6X6_50,
|
||||
"6X6_100": aruco.DICT_6X6_100,
|
||||
"6X6_250": aruco.DICT_6X6_250,
|
||||
"6X6_1000": aruco.DICT_6X6_1000,
|
||||
}
|
||||
|
||||
key = dict_name.upper()
|
||||
if key not in mapping:
|
||||
raise RuntimeError(f"Dicionário ArUco não suportado: {dict_name}")
|
||||
|
||||
if hasattr(aruco, "getPredefinedDictionary"):
|
||||
return aruco.getPredefinedDictionary(mapping[key])
|
||||
|
||||
return aruco.Dictionary_get(mapping[key])
|
||||
|
||||
|
||||
def create_charuco_board(squares_x, squares_y, square_length, marker_length, aruco_dict):
|
||||
aruco = cv2.aruco
|
||||
|
||||
if hasattr(aruco, "CharucoBoard"):
|
||||
try:
|
||||
return aruco.CharucoBoard(
|
||||
(squares_x, squares_y),
|
||||
square_length,
|
||||
marker_length,
|
||||
aruco_dict
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if hasattr(aruco, "CharucoBoard_create"):
|
||||
return aruco.CharucoBoard_create(
|
||||
squares_x,
|
||||
squares_y,
|
||||
square_length,
|
||||
marker_length,
|
||||
aruco_dict
|
||||
)
|
||||
|
||||
raise RuntimeError("Sua versão do OpenCV não tem CharucoBoard/CharucoBoard_create.")
|
||||
|
||||
|
||||
def get_board_corners(board):
|
||||
if hasattr(board, "getChessboardCorners"):
|
||||
return np.array(board.getChessboardCorners(), dtype=np.float32)
|
||||
|
||||
if hasattr(board, "chessboardCorners"):
|
||||
return np.array(board.chessboardCorners, dtype=np.float32)
|
||||
|
||||
raise RuntimeError("Não consegui acessar chessboardCorners do ChArUco board.")
|
||||
|
||||
|
||||
def create_detector_params():
|
||||
aruco = cv2.aruco
|
||||
|
||||
if hasattr(aruco, "DetectorParameters"):
|
||||
return aruco.DetectorParameters()
|
||||
|
||||
if hasattr(aruco, "DetectorParameters_create"):
|
||||
return aruco.DetectorParameters_create()
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def detect_charuco(gray: np.ndarray, board, aruco_dict, min_corners: int):
|
||||
"""
|
||||
Retorna:
|
||||
charuco_corners: Nx2 float32
|
||||
charuco_ids: N int32
|
||||
marker_corners, marker_ids
|
||||
"""
|
||||
aruco = cv2.aruco
|
||||
params = create_detector_params()
|
||||
|
||||
# OpenCV novo
|
||||
if hasattr(aruco, "CharucoDetector"):
|
||||
try:
|
||||
detector = aruco.CharucoDetector(board)
|
||||
charuco_corners, charuco_ids, marker_corners, marker_ids = detector.detectBoard(gray)
|
||||
|
||||
if charuco_corners is None or charuco_ids is None:
|
||||
return None, None, marker_corners, marker_ids
|
||||
|
||||
corners = np.array(charuco_corners, dtype=np.float32).reshape(-1, 2)
|
||||
ids = np.array(charuco_ids, dtype=np.int32).reshape(-1)
|
||||
|
||||
if len(ids) < min_corners:
|
||||
return None, None, marker_corners, marker_ids
|
||||
|
||||
return corners, ids, marker_corners, marker_ids
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# OpenCV legado
|
||||
if params is not None:
|
||||
marker_corners, marker_ids, rejected = aruco.detectMarkers(
|
||||
gray,
|
||||
aruco_dict,
|
||||
parameters=params
|
||||
)
|
||||
else:
|
||||
marker_corners, marker_ids, rejected = aruco.detectMarkers(gray, aruco_dict)
|
||||
|
||||
if marker_ids is None or len(marker_ids) == 0:
|
||||
return None, None, marker_corners, marker_ids
|
||||
|
||||
try:
|
||||
aruco.refineDetectedMarkers(gray, board, marker_corners, marker_ids, rejected)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
retval, charuco_corners, charuco_ids = aruco.interpolateCornersCharuco(
|
||||
marker_corners,
|
||||
marker_ids,
|
||||
gray,
|
||||
board
|
||||
)
|
||||
|
||||
if charuco_corners is None or charuco_ids is None:
|
||||
return None, None, marker_corners, marker_ids
|
||||
|
||||
corners = np.array(charuco_corners, dtype=np.float32).reshape(-1, 2)
|
||||
ids = np.array(charuco_ids, dtype=np.int32).reshape(-1)
|
||||
|
||||
if len(ids) < min_corners:
|
||||
return None, None, marker_corners, marker_ids
|
||||
|
||||
return corners, ids, marker_corners, marker_ids
|
||||
|
||||
|
||||
def common_charuco_points(corners_l, ids_l, corners_r, ids_r, board_corners, min_common):
|
||||
map_l = {int(i): corners_l[k] for k, i in enumerate(ids_l)}
|
||||
map_r = {int(i): corners_r[k] for k, i in enumerate(ids_r)}
|
||||
|
||||
common_ids = sorted(set(map_l.keys()) & set(map_r.keys()))
|
||||
|
||||
if len(common_ids) < min_common:
|
||||
return None, None, None, common_ids
|
||||
|
||||
obj = []
|
||||
img_l = []
|
||||
img_r = []
|
||||
|
||||
max_id = len(board_corners) - 1
|
||||
|
||||
for cid in common_ids:
|
||||
if cid < 0 or cid > max_id:
|
||||
continue
|
||||
|
||||
obj.append(board_corners[cid])
|
||||
img_l.append(map_l[cid])
|
||||
img_r.append(map_r[cid])
|
||||
|
||||
if len(obj) < min_common:
|
||||
return None, None, None, common_ids
|
||||
|
||||
obj = np.array(obj, dtype=np.float32).reshape(-1, 1, 3)
|
||||
img_l = np.array(img_l, dtype=np.float32).reshape(-1, 1, 2)
|
||||
img_r = np.array(img_r, dtype=np.float32).reshape(-1, 1, 2)
|
||||
|
||||
return obj, img_l, img_r, common_ids
|
||||
|
||||
|
||||
def draw_debug(gray, charuco_corners, charuco_ids, title):
|
||||
bgr = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
if charuco_corners is not None and charuco_ids is not None:
|
||||
corners_draw = np.array(charuco_corners, dtype=np.float32).reshape(-1, 1, 2)
|
||||
ids_draw = np.array(charuco_ids, dtype=np.int32).reshape(-1, 1)
|
||||
try:
|
||||
cv2.aruco.drawDetectedCornersCharuco(bgr, corners_draw, ids_draw, (0, 255, 0))
|
||||
except Exception:
|
||||
for p in charuco_corners:
|
||||
cv2.circle(bgr, tuple(np.round(p).astype(int)), 3, (0, 255, 0), -1)
|
||||
|
||||
cv2.putText(
|
||||
bgr,
|
||||
title,
|
||||
(20, 35),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.85,
|
||||
(255, 255, 255),
|
||||
2,
|
||||
cv2.LINE_AA
|
||||
)
|
||||
|
||||
return bgr
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Calibration
|
||||
# ============================================================
|
||||
|
||||
def calibrate_single_camera(objpoints, imgpoints, image_size):
|
||||
flags = 0
|
||||
|
||||
ret, K, D, rvecs, tvecs = cv2.calibrateCamera(
|
||||
objectPoints=objpoints,
|
||||
imagePoints=imgpoints,
|
||||
imageSize=image_size,
|
||||
cameraMatrix=None,
|
||||
distCoeffs=None,
|
||||
flags=flags
|
||||
)
|
||||
|
||||
return ret, K, D, rvecs, tvecs
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument("--root_dir", required=True)
|
||||
parser.add_argument("--out_dir", default="calibration/stereo_charuco_calib_out")
|
||||
|
||||
parser.add_argument("--left_cam", default="CAM_C")
|
||||
parser.add_argument("--right_cam", default="CAM_B")
|
||||
|
||||
parser.add_argument("--width", type=int, default=-1)
|
||||
parser.add_argument("--height", type=int, default=-1)
|
||||
|
||||
# ChArUco 7x13
|
||||
parser.add_argument("--squares_x", type=int, default=7)
|
||||
parser.add_argument("--squares_y", type=int, default=13)
|
||||
|
||||
# Use a unidade que quiser. Recomendo metros.
|
||||
# Ex: quadrado de 20 mm => 0.020
|
||||
parser.add_argument("--square_length", type=float, required=True)
|
||||
parser.add_argument("--marker_length", type=float, required=True)
|
||||
|
||||
parser.add_argument("--aruco_dict", default="4X4_50")
|
||||
|
||||
parser.add_argument("--min_corners", type=int, default=12)
|
||||
parser.add_argument("--min_common", type=int, default=10)
|
||||
|
||||
parser.add_argument("--no_clahe", action="store_true")
|
||||
parser.add_argument("--show", action="store_true")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
root_dir = Path(args.root_dir)
|
||||
out_dir = Path(args.out_dir)
|
||||
debug_dir = out_dir / "debug"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
debug_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
left_cam = args.left_cam
|
||||
right_cam = args.right_cam
|
||||
|
||||
w_left, h_left = resolve_width_height(args, left_cam)
|
||||
w_right, h_right = resolve_width_height(args, right_cam)
|
||||
|
||||
if (w_left, h_left) != (w_right, h_right):
|
||||
print(f"[WARN] Resoluções diferentes: {left_cam}=({w_left},{h_left}) {right_cam}=({w_right},{h_right})")
|
||||
print("[WARN] Vou calibrar usando o menor tamanho comum após resize.")
|
||||
|
||||
image_w = min(w_left, w_right)
|
||||
image_h = min(h_left, h_right)
|
||||
image_size = (image_w, image_h)
|
||||
|
||||
print(f"[INFO] root_dir={root_dir}")
|
||||
print(f"[INFO] left={left_cam} {w_left}x{h_left}")
|
||||
print(f"[INFO] right={right_cam} {w_right}x{h_right}")
|
||||
print(f"[INFO] image_size usado={image_size}")
|
||||
print(f"[INFO] ChArUco squares={args.squares_x}x{args.squares_y}")
|
||||
print(f"[INFO] square_length={args.square_length}")
|
||||
print(f"[INFO] marker_length={args.marker_length}")
|
||||
print(f"[INFO] aruco_dict={args.aruco_dict}")
|
||||
|
||||
pairs, left_bins, right_bins, pair_mode = pair_cam_bins(root_dir, left_cam, right_cam)
|
||||
|
||||
print(f"[INFO] arquivos {left_cam}: {len(left_bins)}")
|
||||
print(f"[INFO] arquivos {right_cam}: {len(right_bins)}")
|
||||
print(f"[INFO] pares encontrados: {len(pairs)}")
|
||||
print(f"[INFO] pareamento: {pair_mode}")
|
||||
|
||||
if len(pairs) < 5:
|
||||
raise RuntimeError("Poucos pares encontrados. Capture mais imagens ou verifique nomes dos arquivos.")
|
||||
|
||||
aruco_dict = get_aruco_dict(args.aruco_dict)
|
||||
board = create_charuco_board(
|
||||
args.squares_x,
|
||||
args.squares_y,
|
||||
args.square_length,
|
||||
args.marker_length,
|
||||
aruco_dict
|
||||
)
|
||||
board_corners = get_board_corners(board)
|
||||
|
||||
stereo_objpoints = []
|
||||
stereo_imgpoints_l = []
|
||||
stereo_imgpoints_r = []
|
||||
|
||||
left_objpoints = []
|
||||
left_imgpoints = []
|
||||
|
||||
right_objpoints = []
|
||||
right_imgpoints = []
|
||||
|
||||
accepted = 0
|
||||
rejected = 0
|
||||
|
||||
for idx, (lp, rp) in enumerate(pairs):
|
||||
print(f"[{idx+1}/{len(pairs)}] L={lp.name} | R={rp.name}")
|
||||
|
||||
try:
|
||||
gray_l = raw10_bin_to_gray(lp, w_left, h_left, use_clahe=not args.no_clahe)
|
||||
gray_r = raw10_bin_to_gray(rp, w_right, h_right, use_clahe=not args.no_clahe)
|
||||
|
||||
if gray_l.shape[::-1] != image_size:
|
||||
gray_l = cv2.resize(gray_l, image_size, interpolation=cv2.INTER_AREA)
|
||||
if gray_r.shape[::-1] != image_size:
|
||||
gray_r = cv2.resize(gray_r, image_size, interpolation=cv2.INTER_AREA)
|
||||
|
||||
corners_l, ids_l, _, _ = detect_charuco(
|
||||
gray_l,
|
||||
board,
|
||||
aruco_dict,
|
||||
min_corners=args.min_corners
|
||||
)
|
||||
|
||||
corners_r, ids_r, _, _ = detect_charuco(
|
||||
gray_r,
|
||||
board,
|
||||
aruco_dict,
|
||||
min_corners=args.min_corners
|
||||
)
|
||||
|
||||
n_l = 0 if ids_l is None else len(ids_l)
|
||||
n_r = 0 if ids_r is None else len(ids_r)
|
||||
|
||||
if corners_l is None or corners_r is None:
|
||||
print(f" [REJECT] detect insuficiente: left={n_l}, right={n_r}")
|
||||
rejected += 1
|
||||
continue
|
||||
|
||||
obj, img_l, img_r, common_ids = common_charuco_points(
|
||||
corners_l,
|
||||
ids_l,
|
||||
corners_r,
|
||||
ids_r,
|
||||
board_corners,
|
||||
min_common=args.min_common
|
||||
)
|
||||
|
||||
if obj is None:
|
||||
print(f" [REJECT] comuns insuficientes: common={len(common_ids)}")
|
||||
rejected += 1
|
||||
continue
|
||||
|
||||
stereo_objpoints.append(obj)
|
||||
stereo_imgpoints_l.append(img_l)
|
||||
stereo_imgpoints_r.append(img_r)
|
||||
|
||||
left_objpoints.append(obj.copy())
|
||||
left_imgpoints.append(img_l.copy())
|
||||
|
||||
right_objpoints.append(obj.copy())
|
||||
right_imgpoints.append(img_r.copy())
|
||||
|
||||
accepted += 1
|
||||
|
||||
dbg_l = draw_debug(gray_l, corners_l, ids_l, f"{left_cam} corners={n_l}")
|
||||
dbg_r = draw_debug(gray_r, corners_r, ids_r, f"{right_cam} corners={n_r}")
|
||||
dbg = np.hstack([dbg_l, dbg_r])
|
||||
|
||||
cv2.putText(
|
||||
dbg,
|
||||
f"COMMON={len(common_ids)} ACCEPTED",
|
||||
(20, dbg.shape[0] - 25),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.9,
|
||||
(0, 255, 0),
|
||||
2,
|
||||
cv2.LINE_AA
|
||||
)
|
||||
|
||||
cv2.imwrite(str(debug_dir / f"pair_{idx:04d}_accepted.png"), dbg)
|
||||
|
||||
if args.show:
|
||||
cv2.imshow("debug", dbg)
|
||||
key = cv2.waitKey(1) & 0xFF
|
||||
if key in [27, ord("q"), ord("Q")]:
|
||||
break
|
||||
|
||||
print(f" [OK] left={n_l}, right={n_r}, common={len(common_ids)}")
|
||||
|
||||
except Exception as e:
|
||||
print(f" [ERRO] {e}")
|
||||
rejected += 1
|
||||
|
||||
if args.show:
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
print("")
|
||||
print(f"[INFO] aceitos: {accepted}")
|
||||
print(f"[INFO] rejeitados: {rejected}")
|
||||
|
||||
if accepted < 8:
|
||||
raise RuntimeError(
|
||||
f"Poucos pares aceitos: {accepted}. Ideal: pelo menos 15-25 bons, melhor 30+."
|
||||
)
|
||||
|
||||
print("[CALIB] Calibrando câmera esquerda...")
|
||||
ret_l, K_l, D_l, rvecs_l, tvecs_l = calibrate_single_camera(
|
||||
left_objpoints,
|
||||
left_imgpoints,
|
||||
image_size
|
||||
)
|
||||
|
||||
print("[CALIB] Calibrando câmera direita...")
|
||||
ret_r, K_r, D_r, rvecs_r, tvecs_r = calibrate_single_camera(
|
||||
right_objpoints,
|
||||
right_imgpoints,
|
||||
image_size
|
||||
)
|
||||
|
||||
print(f"[RESULT] RMS left : {ret_l:.6f}")
|
||||
print(f"[RESULT] RMS right: {ret_r:.6f}")
|
||||
|
||||
print("[CALIB] Calibração estéreo...")
|
||||
|
||||
stereo_flags = cv2.CALIB_FIX_INTRINSIC
|
||||
|
||||
ret_stereo, K_l2, D_l2, K_r2, D_r2, R, T, E, F = cv2.stereoCalibrate(
|
||||
objectPoints=stereo_objpoints,
|
||||
imagePoints1=stereo_imgpoints_l,
|
||||
imagePoints2=stereo_imgpoints_r,
|
||||
cameraMatrix1=K_l,
|
||||
distCoeffs1=D_l,
|
||||
cameraMatrix2=K_r,
|
||||
distCoeffs2=D_r,
|
||||
imageSize=image_size,
|
||||
flags=stereo_flags,
|
||||
criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, 1e-7)
|
||||
)
|
||||
|
||||
print(f"[RESULT] RMS stereo: {ret_stereo:.6f}")
|
||||
print(f"[RESULT] T: {T.ravel()}")
|
||||
|
||||
print("[CALIB] stereoRectify...")
|
||||
|
||||
R1, R2, P1, P2, Q, roi1, roi2 = cv2.stereoRectify(
|
||||
cameraMatrix1=K_l2,
|
||||
distCoeffs1=D_l2,
|
||||
cameraMatrix2=K_r2,
|
||||
distCoeffs2=D_r2,
|
||||
imageSize=image_size,
|
||||
R=R,
|
||||
T=T,
|
||||
flags=cv2.CALIB_ZERO_DISPARITY,
|
||||
alpha=0
|
||||
)
|
||||
|
||||
map1x, map1y = cv2.initUndistortRectifyMap(
|
||||
K_l2,
|
||||
D_l2,
|
||||
R1,
|
||||
P1,
|
||||
image_size,
|
||||
cv2.CV_32FC1
|
||||
)
|
||||
|
||||
map2x, map2y = cv2.initUndistortRectifyMap(
|
||||
K_r2,
|
||||
D_r2,
|
||||
R2,
|
||||
P2,
|
||||
image_size,
|
||||
cv2.CV_32FC1
|
||||
)
|
||||
|
||||
out_path = out_dir / f"stereo_calib_{left_cam}_{right_cam}.npz"
|
||||
|
||||
np.savez_compressed(
|
||||
out_path,
|
||||
left_cam=left_cam,
|
||||
right_cam=right_cam,
|
||||
image_size=np.array(image_size, dtype=np.int32),
|
||||
|
||||
squares_x=args.squares_x,
|
||||
squares_y=args.squares_y,
|
||||
square_length=args.square_length,
|
||||
marker_length=args.marker_length,
|
||||
aruco_dict=args.aruco_dict,
|
||||
|
||||
rms_left=ret_l,
|
||||
rms_right=ret_r,
|
||||
rms_stereo=ret_stereo,
|
||||
|
||||
K_left=K_l2,
|
||||
D_left=D_l2,
|
||||
K_right=K_r2,
|
||||
D_right=D_r2,
|
||||
|
||||
R=R,
|
||||
T=T,
|
||||
E=E,
|
||||
F=F,
|
||||
|
||||
R1=R1,
|
||||
R2=R2,
|
||||
P1=P1,
|
||||
P2=P2,
|
||||
Q=Q,
|
||||
roi1=np.array(roi1),
|
||||
roi2=np.array(roi2),
|
||||
|
||||
map1x=map1x,
|
||||
map1y=map1y,
|
||||
map2x=map2x,
|
||||
map2y=map2y,
|
||||
|
||||
accepted=accepted,
|
||||
rejected=rejected
|
||||
)
|
||||
|
||||
print("")
|
||||
print(f"[OK] calibração salva em: {out_path}")
|
||||
print(f"[OK] debug salvo em: {debug_dir}")
|
||||
print("")
|
||||
print("Resumo:")
|
||||
print(f" RMS left = {ret_l:.6f}")
|
||||
print(f" RMS right = {ret_r:.6f}")
|
||||
print(f" RMS stereo = {ret_stereo:.6f}")
|
||||
print(f" T = {T.ravel()}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,705 @@
|
|||
import argparse
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
# ============================================================
|
||||
# RAW10 / image conversion
|
||||
# ============================================================
|
||||
|
||||
def unpack_raw10_packed(raw: bytes, width: int, height: int) -> np.ndarray:
|
||||
arr = np.frombuffer(raw, dtype=np.uint8)
|
||||
|
||||
pixel_count = width * height
|
||||
expected_bytes = (pixel_count // 4) * 5
|
||||
|
||||
if pixel_count % 4 != 0:
|
||||
raise RuntimeError(f"width*height precisa ser múltiplo de 4. Recebido: {pixel_count}")
|
||||
|
||||
if arr.size < expected_bytes:
|
||||
raise RuntimeError(
|
||||
f"RAW10 menor que esperado. bytes={arr.size}, esperado={expected_bytes}, "
|
||||
f"width={width}, height={height}"
|
||||
)
|
||||
|
||||
arr = arr[:expected_bytes]
|
||||
groups = arr.reshape(-1, 5).astype(np.uint16)
|
||||
|
||||
p0 = (groups[:, 0] << 2) | ((groups[:, 4] >> 0) & 0x03)
|
||||
p1 = (groups[:, 1] << 2) | ((groups[:, 4] >> 2) & 0x03)
|
||||
p2 = (groups[:, 2] << 2) | ((groups[:, 4] >> 4) & 0x03)
|
||||
p3 = (groups[:, 3] << 2) | ((groups[:, 4] >> 6) & 0x03)
|
||||
|
||||
out = np.empty(groups.shape[0] * 4, dtype=np.uint16)
|
||||
out[0::4] = p0
|
||||
out[1::4] = p1
|
||||
out[2::4] = p2
|
||||
out[3::4] = p3
|
||||
|
||||
return out.reshape(height, width)
|
||||
|
||||
|
||||
def normalize_to_u8(img: np.ndarray, p_low=1.0, p_high=99.0) -> np.ndarray:
|
||||
arr = img.astype(np.float32)
|
||||
valid = np.isfinite(arr)
|
||||
|
||||
if np.count_nonzero(valid) < 20:
|
||||
return np.zeros(arr.shape[:2], dtype=np.uint8)
|
||||
|
||||
vals = arr[valid]
|
||||
lo = np.percentile(vals, p_low)
|
||||
hi = np.percentile(vals, p_high)
|
||||
|
||||
out = (arr - lo) / (hi - lo + 1e-6)
|
||||
out = np.clip(out, 0.0, 1.0)
|
||||
return (out * 255).astype(np.uint8)
|
||||
|
||||
|
||||
def debayer_raw10_to_bgr_u8(raw10: np.ndarray, bayer: str) -> np.ndarray:
|
||||
gray_u8 = normalize_to_u8(raw10)
|
||||
bayer = bayer.upper()
|
||||
|
||||
code_map = {
|
||||
"RGGB": cv2.COLOR_BayerRG2BGR,
|
||||
"BGGR": cv2.COLOR_BayerBG2BGR,
|
||||
"GRBG": cv2.COLOR_BayerGR2BGR,
|
||||
"GBRG": cv2.COLOR_BayerGB2BGR,
|
||||
}
|
||||
|
||||
if bayer not in code_map:
|
||||
raise RuntimeError(f"Bayer pattern não suportado: {bayer}")
|
||||
|
||||
return cv2.cvtColor(gray_u8, code_map[bayer])
|
||||
|
||||
|
||||
def read_rgb(path: Path, width: int, height: int, bayer: str, rgb_view: str, use_clahe: bool):
|
||||
raw = path.read_bytes()
|
||||
raw10 = unpack_raw10_packed(raw, width, height)
|
||||
mode = rgb_view.lower().strip()
|
||||
|
||||
if mode == "color":
|
||||
bgr = debayer_raw10_to_bgr_u8(raw10, bayer)
|
||||
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
|
||||
return bgr, gray
|
||||
|
||||
if mode == "gray":
|
||||
bgr_color = debayer_raw10_to_bgr_u8(raw10, bayer)
|
||||
gray = cv2.cvtColor(bgr_color, cv2.COLOR_BGR2GRAY)
|
||||
elif mode == "raw_bayer_gray":
|
||||
gray = normalize_to_u8(raw10)
|
||||
else:
|
||||
raise RuntimeError(f"rgb_view inválido: {rgb_view}")
|
||||
|
||||
if use_clahe:
|
||||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
gray = clahe.apply(gray)
|
||||
|
||||
bgr = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
|
||||
return bgr, gray
|
||||
|
||||
|
||||
def read_mono(path: Path, width: int, height: int, use_clahe: bool):
|
||||
raw = path.read_bytes()
|
||||
raw10 = unpack_raw10_packed(raw, width, height)
|
||||
gray = normalize_to_u8(raw10)
|
||||
|
||||
if use_clahe:
|
||||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
gray = clahe.apply(gray)
|
||||
|
||||
bgr = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
|
||||
return bgr, gray
|
||||
|
||||
|
||||
def to_gray_u8(img_bgr: np.ndarray):
|
||||
if img_bgr.ndim == 2:
|
||||
return img_bgr
|
||||
return cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Triplet pairing
|
||||
# ============================================================
|
||||
|
||||
def clean_stem_for_pair(path: Path, cam_key: str):
|
||||
s = path.stem
|
||||
variants = [
|
||||
cam_key,
|
||||
cam_key.lower(),
|
||||
cam_key.replace("_", ""),
|
||||
cam_key.replace("_", "").lower(),
|
||||
]
|
||||
|
||||
for v in variants:
|
||||
s = s.replace(v, "")
|
||||
|
||||
s = re.sub(r"[_\-\s]+", "_", s).strip("_").lower()
|
||||
return s
|
||||
|
||||
|
||||
def find_cam_bins(root_dir: Path, cam_key: str):
|
||||
return sorted([p for p in root_dir.rglob("*.bin") if cam_key.lower() in p.name.lower()])
|
||||
|
||||
|
||||
def find_triplets(root_dir: Path, cams: list[str]):
|
||||
by_cam = {cam: find_cam_bins(root_dir, cam) for cam in cams}
|
||||
key_maps = {}
|
||||
|
||||
for cam, paths in by_cam.items():
|
||||
m = {}
|
||||
for p in paths:
|
||||
key = clean_stem_for_pair(p, cam)
|
||||
m[(p.parent, key)] = p
|
||||
m.setdefault((None, key), p)
|
||||
key_maps[cam] = m
|
||||
|
||||
ref_cam = cams[0]
|
||||
triplets = []
|
||||
|
||||
for ref_path in by_cam[ref_cam]:
|
||||
key = clean_stem_for_pair(ref_path, ref_cam)
|
||||
folder = ref_path.parent
|
||||
item = {ref_cam: ref_path}
|
||||
ok = True
|
||||
|
||||
for cam in cams[1:]:
|
||||
p = key_maps[cam].get((folder, key)) or key_maps[cam].get((None, key))
|
||||
if p is None:
|
||||
same_folder = [x for x in by_cam[cam] if x.parent == folder]
|
||||
if len(same_folder) == 1:
|
||||
p = same_folder[0]
|
||||
|
||||
if p is None:
|
||||
ok = False
|
||||
break
|
||||
|
||||
item[cam] = p
|
||||
|
||||
if ok:
|
||||
triplets.append(item)
|
||||
|
||||
return triplets, by_cam
|
||||
|
||||
|
||||
# ============================================================
|
||||
# ChArUco helpers
|
||||
# ============================================================
|
||||
|
||||
def get_aruco_dict(dict_name: str):
|
||||
aruco = cv2.aruco
|
||||
mapping = {
|
||||
"4X4_50": aruco.DICT_4X4_50,
|
||||
"4X4_100": aruco.DICT_4X4_100,
|
||||
"4X4_250": aruco.DICT_4X4_250,
|
||||
"4X4_1000": aruco.DICT_4X4_1000,
|
||||
"5X5_50": aruco.DICT_5X5_50,
|
||||
"5X5_100": aruco.DICT_5X5_100,
|
||||
"5X5_250": aruco.DICT_5X5_250,
|
||||
"5X5_1000": aruco.DICT_5X5_1000,
|
||||
"6X6_50": aruco.DICT_6X6_50,
|
||||
"6X6_100": aruco.DICT_6X6_100,
|
||||
"6X6_250": aruco.DICT_6X6_250,
|
||||
"6X6_1000": aruco.DICT_6X6_1000,
|
||||
}
|
||||
|
||||
key = dict_name.upper()
|
||||
if key not in mapping:
|
||||
raise RuntimeError(f"Dicionário ArUco não suportado: {dict_name}")
|
||||
|
||||
if hasattr(aruco, "getPredefinedDictionary"):
|
||||
return aruco.getPredefinedDictionary(mapping[key])
|
||||
return aruco.Dictionary_get(mapping[key])
|
||||
|
||||
|
||||
def create_charuco_board(squares_x, squares_y, square_length, marker_length, aruco_dict):
|
||||
aruco = cv2.aruco
|
||||
|
||||
if hasattr(aruco, "CharucoBoard"):
|
||||
try:
|
||||
return aruco.CharucoBoard((squares_x, squares_y), square_length, marker_length, aruco_dict)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if hasattr(aruco, "CharucoBoard_create"):
|
||||
return aruco.CharucoBoard_create(squares_x, squares_y, square_length, marker_length, aruco_dict)
|
||||
|
||||
raise RuntimeError("Sua versão do OpenCV não tem CharucoBoard/CharucoBoard_create.")
|
||||
|
||||
|
||||
def create_detector_params():
|
||||
aruco = cv2.aruco
|
||||
if hasattr(aruco, "DetectorParameters"):
|
||||
return aruco.DetectorParameters()
|
||||
if hasattr(aruco, "DetectorParameters_create"):
|
||||
return aruco.DetectorParameters_create()
|
||||
return None
|
||||
|
||||
|
||||
def detect_charuco(gray: np.ndarray, board, aruco_dict, min_corners: int):
|
||||
aruco = cv2.aruco
|
||||
params = create_detector_params()
|
||||
|
||||
if hasattr(aruco, "CharucoDetector"):
|
||||
try:
|
||||
detector = aruco.CharucoDetector(board)
|
||||
charuco_corners, charuco_ids, marker_corners, marker_ids = detector.detectBoard(gray)
|
||||
if charuco_corners is None or charuco_ids is None:
|
||||
return None, None
|
||||
corners = np.array(charuco_corners, dtype=np.float32).reshape(-1, 2)
|
||||
ids = np.array(charuco_ids, dtype=np.int32).reshape(-1)
|
||||
if len(ids) < min_corners:
|
||||
return None, None
|
||||
return corners, ids
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if params is not None:
|
||||
marker_corners, marker_ids, rejected = aruco.detectMarkers(gray, aruco_dict, parameters=params)
|
||||
else:
|
||||
marker_corners, marker_ids, rejected = aruco.detectMarkers(gray, aruco_dict)
|
||||
|
||||
if marker_ids is None or len(marker_ids) == 0:
|
||||
return None, None
|
||||
|
||||
try:
|
||||
aruco.refineDetectedMarkers(gray, board, marker_corners, marker_ids, rejected)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
retval, charuco_corners, charuco_ids = aruco.interpolateCornersCharuco(marker_corners, marker_ids, gray, board)
|
||||
|
||||
if charuco_corners is None or charuco_ids is None:
|
||||
return None, None
|
||||
|
||||
corners = np.array(charuco_corners, dtype=np.float32).reshape(-1, 2)
|
||||
ids = np.array(charuco_ids, dtype=np.int32).reshape(-1)
|
||||
|
||||
if len(ids) < min_corners:
|
||||
return None, None
|
||||
|
||||
return corners, ids
|
||||
|
||||
|
||||
def common_points_2cam(corners_a, ids_a, corners_b, ids_b, min_common: int):
|
||||
map_a = {int(i): corners_a[k] for k, i in enumerate(ids_a)}
|
||||
map_b = {int(i): corners_b[k] for k, i in enumerate(ids_b)}
|
||||
|
||||
common_ids = sorted(set(map_a.keys()) & set(map_b.keys()))
|
||||
if len(common_ids) < min_common:
|
||||
return None, None, common_ids
|
||||
|
||||
pts_a = np.array([map_a[i] for i in common_ids], dtype=np.float32).reshape(-1, 2)
|
||||
pts_b = np.array([map_b[i] for i in common_ids], dtype=np.float32).reshape(-1, 2)
|
||||
|
||||
return pts_a, pts_b, common_ids
|
||||
|
||||
|
||||
def compute_homography_to_rgb(rgb_det, other_det, min_common: int, ransac_thresh: float):
|
||||
rgb_corners, rgb_ids = rgb_det
|
||||
other_corners, other_ids = other_det
|
||||
|
||||
if rgb_corners is None or rgb_ids is None or other_corners is None or other_ids is None:
|
||||
return None, None, []
|
||||
|
||||
pts_other, pts_rgb, common_ids = common_points_2cam(
|
||||
other_corners,
|
||||
other_ids,
|
||||
rgb_corners,
|
||||
rgb_ids,
|
||||
min_common=min_common,
|
||||
)
|
||||
|
||||
if pts_other is None:
|
||||
return None, None, common_ids
|
||||
|
||||
H, mask = cv2.findHomography(pts_other, pts_rgb, cv2.RANSAC, ransac_thresh)
|
||||
if H is None:
|
||||
return None, None, common_ids
|
||||
|
||||
inliers = int(np.count_nonzero(mask)) if mask is not None else 0
|
||||
return H, inliers, common_ids
|
||||
|
||||
|
||||
def draw_charuco_debug(img_bgr, corners, ids, label):
|
||||
out = img_bgr.copy()
|
||||
if corners is not None and ids is not None:
|
||||
corners_draw = np.array(corners, dtype=np.float32).reshape(-1, 1, 2)
|
||||
ids_draw = np.array(ids, dtype=np.int32).reshape(-1, 1)
|
||||
try:
|
||||
cv2.aruco.drawDetectedCornersCharuco(out, corners_draw, ids_draw, (0, 255, 0))
|
||||
except Exception:
|
||||
for p in corners:
|
||||
cv2.circle(out, tuple(np.round(p).astype(int)), 3, (0, 255, 0), -1)
|
||||
n = 0 if ids is None else len(ids)
|
||||
cv2.putText(out, f"{label} corners={n}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 1, cv2.LINE_AA)
|
||||
return out
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Visualization
|
||||
# ============================================================
|
||||
|
||||
def draw_lines(img_bgr, step=40):
|
||||
out = img_bgr.copy()
|
||||
h, w = out.shape[:2]
|
||||
for y in range(0, h, step):
|
||||
color = (0, 255, 255) if (y // step) % 2 == 0 else (255, 255, 0)
|
||||
cv2.line(out, (0, y), (w, y), color, 1, cv2.LINE_AA)
|
||||
return out
|
||||
|
||||
|
||||
def resize_to_height(img, target_h):
|
||||
h, w = img.shape[:2]
|
||||
if h == target_h:
|
||||
return img
|
||||
scale = target_h / h
|
||||
new_w = max(1, int(w * scale))
|
||||
return cv2.resize(img, (new_w, target_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
|
||||
def put_label(img, text, color=(255, 255, 255)):
|
||||
out = img.copy()
|
||||
cv2.rectangle(out, (0, 0), (out.shape[1], 34), (0, 0, 0), -1)
|
||||
cv2.putText(out, text, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.65, color, 1, cv2.LINE_AA)
|
||||
return out
|
||||
|
||||
|
||||
def draw_header(canvas, lines):
|
||||
header_h = 24 + 24 * len(lines)
|
||||
cv2.rectangle(canvas, (0, 0), (canvas.shape[1], header_h), (0, 0, 0), -1)
|
||||
y = 24
|
||||
for line in lines:
|
||||
cv2.putText(canvas, line, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA)
|
||||
y += 24
|
||||
return canvas
|
||||
|
||||
|
||||
def compose_three(left, center, right, lines, args):
|
||||
if args.lines:
|
||||
left = draw_lines(left, args.line_step)
|
||||
center = draw_lines(center, args.line_step)
|
||||
right = draw_lines(right, args.line_step)
|
||||
|
||||
imgs = [resize_to_height(x, args.view_h) for x in [left, center, right]]
|
||||
h = min(x.shape[0] for x in imgs)
|
||||
imgs = [x[:h] for x in imgs]
|
||||
canvas = np.hstack(imgs)
|
||||
return draw_header(canvas, lines)
|
||||
|
||||
|
||||
def absdiff_bgr(a, b):
|
||||
ag = to_gray_u8(a)
|
||||
bg = to_gray_u8(b)
|
||||
diff = cv2.absdiff(ag, bg)
|
||||
return cv2.cvtColor(diff, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
|
||||
def overlay_bgr(base, layer, alpha=0.45):
|
||||
layer = cv2.resize(layer, (base.shape[1], base.shape[0]), interpolation=cv2.INTER_AREA)
|
||||
return cv2.addWeighted(base, 1.0 - alpha, layer, alpha, 0)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Main processing
|
||||
# ============================================================
|
||||
|
||||
def load_triplet(item, args):
|
||||
rgb_bgr, rgb_gray = read_rgb(
|
||||
item[args.rgb_cam],
|
||||
args.width,
|
||||
args.height,
|
||||
args.rgb_bayer,
|
||||
args.rgb_view,
|
||||
use_clahe=not args.no_clahe,
|
||||
)
|
||||
re_bgr, re_gray = read_mono(item[args.re_cam], args.width, args.height, use_clahe=not args.no_clahe)
|
||||
nir_bgr, nir_gray = read_mono(item[args.nir_cam], args.width, args.height, use_clahe=not args.no_clahe)
|
||||
|
||||
return {
|
||||
"rgb_bgr": rgb_bgr,
|
||||
"rgb_gray": rgb_gray,
|
||||
"re_bgr": re_bgr,
|
||||
"re_gray": re_gray,
|
||||
"nir_bgr": nir_bgr,
|
||||
"nir_gray": nir_gray,
|
||||
}
|
||||
|
||||
|
||||
def build_views(item, board, aruco_dict, args):
|
||||
frames = load_triplet(item, args)
|
||||
|
||||
rgb_det = detect_charuco(frames["rgb_gray"], board, aruco_dict, args.min_corners)
|
||||
re_det = detect_charuco(frames["re_gray"], board, aruco_dict, args.min_corners)
|
||||
nir_det = detect_charuco(frames["nir_gray"], board, aruco_dict, args.min_corners)
|
||||
|
||||
H_re, inliers_re, common_re = compute_homography_to_rgb(
|
||||
rgb_det,
|
||||
re_det,
|
||||
min_common=args.min_common,
|
||||
ransac_thresh=args.ransac_thresh,
|
||||
)
|
||||
|
||||
H_nir, inliers_nir, common_nir = compute_homography_to_rgb(
|
||||
rgb_det,
|
||||
nir_det,
|
||||
min_common=args.min_common,
|
||||
ransac_thresh=args.ransac_thresh,
|
||||
)
|
||||
|
||||
h, w = frames["rgb_bgr"].shape[:2]
|
||||
|
||||
if H_re is not None:
|
||||
re_to_rgb = cv2.warpPerspective(
|
||||
frames["re_bgr"],
|
||||
H_re,
|
||||
(w, h),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
else:
|
||||
re_to_rgb = np.zeros_like(frames["rgb_bgr"])
|
||||
|
||||
if H_nir is not None:
|
||||
nir_to_rgb = cv2.warpPerspective(
|
||||
frames["nir_bgr"],
|
||||
H_nir,
|
||||
(w, h),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
else:
|
||||
nir_to_rgb = np.zeros_like(frames["rgb_bgr"])
|
||||
|
||||
stats = {
|
||||
"rgb_corners": 0 if rgb_det[1] is None else len(rgb_det[1]),
|
||||
"re_corners": 0 if re_det[1] is None else len(re_det[1]),
|
||||
"nir_corners": 0 if nir_det[1] is None else len(nir_det[1]),
|
||||
"common_re": len(common_re),
|
||||
"common_nir": len(common_nir),
|
||||
"inliers_re": 0 if inliers_re is None else int(inliers_re),
|
||||
"inliers_nir": 0 if inliers_nir is None else int(inliers_nir),
|
||||
"ok_re": H_re is not None,
|
||||
"ok_nir": H_nir is not None,
|
||||
}
|
||||
|
||||
return {
|
||||
**frames,
|
||||
"rgb_det": rgb_det,
|
||||
"re_det": re_det,
|
||||
"nir_det": nir_det,
|
||||
"re_to_rgb": re_to_rgb,
|
||||
"nir_to_rgb": nir_to_rgb,
|
||||
"diff_re": absdiff_bgr(frames["rgb_bgr"], re_to_rgb),
|
||||
"diff_nir": absdiff_bgr(frames["rgb_bgr"], nir_to_rgb),
|
||||
"overlay_re": overlay_bgr(frames["rgb_bgr"], re_to_rgb, args.overlay_alpha),
|
||||
"overlay_nir": overlay_bgr(frames["rgb_bgr"], nir_to_rgb, args.overlay_alpha),
|
||||
"stats": stats,
|
||||
}
|
||||
|
||||
|
||||
def compose_mode(views, item, idx, total, mode, args):
|
||||
stats = views["stats"]
|
||||
name = item[args.rgb_cam].name
|
||||
|
||||
status = (
|
||||
f"RGB={stats['rgb_corners']} RE={stats['re_corners']} NIR={stats['nir_corners']} | "
|
||||
f"RE common/inliers={stats['common_re']}/{stats['inliers_re']} | "
|
||||
f"NIR common/inliers={stats['common_nir']}/{stats['inliers_nir']}"
|
||||
)
|
||||
|
||||
if mode == "native":
|
||||
left = put_label(views["re_bgr"], "RE native")
|
||||
center = put_label(views["rgb_bgr"], "RGB REF native")
|
||||
right = put_label(views["nir_bgr"], "NIR native")
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | {name}",
|
||||
f"mode=native | {status}",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | S salvar | Q sair",
|
||||
]
|
||||
return compose_three(left, center, right, lines, args)
|
||||
|
||||
if mode == "rgb_ref":
|
||||
left = put_label(views["re_to_rgb"], "RE -> RGB plane")
|
||||
center = put_label(views["rgb_bgr"], "RGB REF")
|
||||
right = put_label(views["nir_to_rgb"], "NIR -> RGB plane")
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | {name}",
|
||||
f"mode=rgb_ref planar homography | {status}",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | S salvar | Q sair",
|
||||
]
|
||||
return compose_three(left, center, right, lines, args)
|
||||
|
||||
if mode == "diff":
|
||||
left = put_label(views["diff_re"], "diff RGB vs RE_to_RGB")
|
||||
center = put_label(views["rgb_bgr"], "RGB REF")
|
||||
right = put_label(views["diff_nir"], "diff RGB vs NIR_to_RGB")
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | {name}",
|
||||
f"mode=diff | {status}",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | S salvar | Q sair",
|
||||
]
|
||||
return compose_three(left, center, right, lines, args)
|
||||
|
||||
if mode == "overlay":
|
||||
left = put_label(views["overlay_re"], "overlay RGB + RE_to_RGB")
|
||||
center = put_label(views["rgb_bgr"], "RGB REF")
|
||||
right = put_label(views["overlay_nir"], "overlay RGB + NIR_to_RGB")
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | {name}",
|
||||
f"mode=overlay alpha={args.overlay_alpha:.2f} | {status}",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | S salvar | Q sair",
|
||||
]
|
||||
return compose_three(left, center, right, lines, args)
|
||||
|
||||
if mode == "debug_corners":
|
||||
left = draw_charuco_debug(views["re_bgr"], views["re_det"][0], views["re_det"][1], "RE native")
|
||||
center = draw_charuco_debug(views["rgb_bgr"], views["rgb_det"][0], views["rgb_det"][1], "RGB native")
|
||||
right = draw_charuco_debug(views["nir_bgr"], views["nir_det"][0], views["nir_det"][1], "NIR native")
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | {name}",
|
||||
f"mode=debug_corners | {status}",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | S salvar | Q sair",
|
||||
]
|
||||
return compose_three(left, center, right, lines, args)
|
||||
|
||||
raise RuntimeError(f"Modo desconhecido: {mode}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument("--root_dir", required=True)
|
||||
parser.add_argument("--rgb_cam", default="CAM_A")
|
||||
parser.add_argument("--re_cam", default="CAM_B")
|
||||
parser.add_argument("--nir_cam", default="CAM_C")
|
||||
|
||||
parser.add_argument("--width", type=int, default=1280)
|
||||
parser.add_argument("--height", type=int, default=800)
|
||||
parser.add_argument("--rgb_bayer", default="BGGR")
|
||||
parser.add_argument(
|
||||
"--rgb_view",
|
||||
default="gray",
|
||||
choices=["color", "gray", "raw_bayer_gray"],
|
||||
help="Como processar CAM_A/RGB no viewer.",
|
||||
)
|
||||
|
||||
parser.add_argument("--squares_x", type=int, default=13)
|
||||
parser.add_argument("--squares_y", type=int, default=7)
|
||||
parser.add_argument("--square_length", type=float, default=0.031)
|
||||
parser.add_argument("--marker_length", type=float, default=0.023)
|
||||
parser.add_argument("--aruco_dict", default="4X4_50")
|
||||
|
||||
parser.add_argument("--min_corners", type=int, default=30)
|
||||
parser.add_argument("--min_common", type=int, default=20)
|
||||
parser.add_argument("--ransac_thresh", type=float, default=3.0)
|
||||
|
||||
parser.add_argument("--view_h", type=int, default=420)
|
||||
parser.add_argument("--no_clahe", action="store_true")
|
||||
parser.add_argument("--lines", action="store_true")
|
||||
parser.add_argument("--line_step", type=int, default=40)
|
||||
parser.add_argument("--overlay_alpha", type=float, default=0.45)
|
||||
|
||||
parser.add_argument("--save_dir", default="rgb_reference_homography_saves")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
root_dir = Path(args.root_dir)
|
||||
save_dir = Path(args.save_dir)
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
cams = [args.rgb_cam, args.re_cam, args.nir_cam]
|
||||
triplets, by_cam = find_triplets(root_dir, cams)
|
||||
|
||||
print(f"[INFO] root_dir={root_dir}")
|
||||
print(f"[INFO] cams RGB={args.rgb_cam} RE={args.re_cam} NIR={args.nir_cam}")
|
||||
print(f"[INFO] rgb_bayer={args.rgb_bayer} rgb_view={args.rgb_view}")
|
||||
print(f"[INFO] ChArUco squares={args.squares_x}x{args.squares_y}")
|
||||
for cam in cams:
|
||||
print(f"[INFO] arquivos {cam}: {len(by_cam[cam])}")
|
||||
print(f"[INFO] triplets: {len(triplets)}")
|
||||
|
||||
if not triplets:
|
||||
raise RuntimeError("Nenhum triplet encontrado.")
|
||||
|
||||
aruco_dict = get_aruco_dict(args.aruco_dict)
|
||||
board = create_charuco_board(args.squares_x, args.squares_y, args.square_length, args.marker_length, aruco_dict)
|
||||
|
||||
modes = ["native", "debug_corners", "rgb_ref", "overlay", "diff"]
|
||||
mode_idx = 0
|
||||
idx = 0
|
||||
|
||||
cached_key = None
|
||||
cached_views = None
|
||||
|
||||
cv2.namedWindow("RGB Reference Homography Viewer", cv2.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
item = triplets[idx]
|
||||
mode = modes[mode_idx]
|
||||
|
||||
key_cache = (
|
||||
tuple(str(item[cam]) for cam in cams),
|
||||
args.width,
|
||||
args.height,
|
||||
args.rgb_bayer,
|
||||
args.rgb_view,
|
||||
args.no_clahe,
|
||||
args.min_corners,
|
||||
args.min_common,
|
||||
args.ransac_thresh,
|
||||
args.overlay_alpha,
|
||||
)
|
||||
|
||||
if key_cache != cached_key:
|
||||
print(f"[RUN] {idx + 1}/{len(triplets)} - {item[args.rgb_cam].name}")
|
||||
try:
|
||||
cached_views = build_views(item, board, aruco_dict, args)
|
||||
cached_key = key_cache
|
||||
except Exception as e:
|
||||
print("[ERRO] Falha processando triplet:")
|
||||
for cam in cams:
|
||||
print(f" {cam}={item[cam]}")
|
||||
print(f" erro={e}")
|
||||
idx = min(idx + 1, len(triplets) - 1)
|
||||
cached_key = None
|
||||
cached_views = None
|
||||
continue
|
||||
|
||||
view = compose_mode(cached_views, item, idx, len(triplets), mode, args)
|
||||
cv2.imshow("RGB Reference Homography Viewer", view)
|
||||
key = cv2.waitKey(0) & 0xFF
|
||||
|
||||
if key in [27, ord("q"), ord("Q")]:
|
||||
break
|
||||
|
||||
elif key in [ord("n"), ord("N"), 32]:
|
||||
idx = min(idx + 1, len(triplets) - 1)
|
||||
cached_key = None
|
||||
|
||||
elif key in [ord("a"), ord("A")]:
|
||||
idx = max(idx - 1, 0)
|
||||
cached_key = None
|
||||
|
||||
elif key in [ord("m"), ord("M")]:
|
||||
mode_idx = (mode_idx + 1) % len(modes)
|
||||
print(f"[PARAM] mode={modes[mode_idx]}")
|
||||
|
||||
elif key in [ord("l"), ord("L")]:
|
||||
args.lines = not args.lines
|
||||
print(f"[PARAM] lines={args.lines}")
|
||||
|
||||
elif key in [ord("s"), ord("S")]:
|
||||
out_path = save_dir / f"rgb_ref_homography_{idx:04d}_{mode}.png"
|
||||
cv2.imwrite(str(out_path), view)
|
||||
print(f"[SAVE] {out_path}")
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,653 @@
|
|||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from core.oak_fcc3_client import OakFcc3Client as MultiSpectralClient
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Config padrão do projeto
|
||||
# ============================================================
|
||||
|
||||
with open("config.json", "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
RAW_SIZE = config.get("raw_size", [1280, 800])
|
||||
MODULE_PARAMS = config.get("module_params_json")
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Visual helpers
|
||||
# ============================================================
|
||||
|
||||
def normalize_to_u8(img: np.ndarray, p_low=1.0, p_high=99.0) -> np.ndarray:
|
||||
arr = img.astype(np.float32)
|
||||
valid = np.isfinite(arr)
|
||||
|
||||
if np.count_nonzero(valid) < 20:
|
||||
return np.zeros(arr.shape[:2], dtype=np.uint8)
|
||||
|
||||
vals = arr[valid]
|
||||
lo = np.percentile(vals, p_low)
|
||||
hi = np.percentile(vals, p_high)
|
||||
|
||||
out = (arr - lo) / (hi - lo + 1e-6)
|
||||
out = np.clip(out, 0.0, 1.0)
|
||||
return (out * 255).astype(np.uint8)
|
||||
|
||||
|
||||
def ensure_bgr_u8(img) -> np.ndarray:
|
||||
if img is None:
|
||||
return None
|
||||
|
||||
arr = np.asarray(img)
|
||||
|
||||
if arr.ndim == 2:
|
||||
if arr.dtype != np.uint8:
|
||||
arr = normalize_to_u8(arr)
|
||||
return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
if arr.ndim == 3 and arr.shape[2] == 3:
|
||||
if arr.dtype == np.uint8:
|
||||
return arr.copy()
|
||||
arr = np.clip(arr.astype(np.float32), 0.0, 1.0)
|
||||
return (arr * 255.0).astype(np.uint8)
|
||||
|
||||
raise RuntimeError(f"Imagem inválida para BGR: shape={arr.shape}, dtype={arr.dtype}")
|
||||
|
||||
|
||||
def to_gray_u8(img_bgr: np.ndarray) -> np.ndarray:
|
||||
if img_bgr.ndim == 2:
|
||||
return img_bgr if img_bgr.dtype == np.uint8 else normalize_to_u8(img_bgr)
|
||||
return cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
|
||||
def resize_to_fit(img: np.ndarray, w: int, h: int) -> np.ndarray:
|
||||
return cv2.resize(img, (w, h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
|
||||
def put_label(img: np.ndarray, text: str, color=(255, 255, 255)) -> np.ndarray:
|
||||
out = img.copy()
|
||||
cv2.rectangle(out, (0, 0), (out.shape[1], 34), (0, 0, 0), -1)
|
||||
cv2.putText(out, text, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.65, color, 1, cv2.LINE_AA)
|
||||
return out
|
||||
|
||||
|
||||
def overlay_hud(img_bgr: np.ndarray, lines: list[str]):
|
||||
y = 24
|
||||
for s in lines:
|
||||
cv2.putText(img_bgr, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 3, cv2.LINE_AA)
|
||||
cv2.putText(img_bgr, s, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA)
|
||||
y += 24
|
||||
|
||||
|
||||
def colorize_scalar_0_1(x: np.ndarray, cmap=cv2.COLORMAP_TURBO) -> np.ndarray:
|
||||
x8 = np.clip(x * 255.0, 0, 255).astype(np.uint8)
|
||||
return cv2.applyColorMap(x8, cmap)
|
||||
|
||||
|
||||
def make_2x2(a, b, c, d, cell_w=640, cell_h=400, hud_lines=None) -> np.ndarray:
|
||||
a = resize_to_fit(a, cell_w, cell_h)
|
||||
b = resize_to_fit(b, cell_w, cell_h)
|
||||
c = resize_to_fit(c, cell_w, cell_h)
|
||||
d = resize_to_fit(d, cell_w, cell_h)
|
||||
|
||||
top = np.hstack([a, b])
|
||||
bot = np.hstack([c, d])
|
||||
canvas = np.vstack([top, bot])
|
||||
|
||||
if hud_lines:
|
||||
overlay_hud(canvas, hud_lines)
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def make_3x2(a, b, c, d, e, f, cell_w=520, cell_h=320, hud_lines=None) -> np.ndarray:
|
||||
imgs = [resize_to_fit(x, cell_w, cell_h) for x in [a, b, c, d, e, f]]
|
||||
row1 = np.hstack(imgs[:3])
|
||||
row2 = np.hstack(imgs[3:])
|
||||
canvas = np.vstack([row1, row2])
|
||||
|
||||
if hud_lines:
|
||||
overlay_hud(canvas, hud_lines)
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Calibração product bundle
|
||||
# ============================================================
|
||||
|
||||
def scalar_str(x):
|
||||
arr = np.array(x)
|
||||
if arr.shape == ():
|
||||
return str(arr.item())
|
||||
return str(x)
|
||||
|
||||
|
||||
def load_calibration_bundle(path: str | Path) -> dict:
|
||||
data = np.load(str(path), allow_pickle=True)
|
||||
keys = set(data.files)
|
||||
|
||||
required = [
|
||||
"rgb_cam",
|
||||
"re_cam",
|
||||
"nir_cam",
|
||||
"image_size",
|
||||
"H_RE_to_RGB",
|
||||
"H_NIR_to_RGB",
|
||||
]
|
||||
|
||||
for k in required:
|
||||
if k not in keys:
|
||||
raise RuntimeError(f"Calibração sem chave obrigatória: {k}")
|
||||
|
||||
calib = {
|
||||
"data": data,
|
||||
"keys": keys,
|
||||
"rgb_cam": scalar_str(data["rgb_cam"]),
|
||||
"re_cam": scalar_str(data["re_cam"]),
|
||||
"nir_cam": scalar_str(data["nir_cam"]),
|
||||
"image_size": tuple(data["image_size"].astype(int).tolist()),
|
||||
"H_RE_to_RGB": data["H_RE_to_RGB"].astype(np.float64),
|
||||
"H_NIR_to_RGB": data["H_NIR_to_RGB"].astype(np.float64),
|
||||
"overlap_RE_to_RGB": data["overlap_RE_to_RGB"] if "overlap_RE_to_RGB" in keys else None,
|
||||
"overlap_NIR_to_RGB": data["overlap_NIR_to_RGB"] if "overlap_NIR_to_RGB" in keys else None,
|
||||
"overlap_common_RGB": data["overlap_common_RGB"] if "overlap_common_RGB" in keys else None,
|
||||
}
|
||||
|
||||
return calib
|
||||
|
||||
|
||||
def get_pair_prefix(calib: dict, cam1: str, cam2: str):
|
||||
keys = calib["keys"]
|
||||
direct = f"pair_{cam1}_{cam2}"
|
||||
inv = f"pair_{cam2}_{cam1}"
|
||||
|
||||
if f"{direct}_map1x" in keys:
|
||||
return direct, False
|
||||
|
||||
if f"{inv}_map1x" in keys:
|
||||
return inv, True
|
||||
|
||||
raise RuntimeError(f"Par estéreo {cam1}<->{cam2} não encontrado no .npz")
|
||||
|
||||
|
||||
def rectify_pair_gray(gray1, gray2, calib: dict, cam1: str, cam2: str):
|
||||
data = calib["data"]
|
||||
image_w, image_h = calib["image_size"]
|
||||
|
||||
if gray1.shape[::-1] != (image_w, image_h):
|
||||
gray1 = cv2.resize(gray1, (image_w, image_h), interpolation=cv2.INTER_AREA)
|
||||
if gray2.shape[::-1] != (image_w, image_h):
|
||||
gray2 = cv2.resize(gray2, (image_w, image_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
prefix, inverted = get_pair_prefix(calib, cam1, cam2)
|
||||
|
||||
if not inverted:
|
||||
map1x = data[f"{prefix}_map1x"]
|
||||
map1y = data[f"{prefix}_map1y"]
|
||||
map2x = data[f"{prefix}_map2x"]
|
||||
map2y = data[f"{prefix}_map2y"]
|
||||
else:
|
||||
map1x = data[f"{prefix}_map2x"]
|
||||
map1y = data[f"{prefix}_map2y"]
|
||||
map2x = data[f"{prefix}_map1x"]
|
||||
map2y = data[f"{prefix}_map1y"]
|
||||
|
||||
rect1 = cv2.remap(gray1, map1x, map1y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
|
||||
rect2 = cv2.remap(gray2, map2x, map2y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
|
||||
return rect1, rect2
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Frame extraction from OakFcc3Client
|
||||
# ============================================================
|
||||
|
||||
def role_map_from_meta(meta: dict) -> dict:
|
||||
camera_info = meta.get("camera_info", {}) or {}
|
||||
out = {}
|
||||
for cam_id, info in camera_info.items():
|
||||
role = info.get("role") or cam_id
|
||||
out[role] = cam_id
|
||||
return out
|
||||
|
||||
|
||||
def build_live_role_images(cam, frame, meta, calib: dict, beauty_preview=False) -> dict:
|
||||
"""
|
||||
Retorna imagens BGR por role: rgb/re/nir.
|
||||
|
||||
Usa os helpers do OakFcc3Client para reconstruir previews a partir do RAW_BRUTO.
|
||||
O script de referência usa get_next_decoded(...) e build_preview_from_raw_payload(...)
|
||||
para RAW_BRUTO; aqui aproveitamos build_visual_preview_from_raw(...), quando disponível.
|
||||
"""
|
||||
frame_type = meta.get("frame_type", "RAW_BRUTO")
|
||||
|
||||
if frame_type != "RAW_BRUTO" or not isinstance(frame, dict):
|
||||
raise RuntimeError("Este viewer espera frame_type=RAW_BRUTO e frame como dict por câmera.")
|
||||
|
||||
camera_info = meta.get("camera_info", {}) or {}
|
||||
|
||||
previews = None
|
||||
if hasattr(cam, "build_visual_preview_from_raw"):
|
||||
try:
|
||||
previews = cam.build_visual_preview_from_raw(frame, meta)
|
||||
except Exception:
|
||||
previews = None
|
||||
|
||||
if previews is None:
|
||||
preview_bgr, _, preview_source_id = cam.build_preview_from_raw_payload(frame=frame, meta=meta)
|
||||
previews = {preview_source_id: preview_bgr}
|
||||
|
||||
by_role = {}
|
||||
for cam_id, img in previews.items():
|
||||
role = camera_info.get(cam_id, {}).get("role", cam_id)
|
||||
by_role[role] = ensure_bgr_u8(img)
|
||||
|
||||
missing = [r for r in ["rgb", "re", "nir"] if r not in by_role]
|
||||
if missing:
|
||||
raise RuntimeError(f"Previews sem roles necessários: {missing}. Roles disponíveis={list(by_role.keys())}")
|
||||
|
||||
return by_role
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Alignment / depth / confidence
|
||||
# ============================================================
|
||||
|
||||
def warp_spectral_to_rgb(rgb_bgr, re_bgr, nir_bgr, calib: dict):
|
||||
h, w = rgb_bgr.shape[:2]
|
||||
|
||||
re_to_rgb = cv2.warpPerspective(
|
||||
re_bgr,
|
||||
calib["H_RE_to_RGB"],
|
||||
(w, h),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
|
||||
nir_to_rgb = cv2.warpPerspective(
|
||||
nir_bgr,
|
||||
calib["H_NIR_to_RGB"],
|
||||
(w, h),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
|
||||
return re_to_rgb, nir_to_rgb
|
||||
|
||||
|
||||
def make_sgbm(args):
|
||||
num_disp = max(16, int(round(args.num_disp / 16)) * 16)
|
||||
|
||||
block_size = max(3, int(args.block_size))
|
||||
if block_size % 2 == 0:
|
||||
block_size += 1
|
||||
|
||||
matcher = cv2.StereoSGBM_create(
|
||||
minDisparity=args.min_disp,
|
||||
numDisparities=num_disp,
|
||||
blockSize=block_size,
|
||||
P1=8 * block_size * block_size,
|
||||
P2=32 * block_size * block_size,
|
||||
disp12MaxDiff=1,
|
||||
uniquenessRatio=args.uniqueness,
|
||||
speckleWindowSize=args.speckle_window,
|
||||
speckleRange=args.speckle_range,
|
||||
preFilterCap=63,
|
||||
mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY,
|
||||
)
|
||||
|
||||
return matcher, num_disp, block_size
|
||||
|
||||
|
||||
def compute_stereo_disparity_and_confidence(re_bgr, nir_bgr, calib: dict, args):
|
||||
re_cam = calib["re_cam"]
|
||||
nir_cam = calib["nir_cam"]
|
||||
|
||||
re_gray = to_gray_u8(re_bgr)
|
||||
nir_gray = to_gray_u8(nir_bgr)
|
||||
|
||||
# Par estéreo no espaço RE/NIR retificado.
|
||||
re_rect, nir_rect = rectify_pair_gray(re_gray, nir_gray, calib, re_cam, nir_cam)
|
||||
|
||||
matcher, num_disp, block_size = make_sgbm(args)
|
||||
disp = matcher.compute(re_rect, nir_rect).astype(np.float32) / 16.0
|
||||
|
||||
valid = disp > args.min_valid_disp
|
||||
|
||||
# Normalização visual do depth/disparity.
|
||||
disp_vis = disp.copy()
|
||||
disp_vis[~valid] = 0.0
|
||||
|
||||
if np.count_nonzero(valid) > 20:
|
||||
vals = disp_vis[valid]
|
||||
p2 = np.percentile(vals, 2)
|
||||
p98 = np.percentile(vals, 98)
|
||||
disp_norm = (disp_vis - p2) / (p98 - p2 + 1e-6)
|
||||
disp_norm = np.clip(disp_norm, 0.0, 1.0)
|
||||
else:
|
||||
p2 = 0.0
|
||||
p98 = 1.0
|
||||
disp_norm = np.zeros_like(disp_vis, dtype=np.float32)
|
||||
|
||||
disp_color = colorize_scalar_0_1(disp_norm)
|
||||
|
||||
# Confiança geométrica simples:
|
||||
# - válida no SGBM
|
||||
# - penaliza saltos fortes de disparity
|
||||
# - penaliza regiões com baixa textura no par retificado
|
||||
valid_f = valid.astype(np.float32)
|
||||
|
||||
disp_smooth = cv2.GaussianBlur(disp_vis, (5, 5), 0)
|
||||
grad_x = cv2.Sobel(disp_smooth, cv2.CV_32F, 1, 0, ksize=3)
|
||||
grad_y = cv2.Sobel(disp_smooth, cv2.CV_32F, 0, 1, ksize=3)
|
||||
grad_mag = np.sqrt(grad_x * grad_x + grad_y * grad_y)
|
||||
|
||||
grad_penalty = np.clip(grad_mag / max(args.parallax_grad_ref, 1e-6), 0.0, 1.0)
|
||||
|
||||
tex_re = cv2.Laplacian(re_rect, cv2.CV_32F, ksize=3)
|
||||
tex_nir = cv2.Laplacian(nir_rect, cv2.CV_32F, ksize=3)
|
||||
texture = (np.abs(tex_re) + np.abs(tex_nir)) * 0.5
|
||||
texture_norm = np.clip(texture / max(args.texture_ref, 1e-6), 0.0, 1.0)
|
||||
texture_norm = cv2.GaussianBlur(texture_norm, (5, 5), 0)
|
||||
|
||||
confidence_rect = valid_f * (1.0 - grad_penalty) * (0.35 + 0.65 * texture_norm)
|
||||
confidence_rect = np.clip(confidence_rect, 0.0, 1.0)
|
||||
|
||||
# Para exibir junto com RGB, trazemos a confiança do espaço estéreo RE/NIR para RGB usando a homografia RE->RGB.
|
||||
h_rgb, w_rgb = re_bgr.shape[:2]
|
||||
confidence_rgb = cv2.warpPerspective(
|
||||
confidence_rect.astype(np.float32),
|
||||
calib["H_RE_to_RGB"],
|
||||
(w_rgb, h_rgb),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
|
||||
disp_rgb = cv2.warpPerspective(
|
||||
disp_norm.astype(np.float32),
|
||||
calib["H_RE_to_RGB"],
|
||||
(w_rgb, h_rgb),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
|
||||
confidence_color = colorize_scalar_0_1(confidence_rgb, cmap=cv2.COLORMAP_VIRIDIS)
|
||||
disp_rgb_color = colorize_scalar_0_1(disp_rgb, cmap=cv2.COLORMAP_TURBO)
|
||||
|
||||
stats = {
|
||||
"num_disp": num_disp,
|
||||
"block_size": block_size,
|
||||
"valid_pct": float(np.mean(valid) * 100.0),
|
||||
"conf_mean": float(np.mean(confidence_rect)),
|
||||
"disp_p02": float(p2),
|
||||
"disp_p98": float(p98),
|
||||
"disp_p50": float(np.percentile(disp_vis[valid], 50)) if np.count_nonzero(valid) > 20 else 0.0,
|
||||
}
|
||||
|
||||
return {
|
||||
"re_rect": cv2.cvtColor(re_rect, cv2.COLOR_GRAY2BGR),
|
||||
"nir_rect": cv2.cvtColor(nir_rect, cv2.COLOR_GRAY2BGR),
|
||||
"disp_rect_color": disp_color,
|
||||
"disp_rgb_color": disp_rgb_color,
|
||||
"confidence_rgb": confidence_rgb,
|
||||
"confidence_color": confidence_color,
|
||||
"stats": stats,
|
||||
}
|
||||
|
||||
|
||||
def apply_overlap_mask(img_bgr, calib: dict, enabled=True):
|
||||
if not enabled:
|
||||
return img_bgr
|
||||
|
||||
mask = calib.get("overlap_common_RGB")
|
||||
if mask is None:
|
||||
return img_bgr
|
||||
|
||||
if mask.shape[:2] != img_bgr.shape[:2]:
|
||||
mask = cv2.resize(mask, (img_bgr.shape[1], img_bgr.shape[0]), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
mask_bool = mask > 0
|
||||
out = img_bgr.copy()
|
||||
out[~mask_bool] = (out[~mask_bool] * 0.2).astype(np.uint8)
|
||||
return out
|
||||
|
||||
|
||||
def confidence_overlay_on_rgb(rgb_bgr, confidence_rgb, alpha=0.45):
|
||||
conf_color = colorize_scalar_0_1(confidence_rgb, cmap=cv2.COLORMAP_VIRIDIS)
|
||||
return cv2.addWeighted(rgb_bgr, 1.0 - alpha, conf_color, alpha, 0)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Main
|
||||
# ============================================================
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Live preview: RGB referência + RE/NIR alinhados por homografia + depth/confiança por estéreo RE-NIR.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
|
||||
parser.add_argument("--calib_path", required=True, help=".npz com H_RE_to_RGB, H_NIR_to_RGB e mapas estéreo RE-NIR")
|
||||
parser.add_argument("--fps", type=int, default=20)
|
||||
parser.add_argument("--width", type=int, default=RAW_SIZE[0])
|
||||
parser.add_argument("--height", type=int, default=RAW_SIZE[1])
|
||||
parser.add_argument("--bayer", default="BGGR", choices=["GBRG", "GRBG", "RGGB", "BGGR"])
|
||||
parser.add_argument("--module_calibration_json", default=MODULE_PARAMS)
|
||||
|
||||
parser.add_argument("--output_dtype", default="float32", choices=["uint8", "uint16", "float32"])
|
||||
parser.add_argument("--capture_mode", default="TRIPLE", choices=["AUTO", "SINGLE", "DOUBLE", "TRIPLE"])
|
||||
parser.add_argument("--raw_policy", default="require_triple", choices=["allow_single", "require_triple"])
|
||||
|
||||
parser.add_argument("--cell_w", type=int, default=640)
|
||||
parser.add_argument("--cell_h", type=int, default=400)
|
||||
parser.add_argument("--layout", default="2x2", choices=["2x2", "3x2"])
|
||||
parser.add_argument("--mask_overlap", action="store_true", help="Escurece área fora da interseção RE/NIR->RGB salva no .npz")
|
||||
parser.add_argument("--show_conf_overlay", action="store_true")
|
||||
|
||||
parser.add_argument("--num_disp", type=int, default=128)
|
||||
parser.add_argument("--block_size", type=int, default=7)
|
||||
parser.add_argument("--min_disp", type=int, default=0)
|
||||
parser.add_argument("--min_valid_disp", type=float, default=1.0)
|
||||
parser.add_argument("--uniqueness", type=int, default=8)
|
||||
parser.add_argument("--speckle_window", type=int, default=80)
|
||||
parser.add_argument("--speckle_range", type=int, default=2)
|
||||
parser.add_argument("--parallax_grad_ref", type=float, default=6.0)
|
||||
parser.add_argument("--texture_ref", type=float, default=25.0)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
calib = load_calibration_bundle(args.calib_path)
|
||||
|
||||
print("============================================")
|
||||
print("Live Multispec Alignment Preview")
|
||||
print(f"calib_path : {args.calib_path}")
|
||||
print(f"RGB cam : {calib['rgb_cam']}")
|
||||
print(f"RE cam : {calib['re_cam']}")
|
||||
print(f"NIR cam : {calib['nir_cam']}")
|
||||
print(f"image_size : {calib['image_size']}")
|
||||
print(f"raw : {args.width}x{args.height} | bayer={args.bayer}")
|
||||
print(f"capture : RAW_BRUTO | {args.capture_mode} | {args.raw_policy}")
|
||||
print("Keys : Q/Esc sair | O overlap | C conf overlay | [ ] numDisp | - + block")
|
||||
print("============================================")
|
||||
|
||||
window_name = "Live RGB/RE/NIR + Stereo Confidence"
|
||||
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
|
||||
|
||||
fps_view = 0.0
|
||||
n_view = 0
|
||||
t_fps = time.time()
|
||||
last_frame_id = -1
|
||||
msg = ""
|
||||
msg_t = 0.0
|
||||
|
||||
mask_overlap = args.mask_overlap
|
||||
conf_overlay = args.show_conf_overlay
|
||||
|
||||
try:
|
||||
with MultiSpectralClient(
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
bayer=args.bayer,
|
||||
fps=args.fps,
|
||||
frame_type="RAW_BRUTO",
|
||||
output_dtype=args.output_dtype,
|
||||
capture_mode=args.capture_mode,
|
||||
raw_policy=args.raw_policy,
|
||||
module_calibration_json=args.module_calibration_json,
|
||||
) as cam:
|
||||
while True:
|
||||
frame, meta, decoded = cam.get_next_decoded(timeout=1.0)
|
||||
|
||||
if meta is None or frame is None:
|
||||
k = cv2.waitKey(1) & 0xFF
|
||||
if k in (ord("q"), ord("Q"), 27):
|
||||
break
|
||||
continue
|
||||
|
||||
frame_id = meta.get("frame_id", -1)
|
||||
if frame_id == last_frame_id:
|
||||
k = cv2.waitKey(1) & 0xFF
|
||||
if k in (ord("q"), ord("Q"), 27):
|
||||
break
|
||||
continue
|
||||
|
||||
last_frame_id = frame_id
|
||||
|
||||
try:
|
||||
role_imgs = build_live_role_images(cam, frame, meta, calib)
|
||||
|
||||
rgb = role_imgs["rgb"]
|
||||
re = role_imgs["re"]
|
||||
nir = role_imgs["nir"]
|
||||
|
||||
# Garante que tudo use o tamanho do RGB como referência visual.
|
||||
rgb_h, rgb_w = rgb.shape[:2]
|
||||
if re.shape[:2] != (rgb_h, rgb_w):
|
||||
re = cv2.resize(re, (rgb_w, rgb_h), interpolation=cv2.INTER_AREA)
|
||||
if nir.shape[:2] != (rgb_h, rgb_w):
|
||||
nir = cv2.resize(nir, (rgb_w, rgb_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
re_rgb, nir_rgb = warp_spectral_to_rgb(rgb, re, nir, calib)
|
||||
stereo = compute_stereo_disparity_and_confidence(re, nir, calib, args)
|
||||
|
||||
rgb_show = apply_overlap_mask(rgb, calib, enabled=mask_overlap)
|
||||
re_show = apply_overlap_mask(re_rgb, calib, enabled=mask_overlap)
|
||||
nir_show = apply_overlap_mask(nir_rgb, calib, enabled=mask_overlap)
|
||||
|
||||
if conf_overlay:
|
||||
rgb_panel = confidence_overlay_on_rgb(rgb_show, stereo["confidence_rgb"])
|
||||
rgb_panel = put_label(rgb_panel, "RGB + confidence overlay")
|
||||
else:
|
||||
rgb_panel = put_label(rgb_show, "RGB reference")
|
||||
|
||||
re_panel = put_label(re_show, "RE -> RGB plane")
|
||||
nir_panel = put_label(nir_show, "NIR -> RGB plane")
|
||||
depth_panel = put_label(stereo["disp_rgb_color"], "Stereo disparity RE/NIR -> RGB")
|
||||
conf_panel = put_label(stereo["confidence_color"], "Spectral confidence / parallax risk")
|
||||
|
||||
s = stereo["stats"]
|
||||
|
||||
n_view += 1
|
||||
now = time.time()
|
||||
dt = now - t_fps
|
||||
if dt >= 1.0:
|
||||
fps_view = n_view / dt
|
||||
n_view = 0
|
||||
t_fps = now
|
||||
|
||||
hud = [
|
||||
f"frame_id={frame_id} | FPS_VIEW={fps_view:.1f} | overlap={'ON' if mask_overlap else 'OFF'} | conf_overlay={'ON' if conf_overlay else 'OFF'}",
|
||||
f"stereo valid={s['valid_pct']:.1f}% | conf_mean={s['conf_mean']:.2f} | disp p50={s['disp_p50']:.2f} | p02/p98={s['disp_p02']:.2f}/{s['disp_p98']:.2f}",
|
||||
f"numDisp={s['num_disp']} | block={s['block_size']} | O overlap | C conf | [ ] numDisp | - + block | Q sair",
|
||||
]
|
||||
|
||||
if msg and (time.time() - msg_t) < 2.0:
|
||||
hud.append(msg)
|
||||
|
||||
if args.layout == "2x2":
|
||||
canvas = make_2x2(
|
||||
rgb_panel,
|
||||
re_panel,
|
||||
nir_panel,
|
||||
depth_panel,
|
||||
cell_w=args.cell_w,
|
||||
cell_h=args.cell_h,
|
||||
hud_lines=hud,
|
||||
)
|
||||
else:
|
||||
canvas = make_3x2(
|
||||
rgb_panel,
|
||||
re_panel,
|
||||
nir_panel,
|
||||
depth_panel,
|
||||
conf_panel,
|
||||
put_label(stereo["disp_rect_color"], "Raw rectified disparity space"),
|
||||
cell_w=args.cell_w,
|
||||
cell_h=args.cell_h,
|
||||
hud_lines=hud,
|
||||
)
|
||||
|
||||
cv2.imshow(window_name, canvas)
|
||||
|
||||
except Exception as e:
|
||||
err = np.zeros((500, 1200, 3), dtype=np.uint8)
|
||||
cv2.putText(err, f"Erro: {e}", (20, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2, cv2.LINE_AA)
|
||||
cv2.imshow(window_name, err)
|
||||
print(f"[ERRO FRAME] {e}")
|
||||
|
||||
k = cv2.waitKey(1) & 0xFF
|
||||
|
||||
if k in (ord("q"), ord("Q"), 27):
|
||||
break
|
||||
|
||||
elif k in (ord("o"), ord("O")):
|
||||
mask_overlap = not mask_overlap
|
||||
msg = f"overlap mask -> {mask_overlap}"
|
||||
msg_t = time.time()
|
||||
|
||||
elif k in (ord("c"), ord("C")):
|
||||
conf_overlay = not conf_overlay
|
||||
msg = f"confidence overlay -> {conf_overlay}"
|
||||
msg_t = time.time()
|
||||
|
||||
elif k == ord("["):
|
||||
args.num_disp = max(16, args.num_disp - 16)
|
||||
msg = f"num_disp -> {args.num_disp}"
|
||||
msg_t = time.time()
|
||||
|
||||
elif k == ord("]"):
|
||||
args.num_disp = min(512, args.num_disp + 16)
|
||||
msg = f"num_disp -> {args.num_disp}"
|
||||
msg_t = time.time()
|
||||
|
||||
elif k in (ord("-"), ord("_")):
|
||||
args.block_size = max(3, args.block_size - 2)
|
||||
if args.block_size % 2 == 0:
|
||||
args.block_size -= 1
|
||||
msg = f"block_size -> {args.block_size}"
|
||||
msg_t = time.time()
|
||||
|
||||
elif k in (ord("+"), ord("=")):
|
||||
args.block_size = min(31, args.block_size + 2)
|
||||
if args.block_size % 2 == 0:
|
||||
args.block_size += 1
|
||||
msg = f"block_size -> {args.block_size}"
|
||||
msg_t = time.time()
|
||||
|
||||
finally:
|
||||
cv2.destroyAllWindows()
|
||||
print("Fim do preview live.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,565 @@
|
|||
import argparse
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
# ============================================================
|
||||
# RAW10
|
||||
# ============================================================
|
||||
|
||||
def unpack_raw10_packed(raw: bytes, width: int, height: int) -> np.ndarray:
|
||||
arr = np.frombuffer(raw, dtype=np.uint8)
|
||||
|
||||
pixel_count = width * height
|
||||
expected_bytes = (pixel_count // 4) * 5
|
||||
|
||||
if pixel_count % 4 != 0:
|
||||
raise RuntimeError(f"width*height precisa ser múltiplo de 4. Recebido: {pixel_count}")
|
||||
|
||||
if arr.size < expected_bytes:
|
||||
raise RuntimeError(
|
||||
f"RAW10 menor que esperado. bytes={arr.size}, esperado={expected_bytes}, "
|
||||
f"width={width}, height={height}"
|
||||
)
|
||||
|
||||
arr = arr[:expected_bytes]
|
||||
groups = arr.reshape(-1, 5).astype(np.uint16)
|
||||
|
||||
p0 = (groups[:, 0] << 2) | ((groups[:, 4] >> 0) & 0x03)
|
||||
p1 = (groups[:, 1] << 2) | ((groups[:, 4] >> 2) & 0x03)
|
||||
p2 = (groups[:, 2] << 2) | ((groups[:, 4] >> 4) & 0x03)
|
||||
p3 = (groups[:, 3] << 2) | ((groups[:, 4] >> 6) & 0x03)
|
||||
|
||||
out = np.empty(groups.shape[0] * 4, dtype=np.uint16)
|
||||
out[0::4] = p0
|
||||
out[1::4] = p1
|
||||
out[2::4] = p2
|
||||
out[3::4] = p3
|
||||
|
||||
return out.reshape(height, width)
|
||||
|
||||
|
||||
def normalize_to_u8(img: np.ndarray, p_low=1.0, p_high=99.0) -> np.ndarray:
|
||||
arr = img.astype(np.float32)
|
||||
valid = np.isfinite(arr)
|
||||
|
||||
if np.count_nonzero(valid) < 20:
|
||||
return np.zeros(arr.shape[:2], dtype=np.uint8)
|
||||
|
||||
vals = arr[valid]
|
||||
lo = np.percentile(vals, p_low)
|
||||
hi = np.percentile(vals, p_high)
|
||||
|
||||
out = (arr - lo) / (hi - lo + 1e-6)
|
||||
out = np.clip(out, 0.0, 1.0)
|
||||
|
||||
return (out * 255).astype(np.uint8)
|
||||
|
||||
|
||||
def raw10_bin_to_gray(path: Path, width: int, height: int, use_clahe=True) -> np.ndarray:
|
||||
raw = path.read_bytes()
|
||||
mono10 = unpack_raw10_packed(raw, width, height)
|
||||
gray = normalize_to_u8(mono10)
|
||||
|
||||
if use_clahe:
|
||||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
gray = clahe.apply(gray)
|
||||
|
||||
return gray
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Pairing CAM bins
|
||||
# ============================================================
|
||||
|
||||
def clean_stem_for_pair(path: Path, cam_key: str):
|
||||
s = path.stem
|
||||
|
||||
variants = [
|
||||
cam_key,
|
||||
cam_key.lower(),
|
||||
cam_key.replace("_", ""),
|
||||
cam_key.replace("_", "").lower(),
|
||||
]
|
||||
|
||||
for v in variants:
|
||||
s = s.replace(v, "")
|
||||
|
||||
s = re.sub(r"[_\-\s]+", "_", s).strip("_").lower()
|
||||
return s
|
||||
|
||||
|
||||
def find_cam_bins(root_dir: Path, cam_key: str):
|
||||
return sorted([
|
||||
p for p in root_dir.rglob("*.bin")
|
||||
if cam_key.lower() in p.name.lower()
|
||||
])
|
||||
|
||||
|
||||
def pair_cam_bins(root_dir: Path, left_cam: str, right_cam: str):
|
||||
left_bins = find_cam_bins(root_dir, left_cam)
|
||||
right_bins = find_cam_bins(root_dir, right_cam)
|
||||
|
||||
right_by_folder_key = {}
|
||||
right_by_key = {}
|
||||
|
||||
for rp in right_bins:
|
||||
key = clean_stem_for_pair(rp, right_cam)
|
||||
right_by_folder_key[(rp.parent, key)] = rp
|
||||
right_by_key.setdefault(key, rp)
|
||||
|
||||
pairs = []
|
||||
|
||||
for lp in left_bins:
|
||||
key = clean_stem_for_pair(lp, left_cam)
|
||||
|
||||
rp = right_by_folder_key.get((lp.parent, key))
|
||||
if rp is None:
|
||||
rp = right_by_key.get(key)
|
||||
|
||||
if rp is None:
|
||||
same_folder = [r for r in right_bins if r.parent == lp.parent]
|
||||
if len(same_folder) == 1:
|
||||
rp = same_folder[0]
|
||||
|
||||
if rp is not None:
|
||||
pairs.append((lp, rp))
|
||||
|
||||
if pairs:
|
||||
return pairs, left_bins, right_bins, "name/key"
|
||||
|
||||
n = min(len(left_bins), len(right_bins))
|
||||
return list(zip(left_bins[:n], right_bins[:n])), left_bins, right_bins, "order"
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Calibration loading
|
||||
# ============================================================
|
||||
|
||||
def load_stereo_calib(calib_path: Path):
|
||||
data = np.load(str(calib_path), allow_pickle=True)
|
||||
|
||||
required = ["map1x", "map1y", "map2x", "map2y", "image_size"]
|
||||
for k in required:
|
||||
if k not in data:
|
||||
raise RuntimeError(f"Calibração sem chave obrigatória: {k}")
|
||||
|
||||
calib = {
|
||||
"map1x": data["map1x"],
|
||||
"map1y": data["map1y"],
|
||||
"map2x": data["map2x"],
|
||||
"map2y": data["map2y"],
|
||||
"image_size": tuple(data["image_size"].astype(int).tolist()),
|
||||
}
|
||||
|
||||
for k in ["left_cam", "right_cam", "rms_left", "rms_right", "rms_stereo", "T"]:
|
||||
if k in data:
|
||||
calib[k] = data[k]
|
||||
|
||||
return calib
|
||||
|
||||
|
||||
def rectify_pair(left_gray, right_gray, calib):
|
||||
image_w, image_h = calib["image_size"]
|
||||
|
||||
if left_gray.shape[::-1] != (image_w, image_h):
|
||||
left_gray = cv2.resize(left_gray, (image_w, image_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
if right_gray.shape[::-1] != (image_w, image_h):
|
||||
right_gray = cv2.resize(right_gray, (image_w, image_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
left_rect = cv2.remap(
|
||||
left_gray,
|
||||
calib["map1x"],
|
||||
calib["map1y"],
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
|
||||
right_rect = cv2.remap(
|
||||
right_gray,
|
||||
calib["map2x"],
|
||||
calib["map2y"],
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
|
||||
return left_rect, right_rect
|
||||
|
||||
|
||||
# ============================================================
|
||||
# SGBM
|
||||
# ============================================================
|
||||
|
||||
def make_sgbm(args):
|
||||
num_disp = max(16, int(round(args.num_disp / 16)) * 16)
|
||||
|
||||
block_size = max(3, int(args.block_size))
|
||||
if block_size % 2 == 0:
|
||||
block_size += 1
|
||||
|
||||
matcher = cv2.StereoSGBM_create(
|
||||
minDisparity=args.min_disp,
|
||||
numDisparities=num_disp,
|
||||
blockSize=block_size,
|
||||
P1=8 * block_size * block_size,
|
||||
P2=32 * block_size * block_size,
|
||||
disp12MaxDiff=1,
|
||||
uniquenessRatio=args.uniqueness,
|
||||
speckleWindowSize=args.speckle_window,
|
||||
speckleRange=args.speckle_range,
|
||||
preFilterCap=63,
|
||||
mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY,
|
||||
)
|
||||
|
||||
return matcher, num_disp, block_size
|
||||
|
||||
|
||||
def compute_disparity(left_rect, right_rect, args):
|
||||
matcher, num_disp, block_size = make_sgbm(args)
|
||||
|
||||
disp_raw = matcher.compute(left_rect, right_rect).astype(np.float32) / 16.0
|
||||
|
||||
valid = disp_raw > args.min_valid_disp
|
||||
|
||||
disp_vis = disp_raw.copy()
|
||||
disp_vis[~valid] = 0.0
|
||||
|
||||
if np.count_nonzero(valid) > 20:
|
||||
vals = disp_vis[valid]
|
||||
p2 = np.percentile(vals, 2)
|
||||
p98 = np.percentile(vals, 98)
|
||||
disp_norm = (disp_vis - p2) / (p98 - p2 + 1e-6)
|
||||
disp_norm = np.clip(disp_norm, 0.0, 1.0)
|
||||
else:
|
||||
disp_norm = np.zeros_like(disp_vis, dtype=np.float32)
|
||||
|
||||
disp_color = cv2.applyColorMap((disp_norm * 255).astype(np.uint8), cv2.COLORMAP_TURBO)
|
||||
|
||||
valid_mask = np.zeros_like(disp_color)
|
||||
valid_mask[valid] = (255, 255, 255)
|
||||
|
||||
stats = {
|
||||
"num_disp": num_disp,
|
||||
"block_size": block_size,
|
||||
"valid_pct": float(np.mean(valid) * 100.0),
|
||||
"disp_p05": float(np.percentile(disp_vis[valid], 5)) if np.count_nonzero(valid) > 20 else 0.0,
|
||||
"disp_p50": float(np.percentile(disp_vis[valid], 50)) if np.count_nonzero(valid) > 20 else 0.0,
|
||||
"disp_p95": float(np.percentile(disp_vis[valid], 95)) if np.count_nonzero(valid) > 20 else 0.0,
|
||||
}
|
||||
|
||||
return disp_raw, disp_color, valid_mask, valid, stats
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Visualization
|
||||
# ============================================================
|
||||
|
||||
def draw_epipolar_lines(img_bgr, step=40):
|
||||
out = img_bgr.copy()
|
||||
h, w = out.shape[:2]
|
||||
|
||||
for y in range(0, h, step):
|
||||
color = (0, 255, 255) if (y // step) % 2 == 0 else (255, 255, 0)
|
||||
cv2.line(out, (0, y), (w, y), color, 1, cv2.LINE_AA)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def resize_to_height(img, target_h):
|
||||
h, w = img.shape[:2]
|
||||
if h == target_h:
|
||||
return img
|
||||
|
||||
scale = target_h / h
|
||||
new_w = max(1, int(w * scale))
|
||||
return cv2.resize(img, (new_w, target_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
|
||||
def draw_header(canvas, lines):
|
||||
header_h = 24 + 24 * len(lines)
|
||||
cv2.rectangle(canvas, (0, 0), (canvas.shape[1], header_h), (0, 0, 0), -1)
|
||||
|
||||
y = 24
|
||||
for line in lines:
|
||||
cv2.putText(
|
||||
canvas,
|
||||
line,
|
||||
(12, y),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.55,
|
||||
(255, 255, 255),
|
||||
1,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
y += 24
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def make_overlay(left_bgr, disp_color, alpha=0.45):
|
||||
disp_resized = cv2.resize(
|
||||
disp_color,
|
||||
(left_bgr.shape[1], left_bgr.shape[0]),
|
||||
interpolation=cv2.INTER_AREA,
|
||||
)
|
||||
return cv2.addWeighted(left_bgr, 1.0 - alpha, disp_resized, alpha, 0)
|
||||
|
||||
|
||||
def compose_view(left_rect, right_rect, disp_color, valid_mask, pair, idx, total, mode, stats, args):
|
||||
left_bgr = cv2.cvtColor(left_rect, cv2.COLOR_GRAY2BGR)
|
||||
right_bgr = cv2.cvtColor(right_rect, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
if args.lines:
|
||||
left_bgr = draw_epipolar_lines(left_bgr, step=args.line_step)
|
||||
right_bgr = draw_epipolar_lines(right_bgr, step=args.line_step)
|
||||
|
||||
if mode == "disp":
|
||||
third = disp_color
|
||||
mode_name = "disparity"
|
||||
elif mode == "mask":
|
||||
third = valid_mask
|
||||
mode_name = "valid mask"
|
||||
elif mode == "overlay":
|
||||
third = make_overlay(cv2.cvtColor(left_rect, cv2.COLOR_GRAY2BGR), disp_color)
|
||||
mode_name = "overlay"
|
||||
else:
|
||||
diff = cv2.absdiff(left_rect, right_rect)
|
||||
third = cv2.cvtColor(diff, cv2.COLOR_GRAY2BGR)
|
||||
mode_name = "rect diff"
|
||||
|
||||
left = resize_to_height(left_bgr, args.view_h)
|
||||
right = resize_to_height(right_bgr, args.view_h)
|
||||
third = resize_to_height(third, args.view_h)
|
||||
|
||||
h = min(left.shape[0], right.shape[0], third.shape[0])
|
||||
left = left[:h]
|
||||
right = right[:h]
|
||||
third = third[:h]
|
||||
|
||||
canvas = np.hstack([left, right, third])
|
||||
|
||||
lp, rp = pair
|
||||
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | L={lp.name} | R={rp.name}",
|
||||
f"mode={mode_name} | valid={stats['valid_pct']:.1f}% | disp p05={stats['disp_p05']:.2f} p50={stats['disp_p50']:.2f} p95={stats['disp_p95']:.2f}",
|
||||
f"numDisp={stats['num_disp']} | block={stats['block_size']} | uniqueness={args.uniqueness} | lines={args.lines}",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | [ ] numDisp | - + block | S salvar | Q sair",
|
||||
]
|
||||
|
||||
return draw_header(canvas, lines)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Main
|
||||
# ============================================================
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument("--root_dir", required=True)
|
||||
parser.add_argument("--calib_path", required=True)
|
||||
|
||||
parser.add_argument("--left_cam", default="CAM_C")
|
||||
parser.add_argument("--right_cam", default="CAM_B")
|
||||
|
||||
parser.add_argument("--width", type=int, default=1280)
|
||||
parser.add_argument("--height", type=int, default=800)
|
||||
|
||||
parser.add_argument("--view_h", type=int, default=480)
|
||||
|
||||
parser.add_argument("--num_disp", type=int, default=128)
|
||||
parser.add_argument("--block_size", type=int, default=7)
|
||||
parser.add_argument("--min_disp", type=int, default=0)
|
||||
parser.add_argument("--min_valid_disp", type=float, default=1.0)
|
||||
|
||||
parser.add_argument("--uniqueness", type=int, default=8)
|
||||
parser.add_argument("--speckle_window", type=int, default=80)
|
||||
parser.add_argument("--speckle_range", type=int, default=2)
|
||||
|
||||
parser.add_argument("--no_clahe", action="store_true")
|
||||
parser.add_argument("--lines", action="store_true")
|
||||
parser.add_argument("--line_step", type=int, default=40)
|
||||
|
||||
parser.add_argument("--save_dir", default="stereo_rectified_sgbm_saves")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
root_dir = Path(args.root_dir)
|
||||
calib_path = Path(args.calib_path)
|
||||
save_dir = Path(args.save_dir)
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
calib = load_stereo_calib(calib_path)
|
||||
|
||||
print(f"[INFO] calib_path: {calib_path}")
|
||||
print(f"[INFO] calib image_size: {calib['image_size']}")
|
||||
|
||||
if "rms_left" in calib:
|
||||
print(f"[INFO] rms_left: {float(calib['rms_left']):.6f}")
|
||||
if "rms_right" in calib:
|
||||
print(f"[INFO] rms_right: {float(calib['rms_right']):.6f}")
|
||||
if "rms_stereo" in calib:
|
||||
print(f"[INFO] rms_stereo: {float(calib['rms_stereo']):.6f}")
|
||||
if "T" in calib:
|
||||
print(f"[INFO] T: {np.array(calib['T']).ravel()}")
|
||||
|
||||
pairs, left_bins, right_bins, pair_mode = pair_cam_bins(root_dir, args.left_cam, args.right_cam)
|
||||
|
||||
print(f"[INFO] root_dir: {root_dir}")
|
||||
print(f"[INFO] left_cam: {args.left_cam} | arquivos: {len(left_bins)}")
|
||||
print(f"[INFO] right_cam: {args.right_cam} | arquivos: {len(right_bins)}")
|
||||
print(f"[INFO] pares: {len(pairs)}")
|
||||
print(f"[INFO] pareamento: {pair_mode}")
|
||||
|
||||
if not pairs:
|
||||
raise RuntimeError("Nenhum par encontrado.")
|
||||
|
||||
idx = 0
|
||||
mode = "disp"
|
||||
|
||||
cached_key = None
|
||||
cached_data = None
|
||||
|
||||
cv2.namedWindow("Stereo RAW10 Rectified SGBM Viewer", cv2.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
pair = pairs[idx]
|
||||
lp, rp = pair
|
||||
|
||||
key_cache = (
|
||||
str(lp),
|
||||
str(rp),
|
||||
args.num_disp,
|
||||
args.block_size,
|
||||
args.uniqueness,
|
||||
args.speckle_window,
|
||||
args.speckle_range,
|
||||
args.min_disp,
|
||||
args.min_valid_disp,
|
||||
args.no_clahe,
|
||||
args.width,
|
||||
args.height,
|
||||
str(args.calib_path),
|
||||
)
|
||||
|
||||
if key_cache != cached_key:
|
||||
print(f"[RUN] {idx + 1}/{len(pairs)} - L={lp.name} | R={rp.name}")
|
||||
|
||||
try:
|
||||
left_gray = raw10_bin_to_gray(
|
||||
lp,
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
use_clahe=not args.no_clahe,
|
||||
)
|
||||
|
||||
right_gray = raw10_bin_to_gray(
|
||||
rp,
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
use_clahe=not args.no_clahe,
|
||||
)
|
||||
|
||||
left_rect, right_rect = rectify_pair(left_gray, right_gray, calib)
|
||||
|
||||
_, disp_color, valid_mask, valid, stats = compute_disparity(left_rect, right_rect, args)
|
||||
|
||||
cached_data = (left_rect, right_rect, disp_color, valid_mask, stats)
|
||||
cached_key = key_cache
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERRO] Falha processando par:")
|
||||
print(f" L={lp}")
|
||||
print(f" R={rp}")
|
||||
print(f" erro={e}")
|
||||
idx = min(idx + 1, len(pairs) - 1)
|
||||
cached_key = None
|
||||
cached_data = None
|
||||
continue
|
||||
|
||||
else:
|
||||
left_rect, right_rect, disp_color, valid_mask, stats = cached_data
|
||||
|
||||
view = compose_view(
|
||||
left_rect=left_rect,
|
||||
right_rect=right_rect,
|
||||
disp_color=disp_color,
|
||||
valid_mask=valid_mask,
|
||||
pair=pair,
|
||||
idx=idx,
|
||||
total=len(pairs),
|
||||
mode=mode,
|
||||
stats=stats,
|
||||
args=args,
|
||||
)
|
||||
|
||||
cv2.imshow("Stereo RAW10 Rectified SGBM Viewer", view)
|
||||
key = cv2.waitKey(0) & 0xFF
|
||||
|
||||
if key in [27, ord("q"), ord("Q")]:
|
||||
break
|
||||
|
||||
elif key in [ord("n"), ord("N"), 32]:
|
||||
idx = min(idx + 1, len(pairs) - 1)
|
||||
cached_key = None
|
||||
|
||||
elif key in [ord("a"), ord("A")]:
|
||||
idx = max(idx - 1, 0)
|
||||
cached_key = None
|
||||
|
||||
elif key in [ord("m"), ord("M")]:
|
||||
if mode == "disp":
|
||||
mode = "mask"
|
||||
elif mode == "mask":
|
||||
mode = "overlay"
|
||||
elif mode == "overlay":
|
||||
mode = "diff"
|
||||
else:
|
||||
mode = "disp"
|
||||
|
||||
elif key in [ord("l"), ord("L")]:
|
||||
args.lines = not args.lines
|
||||
print(f"[PARAM] lines={args.lines}")
|
||||
|
||||
elif key == ord("["):
|
||||
args.num_disp = max(16, args.num_disp - 16)
|
||||
cached_key = None
|
||||
print(f"[PARAM] num_disp={args.num_disp}")
|
||||
|
||||
elif key == ord("]"):
|
||||
args.num_disp = min(512, args.num_disp + 16)
|
||||
cached_key = None
|
||||
print(f"[PARAM] num_disp={args.num_disp}")
|
||||
|
||||
elif key in [ord("-"), ord("_")]:
|
||||
args.block_size = max(3, args.block_size - 2)
|
||||
if args.block_size % 2 == 0:
|
||||
args.block_size -= 1
|
||||
cached_key = None
|
||||
print(f"[PARAM] block_size={args.block_size}")
|
||||
|
||||
elif key in [ord("+"), ord("=")]:
|
||||
args.block_size = min(31, args.block_size + 2)
|
||||
if args.block_size % 2 == 0:
|
||||
args.block_size += 1
|
||||
cached_key = None
|
||||
print(f"[PARAM] block_size={args.block_size}")
|
||||
|
||||
elif key in [ord("s"), ord("S")]:
|
||||
out_path = save_dir / f"stereo_rectified_{idx:04d}_{mode}.png"
|
||||
cv2.imwrite(str(out_path), view)
|
||||
print(f"[SAVE] {out_path}")
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,786 @@
|
|||
import argparse
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
# ============================================================
|
||||
# RAW10 / image conversion
|
||||
# ============================================================
|
||||
|
||||
def unpack_raw10_packed(raw: bytes, width: int, height: int) -> np.ndarray:
|
||||
arr = np.frombuffer(raw, dtype=np.uint8)
|
||||
|
||||
pixel_count = width * height
|
||||
expected_bytes = (pixel_count // 4) * 5
|
||||
|
||||
if pixel_count % 4 != 0:
|
||||
raise RuntimeError(f"width*height precisa ser múltiplo de 4. Recebido: {pixel_count}")
|
||||
|
||||
if arr.size < expected_bytes:
|
||||
raise RuntimeError(
|
||||
f"RAW10 menor que esperado. bytes={arr.size}, esperado={expected_bytes}, "
|
||||
f"width={width}, height={height}"
|
||||
)
|
||||
|
||||
arr = arr[:expected_bytes]
|
||||
groups = arr.reshape(-1, 5).astype(np.uint16)
|
||||
|
||||
p0 = (groups[:, 0] << 2) | ((groups[:, 4] >> 0) & 0x03)
|
||||
p1 = (groups[:, 1] << 2) | ((groups[:, 4] >> 2) & 0x03)
|
||||
p2 = (groups[:, 2] << 2) | ((groups[:, 4] >> 4) & 0x03)
|
||||
p3 = (groups[:, 3] << 2) | ((groups[:, 4] >> 6) & 0x03)
|
||||
|
||||
out = np.empty(groups.shape[0] * 4, dtype=np.uint16)
|
||||
out[0::4] = p0
|
||||
out[1::4] = p1
|
||||
out[2::4] = p2
|
||||
out[3::4] = p3
|
||||
|
||||
return out.reshape(height, width)
|
||||
|
||||
|
||||
def normalize_to_u8(img: np.ndarray, p_low=1.0, p_high=99.0) -> np.ndarray:
|
||||
arr = img.astype(np.float32)
|
||||
valid = np.isfinite(arr)
|
||||
|
||||
if np.count_nonzero(valid) < 20:
|
||||
return np.zeros(arr.shape[:2], dtype=np.uint8)
|
||||
|
||||
vals = arr[valid]
|
||||
lo = np.percentile(vals, p_low)
|
||||
hi = np.percentile(vals, p_high)
|
||||
|
||||
out = (arr - lo) / (hi - lo + 1e-6)
|
||||
out = np.clip(out, 0.0, 1.0)
|
||||
return (out * 255).astype(np.uint8)
|
||||
|
||||
|
||||
def debayer_raw10_to_bgr_u8(raw10: np.ndarray, bayer: str) -> np.ndarray:
|
||||
gray_u8 = normalize_to_u8(raw10)
|
||||
bayer = bayer.upper()
|
||||
|
||||
code_map = {
|
||||
"RGGB": cv2.COLOR_BayerRG2BGR,
|
||||
"BGGR": cv2.COLOR_BayerBG2BGR,
|
||||
"GRBG": cv2.COLOR_BayerGR2BGR,
|
||||
"GBRG": cv2.COLOR_BayerGB2BGR,
|
||||
}
|
||||
|
||||
if bayer not in code_map:
|
||||
raise RuntimeError(f"Bayer pattern não suportado: {bayer}")
|
||||
|
||||
return cv2.cvtColor(gray_u8, code_map[bayer])
|
||||
|
||||
|
||||
def read_raw10_rgb_bgr(path: Path, width: int, height: int, bayer: str) -> np.ndarray:
|
||||
raw = path.read_bytes()
|
||||
raw10 = unpack_raw10_packed(raw, width, height)
|
||||
return debayer_raw10_to_bgr_u8(raw10, bayer=bayer)
|
||||
|
||||
|
||||
def read_raw10_rgb_view_bgr(
|
||||
path: Path,
|
||||
width: int,
|
||||
height: int,
|
||||
bayer: str,
|
||||
rgb_view: str,
|
||||
use_clahe: bool = True,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Carrega CAM_A/RGB em três modos:
|
||||
|
||||
color:
|
||||
RAW10 Bayer -> debayer BGR -> visual colorido.
|
||||
|
||||
gray:
|
||||
RAW10 Bayer -> debayer BGR -> grayscale -> CLAHE -> BGR fake.
|
||||
Este é o mais parecido com o caminho usado na calibração ChArUco.
|
||||
|
||||
raw_bayer_gray:
|
||||
RAW10 Bayer -> normalize direto -> CLAHE -> BGR fake.
|
||||
Não faz debayer; útil para testar se a interpolação do debayer está influenciando.
|
||||
"""
|
||||
raw = path.read_bytes()
|
||||
raw10 = unpack_raw10_packed(raw, width, height)
|
||||
mode = rgb_view.lower().strip()
|
||||
|
||||
if mode == "color":
|
||||
return debayer_raw10_to_bgr_u8(raw10, bayer=bayer)
|
||||
|
||||
if mode == "gray":
|
||||
bgr = debayer_raw10_to_bgr_u8(raw10, bayer=bayer)
|
||||
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
elif mode == "raw_bayer_gray":
|
||||
gray = normalize_to_u8(raw10)
|
||||
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"rgb_view inválido: {rgb_view}. Use: color, gray ou raw_bayer_gray"
|
||||
)
|
||||
|
||||
if use_clahe:
|
||||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
gray = clahe.apply(gray)
|
||||
|
||||
return cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
|
||||
def read_raw10_mono_bgr(path: Path, width: int, height: int, use_clahe=True) -> np.ndarray:
|
||||
raw = path.read_bytes()
|
||||
raw10 = unpack_raw10_packed(raw, width, height)
|
||||
gray = normalize_to_u8(raw10)
|
||||
|
||||
if use_clahe:
|
||||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
gray = clahe.apply(gray)
|
||||
|
||||
return cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
|
||||
def to_gray_u8(img_bgr: np.ndarray) -> np.ndarray:
|
||||
if img_bgr.ndim == 2:
|
||||
return img_bgr
|
||||
return cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Triplet pairing
|
||||
# ============================================================
|
||||
|
||||
def clean_stem_for_pair(path: Path, cam_key: str):
|
||||
s = path.stem
|
||||
variants = [
|
||||
cam_key,
|
||||
cam_key.lower(),
|
||||
cam_key.replace("_", ""),
|
||||
cam_key.replace("_", "").lower(),
|
||||
]
|
||||
|
||||
for v in variants:
|
||||
s = s.replace(v, "")
|
||||
|
||||
s = re.sub(r"[_\-\s]+", "_", s).strip("_").lower()
|
||||
return s
|
||||
|
||||
|
||||
def find_cam_bins(root_dir: Path, cam_key: str):
|
||||
return sorted([p for p in root_dir.rglob("*.bin") if cam_key.lower() in p.name.lower()])
|
||||
|
||||
|
||||
def find_triplets(root_dir: Path, cams: list[str]):
|
||||
by_cam = {cam: find_cam_bins(root_dir, cam) for cam in cams}
|
||||
key_maps = {}
|
||||
|
||||
for cam, paths in by_cam.items():
|
||||
m = {}
|
||||
for p in paths:
|
||||
key = clean_stem_for_pair(p, cam)
|
||||
m[(p.parent, key)] = p
|
||||
m.setdefault((None, key), p)
|
||||
key_maps[cam] = m
|
||||
|
||||
ref_cam = cams[0]
|
||||
triplets = []
|
||||
|
||||
for ref_path in by_cam[ref_cam]:
|
||||
key = clean_stem_for_pair(ref_path, ref_cam)
|
||||
folder = ref_path.parent
|
||||
item = {ref_cam: ref_path}
|
||||
ok = True
|
||||
|
||||
for cam in cams[1:]:
|
||||
p = key_maps[cam].get((folder, key)) or key_maps[cam].get((None, key))
|
||||
if p is None:
|
||||
same_folder = [x for x in by_cam[cam] if x.parent == folder]
|
||||
if len(same_folder) == 1:
|
||||
p = same_folder[0]
|
||||
|
||||
if p is None:
|
||||
ok = False
|
||||
break
|
||||
|
||||
item[cam] = p
|
||||
|
||||
if ok:
|
||||
triplets.append(item)
|
||||
|
||||
return triplets, by_cam
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Calibration loading
|
||||
# ============================================================
|
||||
|
||||
def scalar_str(x):
|
||||
arr = np.array(x)
|
||||
if arr.shape == ():
|
||||
return str(arr.item())
|
||||
return str(x)
|
||||
|
||||
|
||||
def load_multicam_calib(calib_path: Path):
|
||||
data = np.load(str(calib_path), allow_pickle=True)
|
||||
keys = set(data.files)
|
||||
|
||||
required = ["image_size", "rgb_cam", "nir_cam", "re_cam", "ref_cam"]
|
||||
for k in required:
|
||||
if k not in keys:
|
||||
raise RuntimeError(f"Calibração multicam sem chave obrigatória: {k}")
|
||||
|
||||
calib = {
|
||||
"data": data,
|
||||
"keys": keys,
|
||||
"image_size": tuple(data["image_size"].astype(int).tolist()),
|
||||
"rgb_cam": scalar_str(data["rgb_cam"]),
|
||||
"nir_cam": scalar_str(data["nir_cam"]),
|
||||
"re_cam": scalar_str(data["re_cam"]),
|
||||
"ref_cam": scalar_str(data["ref_cam"]),
|
||||
}
|
||||
|
||||
return calib
|
||||
|
||||
|
||||
def get_pair_prefix(calib, cam1: str, cam2: str):
|
||||
keys = calib["keys"]
|
||||
direct = f"pair_{cam1}_{cam2}"
|
||||
inv = f"pair_{cam2}_{cam1}"
|
||||
|
||||
if f"{direct}_map1x" in keys:
|
||||
return direct, False
|
||||
|
||||
if f"{inv}_map1x" in keys:
|
||||
return inv, True
|
||||
|
||||
raise RuntimeError(f"Par {cam1}<->{cam2} não encontrado no .npz")
|
||||
|
||||
|
||||
def rectify_pair_from_calib(img1_bgr, img2_bgr, calib, cam1: str, cam2: str):
|
||||
"""
|
||||
Retifica duas imagens usando o par salvo no .npz.
|
||||
|
||||
Retorna imagens na ordem solicitada: cam1_rect, cam2_rect.
|
||||
Se o par salvo estiver invertido, troca map1/map2 automaticamente.
|
||||
"""
|
||||
data = calib["data"]
|
||||
image_w, image_h = calib["image_size"]
|
||||
|
||||
if img1_bgr.shape[1] != image_w or img1_bgr.shape[0] != image_h:
|
||||
img1_bgr = cv2.resize(img1_bgr, (image_w, image_h), interpolation=cv2.INTER_AREA)
|
||||
if img2_bgr.shape[1] != image_w or img2_bgr.shape[0] != image_h:
|
||||
img2_bgr = cv2.resize(img2_bgr, (image_w, image_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
prefix, inverted = get_pair_prefix(calib, cam1, cam2)
|
||||
|
||||
if not inverted:
|
||||
map1x = data[f"{prefix}_map1x"]
|
||||
map1y = data[f"{prefix}_map1y"]
|
||||
map2x = data[f"{prefix}_map2x"]
|
||||
map2y = data[f"{prefix}_map2y"]
|
||||
else:
|
||||
map1x = data[f"{prefix}_map2x"]
|
||||
map1y = data[f"{prefix}_map2y"]
|
||||
map2x = data[f"{prefix}_map1x"]
|
||||
map2y = data[f"{prefix}_map1y"]
|
||||
|
||||
rect1 = cv2.remap(
|
||||
img1_bgr,
|
||||
map1x,
|
||||
map1y,
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
|
||||
rect2 = cv2.remap(
|
||||
img2_bgr,
|
||||
map2x,
|
||||
map2y,
|
||||
interpolation=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
|
||||
return rect1, rect2
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Disparity RE/NIR from multicam npz
|
||||
# ============================================================
|
||||
|
||||
def make_sgbm(args):
|
||||
num_disp = max(16, int(round(args.num_disp / 16)) * 16)
|
||||
|
||||
block_size = max(3, int(args.block_size))
|
||||
if block_size % 2 == 0:
|
||||
block_size += 1
|
||||
|
||||
matcher = cv2.StereoSGBM_create(
|
||||
minDisparity=args.min_disp,
|
||||
numDisparities=num_disp,
|
||||
blockSize=block_size,
|
||||
P1=8 * block_size * block_size,
|
||||
P2=32 * block_size * block_size,
|
||||
disp12MaxDiff=1,
|
||||
uniquenessRatio=args.uniqueness,
|
||||
speckleWindowSize=args.speckle_window,
|
||||
speckleRange=args.speckle_range,
|
||||
preFilterCap=63,
|
||||
mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY,
|
||||
)
|
||||
|
||||
return matcher, num_disp, block_size
|
||||
|
||||
|
||||
def compute_disparity(left_rect_bgr, right_rect_bgr, args):
|
||||
left_gray = to_gray_u8(left_rect_bgr)
|
||||
right_gray = to_gray_u8(right_rect_bgr)
|
||||
|
||||
matcher, num_disp, block_size = make_sgbm(args)
|
||||
disp_raw = matcher.compute(left_gray, right_gray).astype(np.float32) / 16.0
|
||||
valid = disp_raw > args.min_valid_disp
|
||||
|
||||
disp_vis = disp_raw.copy()
|
||||
disp_vis[~valid] = 0.0
|
||||
|
||||
if np.count_nonzero(valid) > 20:
|
||||
vals = disp_vis[valid]
|
||||
p2 = np.percentile(vals, 2)
|
||||
p98 = np.percentile(vals, 98)
|
||||
disp_norm = (disp_vis - p2) / (p98 - p2 + 1e-6)
|
||||
disp_norm = np.clip(disp_norm, 0.0, 1.0)
|
||||
else:
|
||||
disp_norm = np.zeros_like(disp_vis, dtype=np.float32)
|
||||
|
||||
disp_color = cv2.applyColorMap((disp_norm * 255).astype(np.uint8), cv2.COLORMAP_TURBO)
|
||||
|
||||
valid_mask = np.zeros_like(disp_color)
|
||||
valid_mask[valid] = (255, 255, 255)
|
||||
|
||||
stats = {
|
||||
"num_disp": num_disp,
|
||||
"block_size": block_size,
|
||||
"valid_pct": float(np.mean(valid) * 100.0),
|
||||
"disp_p05": float(np.percentile(disp_vis[valid], 5)) if np.count_nonzero(valid) > 20 else 0.0,
|
||||
"disp_p50": float(np.percentile(disp_vis[valid], 50)) if np.count_nonzero(valid) > 20 else 0.0,
|
||||
"disp_p95": float(np.percentile(disp_vis[valid], 95)) if np.count_nonzero(valid) > 20 else 0.0,
|
||||
}
|
||||
|
||||
return disp_raw, disp_color, valid_mask, stats
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Visualization
|
||||
# ============================================================
|
||||
|
||||
def draw_epipolar_lines(img_bgr, step=40):
|
||||
out = img_bgr.copy()
|
||||
h, w = out.shape[:2]
|
||||
|
||||
for y in range(0, h, step):
|
||||
color = (0, 255, 255) if (y // step) % 2 == 0 else (255, 255, 0)
|
||||
cv2.line(out, (0, y), (w, y), color, 1, cv2.LINE_AA)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def resize_to_height(img, target_h):
|
||||
h, w = img.shape[:2]
|
||||
if h == target_h:
|
||||
return img
|
||||
|
||||
scale = target_h / h
|
||||
new_w = max(1, int(w * scale))
|
||||
return cv2.resize(img, (new_w, target_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
|
||||
def draw_header(canvas, lines):
|
||||
header_h = 24 + 24 * len(lines)
|
||||
cv2.rectangle(canvas, (0, 0), (canvas.shape[1], header_h), (0, 0, 0), -1)
|
||||
|
||||
y = 24
|
||||
for line in lines:
|
||||
cv2.putText(
|
||||
canvas,
|
||||
line,
|
||||
(12, y),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.55,
|
||||
(255, 255, 255),
|
||||
1,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
y += 24
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def put_label(img, text, color=(255, 255, 255)):
|
||||
out = img.copy()
|
||||
cv2.rectangle(out, (0, 0), (out.shape[1], 34), (0, 0, 0), -1)
|
||||
cv2.putText(out, text, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.65, color, 1, cv2.LINE_AA)
|
||||
return out
|
||||
|
||||
|
||||
def compose_three(left_bgr, center_bgr, right_bgr, title_lines, args):
|
||||
if args.lines:
|
||||
left_bgr = draw_epipolar_lines(left_bgr, step=args.line_step)
|
||||
center_bgr = draw_epipolar_lines(center_bgr, step=args.line_step)
|
||||
right_bgr = draw_epipolar_lines(right_bgr, step=args.line_step)
|
||||
|
||||
left = resize_to_height(left_bgr, args.view_h)
|
||||
center = resize_to_height(center_bgr, args.view_h)
|
||||
right = resize_to_height(right_bgr, args.view_h)
|
||||
|
||||
h = min(left.shape[0], center.shape[0], right.shape[0])
|
||||
left = left[:h]
|
||||
center = center[:h]
|
||||
right = right[:h]
|
||||
|
||||
canvas = np.hstack([left, center, right])
|
||||
return draw_header(canvas, title_lines)
|
||||
|
||||
|
||||
def compose_four(a_bgr, b_bgr, c_bgr, d_bgr, title_lines, args):
|
||||
if args.lines:
|
||||
a_bgr = draw_epipolar_lines(a_bgr, step=args.line_step)
|
||||
b_bgr = draw_epipolar_lines(b_bgr, step=args.line_step)
|
||||
c_bgr = draw_epipolar_lines(c_bgr, step=args.line_step)
|
||||
d_bgr = draw_epipolar_lines(d_bgr, step=args.line_step)
|
||||
|
||||
imgs = [resize_to_height(x, args.view_h) for x in [a_bgr, b_bgr, c_bgr, d_bgr]]
|
||||
h = min(x.shape[0] for x in imgs)
|
||||
imgs = [x[:h] for x in imgs]
|
||||
canvas = np.hstack(imgs)
|
||||
return draw_header(canvas, title_lines)
|
||||
|
||||
|
||||
def absdiff_bgr(a, b):
|
||||
ag = to_gray_u8(a)
|
||||
bg = to_gray_u8(b)
|
||||
diff = cv2.absdiff(ag, bg)
|
||||
return cv2.cvtColor(diff, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
|
||||
def make_overlay(base_bgr, layer_bgr, alpha=0.45):
|
||||
layer = cv2.resize(layer_bgr, (base_bgr.shape[1], base_bgr.shape[0]), interpolation=cv2.INTER_AREA)
|
||||
return cv2.addWeighted(base_bgr, 1.0 - alpha, layer, alpha, 0)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Frame loading and modes
|
||||
# ============================================================
|
||||
|
||||
def load_triplet_images(item, calib, args):
|
||||
rgb_cam = calib["rgb_cam"]
|
||||
nir_cam = calib["nir_cam"]
|
||||
re_cam = calib["re_cam"]
|
||||
|
||||
rgb = read_raw10_rgb_view_bgr(
|
||||
item[rgb_cam],
|
||||
args.width,
|
||||
args.height,
|
||||
args.rgb_bayer,
|
||||
args.rgb_view,
|
||||
use_clahe=not args.no_clahe,
|
||||
)
|
||||
nir = read_raw10_mono_bgr(item[nir_cam], args.width, args.height, use_clahe=not args.no_clahe)
|
||||
re = read_raw10_mono_bgr(item[re_cam], args.width, args.height, use_clahe=not args.no_clahe)
|
||||
|
||||
return {
|
||||
rgb_cam: rgb,
|
||||
nir_cam: nir,
|
||||
re_cam: re,
|
||||
}
|
||||
|
||||
|
||||
def build_views(images, calib, args):
|
||||
rgb_cam = calib["rgb_cam"]
|
||||
nir_cam = calib["nir_cam"]
|
||||
re_cam = calib["re_cam"]
|
||||
|
||||
rgb = images[rgb_cam]
|
||||
nir = images[nir_cam]
|
||||
re = images[re_cam]
|
||||
|
||||
# RGB <-> NIR
|
||||
rgb_ab, nir_ab = rectify_pair_from_calib(rgb, nir, calib, rgb_cam, nir_cam)
|
||||
|
||||
# RGB <-> RE
|
||||
rgb_ac, re_ac = rectify_pair_from_calib(rgb, re, calib, rgb_cam, re_cam)
|
||||
|
||||
# RE <-> NIR
|
||||
re_cb, nir_cb = rectify_pair_from_calib(re, nir, calib, re_cam, nir_cam)
|
||||
|
||||
_, disp_color, valid_mask, disp_stats = compute_disparity(re_cb, nir_cb, args)
|
||||
|
||||
return {
|
||||
"rgb_native": rgb,
|
||||
"nir_native": nir,
|
||||
"re_native": re,
|
||||
"rgb_ab": rgb_ab,
|
||||
"nir_ab": nir_ab,
|
||||
"rgb_ac": rgb_ac,
|
||||
"re_ac": re_ac,
|
||||
"re_cb": re_cb,
|
||||
"nir_cb": nir_cb,
|
||||
"disp_color": disp_color,
|
||||
"valid_mask": valid_mask,
|
||||
"disp_stats": disp_stats,
|
||||
}
|
||||
|
||||
|
||||
def compose_mode(views, item, idx, total, mode, calib, args):
|
||||
rgb_cam = calib["rgb_cam"]
|
||||
nir_cam = calib["nir_cam"]
|
||||
re_cam = calib["re_cam"]
|
||||
|
||||
rgb_name = item[rgb_cam].name
|
||||
|
||||
if mode == "triple_native":
|
||||
left = put_label(views["re_native"], "RE native")
|
||||
center = put_label(views["rgb_native"], "RGB native REF")
|
||||
right = put_label(views["nir_native"], "NIR native")
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | {rgb_name}",
|
||||
"mode=triple_native | sem retificação, RGB no centro",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | S salvar | Q sair",
|
||||
]
|
||||
return compose_three(left, center, right, lines, args)
|
||||
|
||||
if mode == "rgb_re":
|
||||
left = put_label(views["re_ac"], "RE rectificado no par RGB-RE")
|
||||
center = put_label(views["rgb_ac"], "RGB rectificado no par RGB-RE")
|
||||
right = absdiff_bgr(views["rgb_ac"], views["re_ac"])
|
||||
right = put_label(right, "diff RGB-RE")
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | {rgb_name}",
|
||||
"mode=rgb_re | valida alinhamento epipolar RGB<->RE",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | S salvar | Q sair",
|
||||
]
|
||||
return compose_three(left, center, right, lines, args)
|
||||
|
||||
if mode == "rgb_nir":
|
||||
left = absdiff_bgr(views["rgb_ab"], views["nir_ab"])
|
||||
left = put_label(left, "diff RGB-NIR")
|
||||
center = put_label(views["rgb_ab"], "RGB rectificado no par RGB-NIR")
|
||||
right = put_label(views["nir_ab"], "NIR rectificado no par RGB-NIR")
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | {rgb_name}",
|
||||
"mode=rgb_nir | valida alinhamento epipolar RGB<->NIR",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | S salvar | Q sair",
|
||||
]
|
||||
return compose_three(left, center, right, lines, args)
|
||||
|
||||
if mode == "re_nir":
|
||||
left = put_label(views["re_cb"], "RE rectificado no par RE-NIR")
|
||||
center = put_label(views["disp_color"], "disparity RE-NIR")
|
||||
right = put_label(views["nir_cb"], "NIR rectificado no par RE-NIR")
|
||||
st = views["disp_stats"]
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | {rgb_name}",
|
||||
f"mode=re_nir | valid={st['valid_pct']:.1f}% | disp p05={st['disp_p05']:.2f} p50={st['disp_p50']:.2f} p95={st['disp_p95']:.2f}",
|
||||
f"numDisp={st['num_disp']} | block={st['block_size']} | linhas={args.lines}",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | [ ] numDisp | - + block | S salvar | Q sair",
|
||||
]
|
||||
return compose_three(left, center, right, lines, args)
|
||||
|
||||
if mode == "quad_pairs":
|
||||
a = put_label(views["re_ac"], "RE em RGB-RE")
|
||||
b = put_label(views["rgb_ac"], "RGB em RGB-RE")
|
||||
c = put_label(views["rgb_ab"], "RGB em RGB-NIR")
|
||||
d = put_label(views["nir_ab"], "NIR em RGB-NIR")
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | {rgb_name}",
|
||||
"mode=quad_pairs | mostra os dois mundos retificados que usam RGB",
|
||||
"N/SPACE prox | A ant | M modo | L linhas | S salvar | Q sair",
|
||||
]
|
||||
return compose_four(a, b, c, d, lines, args)
|
||||
|
||||
raise RuntimeError(f"Modo desconhecido: {mode}")
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Main
|
||||
# ============================================================
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument("--root_dir", required=True)
|
||||
parser.add_argument("--calib_path", required=True)
|
||||
|
||||
parser.add_argument("--width", type=int, default=1280)
|
||||
parser.add_argument("--height", type=int, default=800)
|
||||
parser.add_argument("--rgb_bayer", default="BGGR", help="Use o mesmo padrão que funcionou na calibração. Ex: BGGR ou RGGB")
|
||||
parser.add_argument(
|
||||
"--rgb_view",
|
||||
default="gray",
|
||||
choices=["color", "gray", "raw_bayer_gray"],
|
||||
help="Como mostrar/processar CAM_A no viewer. gray replica melhor a calibração.",
|
||||
)
|
||||
|
||||
parser.add_argument("--view_h", type=int, default=420)
|
||||
parser.add_argument("--no_clahe", action="store_true")
|
||||
parser.add_argument("--lines", action="store_true")
|
||||
parser.add_argument("--line_step", type=int, default=40)
|
||||
|
||||
parser.add_argument("--num_disp", type=int, default=128)
|
||||
parser.add_argument("--block_size", type=int, default=7)
|
||||
parser.add_argument("--min_disp", type=int, default=0)
|
||||
parser.add_argument("--min_valid_disp", type=float, default=1.0)
|
||||
parser.add_argument("--uniqueness", type=int, default=8)
|
||||
parser.add_argument("--speckle_window", type=int, default=80)
|
||||
parser.add_argument("--speckle_range", type=int, default=2)
|
||||
|
||||
parser.add_argument("--save_dir", default="multicam_rectified_viewer_saves")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
root_dir = Path(args.root_dir)
|
||||
calib_path = Path(args.calib_path)
|
||||
save_dir = Path(args.save_dir)
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
calib = load_multicam_calib(calib_path)
|
||||
|
||||
rgb_cam = calib["rgb_cam"]
|
||||
nir_cam = calib["nir_cam"]
|
||||
re_cam = calib["re_cam"]
|
||||
cams = [rgb_cam, nir_cam, re_cam]
|
||||
|
||||
print(f"[INFO] calib_path: {calib_path}")
|
||||
print(f"[INFO] image_size: {calib['image_size']}")
|
||||
print(f"[INFO] rgb_cam={rgb_cam} nir_cam={nir_cam} re_cam={re_cam} ref_cam={calib['ref_cam']}")
|
||||
|
||||
for k in calib["keys"]:
|
||||
if k.startswith("pair_") and k.endswith("_rms"):
|
||||
print(f"[INFO] {k}: {float(calib['data'][k]):.6f}")
|
||||
if k.startswith("pair_") and k.endswith("_T"):
|
||||
print(f"[INFO] {k}: {np.array(calib['data'][k]).ravel()}")
|
||||
|
||||
triplets, by_cam = find_triplets(root_dir, cams)
|
||||
|
||||
for cam in cams:
|
||||
print(f"[INFO] arquivos {cam}: {len(by_cam[cam])}")
|
||||
print(f"[INFO] triplets: {len(triplets)}")
|
||||
|
||||
if not triplets:
|
||||
raise RuntimeError("Nenhum triplet CAM_A/CAM_B/CAM_C encontrado.")
|
||||
|
||||
modes = ["triple_native", "rgb_re", "rgb_nir", "re_nir", "quad_pairs"]
|
||||
mode_idx = 0
|
||||
|
||||
idx = 0
|
||||
cached_key = None
|
||||
cached_views = None
|
||||
|
||||
cv2.namedWindow("Multicam RAW10 Rectified Viewer", cv2.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
item = triplets[idx]
|
||||
mode = modes[mode_idx]
|
||||
|
||||
key_cache = (
|
||||
tuple(str(item[cam]) for cam in cams),
|
||||
args.width,
|
||||
args.height,
|
||||
args.rgb_bayer,
|
||||
args.rgb_view,
|
||||
args.no_clahe,
|
||||
str(calib_path),
|
||||
args.num_disp,
|
||||
args.block_size,
|
||||
args.min_disp,
|
||||
args.min_valid_disp,
|
||||
args.uniqueness,
|
||||
args.speckle_window,
|
||||
args.speckle_range,
|
||||
)
|
||||
|
||||
if key_cache != cached_key:
|
||||
print(f"[RUN] {idx + 1}/{len(triplets)} - {item[rgb_cam].name}")
|
||||
|
||||
try:
|
||||
images = load_triplet_images(item, calib, args)
|
||||
cached_views = build_views(images, calib, args)
|
||||
cached_key = key_cache
|
||||
except Exception as e:
|
||||
print("[ERRO] Falha processando triplet:")
|
||||
for cam in cams:
|
||||
print(f" {cam}={item[cam]}")
|
||||
print(f" erro={e}")
|
||||
idx = min(idx + 1, len(triplets) - 1)
|
||||
cached_key = None
|
||||
cached_views = None
|
||||
continue
|
||||
|
||||
view = compose_mode(
|
||||
views=cached_views,
|
||||
item=item,
|
||||
idx=idx,
|
||||
total=len(triplets),
|
||||
mode=mode,
|
||||
calib=calib,
|
||||
args=args,
|
||||
)
|
||||
|
||||
cv2.imshow("Multicam RAW10 Rectified Viewer", view)
|
||||
key = cv2.waitKey(0) & 0xFF
|
||||
|
||||
if key in [27, ord("q"), ord("Q")]:
|
||||
break
|
||||
|
||||
elif key in [ord("n"), ord("N"), 32]:
|
||||
idx = min(idx + 1, len(triplets) - 1)
|
||||
cached_key = None
|
||||
|
||||
elif key in [ord("a"), ord("A")]:
|
||||
idx = max(idx - 1, 0)
|
||||
cached_key = None
|
||||
|
||||
elif key in [ord("m"), ord("M")]:
|
||||
mode_idx = (mode_idx + 1) % len(modes)
|
||||
print(f"[PARAM] mode={modes[mode_idx]}")
|
||||
|
||||
elif key in [ord("l"), ord("L")]:
|
||||
args.lines = not args.lines
|
||||
print(f"[PARAM] lines={args.lines}")
|
||||
|
||||
elif key == ord("["):
|
||||
args.num_disp = max(16, args.num_disp - 16)
|
||||
cached_key = None
|
||||
print(f"[PARAM] num_disp={args.num_disp}")
|
||||
|
||||
elif key == ord("]"):
|
||||
args.num_disp = min(512, args.num_disp + 16)
|
||||
cached_key = None
|
||||
print(f"[PARAM] num_disp={args.num_disp}")
|
||||
|
||||
elif key in [ord("-"), ord("_")]:
|
||||
args.block_size = max(3, args.block_size - 2)
|
||||
if args.block_size % 2 == 0:
|
||||
args.block_size -= 1
|
||||
cached_key = None
|
||||
print(f"[PARAM] block_size={args.block_size}")
|
||||
|
||||
elif key in [ord("+"), ord("=")]:
|
||||
args.block_size = min(31, args.block_size + 2)
|
||||
if args.block_size % 2 == 0:
|
||||
args.block_size += 1
|
||||
cached_key = None
|
||||
print(f"[PARAM] block_size={args.block_size}")
|
||||
|
||||
elif key in [ord("s"), ord("S")]:
|
||||
out_path = save_dir / f"multicam_{idx:04d}_{mode}.png"
|
||||
cv2.imwrite(str(out_path), view)
|
||||
print(f"[SAVE] {out_path}")
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,638 @@
|
|||
import argparse
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
IMG_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# RAW10 unpack
|
||||
# ============================================================
|
||||
|
||||
def unpack_raw10_packed(raw: bytes, width: int, height: int) -> np.ndarray:
|
||||
"""
|
||||
Desempacota RAW10 packed padrão:
|
||||
a cada 5 bytes = 4 pixels de 10 bits.
|
||||
|
||||
Layout comum:
|
||||
b0 = p0[9:2]
|
||||
b1 = p1[9:2]
|
||||
b2 = p2[9:2]
|
||||
b3 = p3[9:2]
|
||||
b4 = p0[1:0] | p1[1:0]<<2 | p2[1:0]<<4 | p3[1:0]<<6
|
||||
|
||||
Retorna uint16 HxW com valores 0..1023.
|
||||
"""
|
||||
arr = np.frombuffer(raw, dtype=np.uint8)
|
||||
|
||||
expected_groups = (width * height) // 4
|
||||
expected_bytes = expected_groups * 5
|
||||
|
||||
if arr.size < expected_bytes:
|
||||
raise RuntimeError(
|
||||
f"RAW10 menor que esperado. bytes={arr.size}, esperado={expected_bytes}, "
|
||||
f"width={width}, height={height}"
|
||||
)
|
||||
|
||||
arr = arr[:expected_bytes]
|
||||
groups = arr.reshape(-1, 5).astype(np.uint16)
|
||||
|
||||
p0 = (groups[:, 0] << 2) | ((groups[:, 4] >> 0) & 0x03)
|
||||
p1 = (groups[:, 1] << 2) | ((groups[:, 4] >> 2) & 0x03)
|
||||
p2 = (groups[:, 2] << 2) | ((groups[:, 4] >> 4) & 0x03)
|
||||
p3 = (groups[:, 3] << 2) | ((groups[:, 4] >> 6) & 0x03)
|
||||
|
||||
out = np.empty(groups.shape[0] * 4, dtype=np.uint16)
|
||||
out[0::4] = p0
|
||||
out[1::4] = p1
|
||||
out[2::4] = p2
|
||||
out[3::4] = p3
|
||||
|
||||
out = out.reshape(height, width)
|
||||
return out
|
||||
|
||||
|
||||
def normalize_to_u8(img: np.ndarray, p_low: float = 1.0, p_high: float = 99.0) -> np.ndarray:
|
||||
arr = img.astype(np.float32)
|
||||
valid = np.isfinite(arr)
|
||||
|
||||
if np.count_nonzero(valid) < 20:
|
||||
return np.zeros(arr.shape[:2], dtype=np.uint8)
|
||||
|
||||
vals = arr[valid]
|
||||
lo = np.percentile(vals, p_low)
|
||||
hi = np.percentile(vals, p_high)
|
||||
|
||||
out = (arr - lo) / (hi - lo + 1e-6)
|
||||
out = np.clip(out, 0.0, 1.0)
|
||||
return (out * 255).astype(np.uint8)
|
||||
|
||||
|
||||
def raw10_bin_to_gray_u8(path: Path, width: int, height: int, clahe: bool = True) -> tuple[np.ndarray, np.ndarray]:
|
||||
raw = path.read_bytes()
|
||||
mono10 = unpack_raw10_packed(raw, width=width, height=height)
|
||||
|
||||
gray_u8 = normalize_to_u8(mono10)
|
||||
|
||||
if clahe:
|
||||
eq = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
gray_u8 = eq.apply(gray_u8)
|
||||
|
||||
bgr = cv2.cvtColor(gray_u8, cv2.COLOR_GRAY2BGR)
|
||||
return gray_u8, bgr
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Metadata
|
||||
# ============================================================
|
||||
|
||||
def find_meta_for_bin(bin_path: Path) -> Path | None:
|
||||
"""
|
||||
Procura meta.json na pasta do .bin ou nas pastas acima próximas.
|
||||
"""
|
||||
candidates = [
|
||||
bin_path.parent / "meta.json",
|
||||
bin_path.parent / "metadata.json",
|
||||
bin_path.parent.parent / "meta.json",
|
||||
bin_path.parent.parent / "metadata.json",
|
||||
]
|
||||
|
||||
for c in candidates:
|
||||
if c.exists():
|
||||
return c
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def extract_camera_info_from_meta(meta: dict, cam_key: str):
|
||||
"""
|
||||
Tenta extrair width/height/raw_format para CAM_B ou CAM_C em vários formatos de meta.
|
||||
"""
|
||||
# Caso meta["camera_info"]["CAM_B"]
|
||||
for root_key in ["camera_info", "cameras", "camera_meta", "payload_sources_info"]:
|
||||
root = meta.get(root_key)
|
||||
if isinstance(root, dict):
|
||||
info = root.get(cam_key)
|
||||
if isinstance(info, dict):
|
||||
return info
|
||||
|
||||
# Caso meta tenha lista de câmeras
|
||||
for root_key in ["camera_info", "cameras", "sources"]:
|
||||
root = meta.get(root_key)
|
||||
if isinstance(root, list):
|
||||
for item in root:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
name = item.get("camera") or item.get("name") or item.get("id") or item.get("socket")
|
||||
if name == cam_key:
|
||||
return item
|
||||
|
||||
# Fallback: procura recursivamente dicionário que tenha CAM_B/C
|
||||
stack = [meta]
|
||||
while stack:
|
||||
obj = stack.pop()
|
||||
if isinstance(obj, dict):
|
||||
if cam_key in obj and isinstance(obj[cam_key], dict):
|
||||
return obj[cam_key]
|
||||
for v in obj.values():
|
||||
if isinstance(v, (dict, list)):
|
||||
stack.append(v)
|
||||
elif isinstance(obj, list):
|
||||
for v in obj:
|
||||
if isinstance(v, (dict, list)):
|
||||
stack.append(v)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def parse_width_height_from_meta(meta_path: Path, cam_key: str):
|
||||
try:
|
||||
meta = json.loads(meta_path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
meta = json.loads(meta_path.read_text(encoding="latin-1"))
|
||||
|
||||
info = extract_camera_info_from_meta(meta, cam_key)
|
||||
|
||||
if not info:
|
||||
return None, None
|
||||
|
||||
width = (
|
||||
info.get("width")
|
||||
or info.get("w")
|
||||
or info.get("sensor_width")
|
||||
or info.get("frame_width")
|
||||
)
|
||||
|
||||
height = (
|
||||
info.get("height")
|
||||
or info.get("h")
|
||||
or info.get("sensor_height")
|
||||
or info.get("frame_height")
|
||||
)
|
||||
|
||||
if width is None or height is None:
|
||||
return None, None
|
||||
|
||||
return int(width), int(height)
|
||||
|
||||
|
||||
def resolve_width_height(bin_path: Path, cam_key: str, args):
|
||||
if args.width > 0 and args.height > 0:
|
||||
return args.width, args.height
|
||||
|
||||
meta_path = find_meta_for_bin(bin_path)
|
||||
if meta_path:
|
||||
w, h = parse_width_height_from_meta(meta_path, cam_key)
|
||||
if w and h:
|
||||
return w, h
|
||||
|
||||
raise RuntimeError(
|
||||
f"Não consegui descobrir width/height para {bin_path}. "
|
||||
f"Passe manualmente: --width 1280 --height 800"
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Pairing CAM_B/C
|
||||
# ============================================================
|
||||
|
||||
def clean_stem_for_pair(path: Path, cam_key: str):
|
||||
"""
|
||||
Remove CAM_B/C do nome para tentar parear frames.
|
||||
"""
|
||||
stem = path.stem
|
||||
s = stem
|
||||
|
||||
patterns = [
|
||||
cam_key,
|
||||
cam_key.lower(),
|
||||
cam_key.replace("_", ""),
|
||||
cam_key.replace("_", "").lower(),
|
||||
]
|
||||
|
||||
for p in patterns:
|
||||
s = s.replace(p, "")
|
||||
|
||||
s = re.sub(r"[_\-\s]+", "_", s).strip("_").lower()
|
||||
return s
|
||||
|
||||
|
||||
def find_cam_bins(root_dir: Path, cam_key: str):
|
||||
bins = []
|
||||
for p in root_dir.rglob("*.bin"):
|
||||
name = p.name.lower()
|
||||
if cam_key.lower() in name:
|
||||
bins.append(p)
|
||||
return sorted(bins)
|
||||
|
||||
|
||||
def pair_cam_bins(root_dir: Path, left_cam: str, right_cam: str):
|
||||
left_bins = find_cam_bins(root_dir, left_cam)
|
||||
right_bins = find_cam_bins(root_dir, right_cam)
|
||||
|
||||
right_by_folder_and_key = {}
|
||||
right_by_key = {}
|
||||
|
||||
for rp in right_bins:
|
||||
key = clean_stem_for_pair(rp, right_cam)
|
||||
right_by_folder_and_key[(rp.parent, key)] = rp
|
||||
right_by_key.setdefault(key, rp)
|
||||
|
||||
pairs = []
|
||||
|
||||
for lp in left_bins:
|
||||
key = clean_stem_for_pair(lp, left_cam)
|
||||
|
||||
rp = right_by_folder_and_key.get((lp.parent, key))
|
||||
if rp is None:
|
||||
rp = right_by_key.get(key)
|
||||
|
||||
# Fallback comum: CAM_B/C dentro da mesma pasta, mas nomes não batem
|
||||
if rp is None:
|
||||
candidates_same_folder = [r for r in right_bins if r.parent == lp.parent]
|
||||
if len(candidates_same_folder) == 1:
|
||||
rp = candidates_same_folder[0]
|
||||
|
||||
if rp is not None:
|
||||
pairs.append((lp, rp))
|
||||
|
||||
if pairs:
|
||||
return pairs, left_bins, right_bins, "cam-key"
|
||||
|
||||
# Fallback por ordem
|
||||
n = min(len(left_bins), len(right_bins))
|
||||
pairs = list(zip(left_bins[:n], right_bins[:n]))
|
||||
return pairs, left_bins, right_bins, "order"
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Stereo SGBM
|
||||
# ============================================================
|
||||
|
||||
def ensure_same_size(left, right):
|
||||
h = min(left.shape[0], right.shape[0])
|
||||
w = min(left.shape[1], right.shape[1])
|
||||
|
||||
left2 = cv2.resize(left, (w, h), interpolation=cv2.INTER_AREA)
|
||||
right2 = cv2.resize(right, (w, h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
return left2, right2
|
||||
|
||||
|
||||
def make_sgbm(num_disp, block_size, min_disp=0, uniqueness=8, speckle_window=80, speckle_range=2):
|
||||
num_disp = max(16, int(round(num_disp / 16)) * 16)
|
||||
|
||||
block_size = max(3, int(block_size))
|
||||
if block_size % 2 == 0:
|
||||
block_size += 1
|
||||
|
||||
matcher = cv2.StereoSGBM_create(
|
||||
minDisparity=min_disp,
|
||||
numDisparities=num_disp,
|
||||
blockSize=block_size,
|
||||
P1=8 * block_size * block_size,
|
||||
P2=32 * block_size * block_size,
|
||||
disp12MaxDiff=1,
|
||||
uniquenessRatio=uniqueness,
|
||||
speckleWindowSize=speckle_window,
|
||||
speckleRange=speckle_range,
|
||||
preFilterCap=63,
|
||||
mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY,
|
||||
)
|
||||
|
||||
return matcher, num_disp, block_size
|
||||
|
||||
|
||||
def compute_disparity(left_gray, right_gray, args):
|
||||
matcher, num_disp, block_size = make_sgbm(
|
||||
num_disp=args.num_disp,
|
||||
block_size=args.block_size,
|
||||
min_disp=args.min_disp,
|
||||
uniqueness=args.uniqueness,
|
||||
speckle_window=args.speckle_window,
|
||||
speckle_range=args.speckle_range,
|
||||
)
|
||||
|
||||
disp_raw = matcher.compute(left_gray, right_gray).astype(np.float32) / 16.0
|
||||
|
||||
valid = disp_raw > args.min_valid_disp
|
||||
|
||||
disp_vis = disp_raw.copy()
|
||||
disp_vis[~valid] = 0.0
|
||||
|
||||
if np.count_nonzero(valid) > 20:
|
||||
vals = disp_vis[valid]
|
||||
p2 = np.percentile(vals, 2)
|
||||
p98 = np.percentile(vals, 98)
|
||||
disp_norm = (disp_vis - p2) / (p98 - p2 + 1e-6)
|
||||
disp_norm = np.clip(disp_norm, 0.0, 1.0)
|
||||
else:
|
||||
disp_norm = np.zeros_like(disp_vis, dtype=np.float32)
|
||||
|
||||
disp_color = cv2.applyColorMap((disp_norm * 255).astype(np.uint8), cv2.COLORMAP_TURBO)
|
||||
|
||||
valid_mask = np.zeros_like(disp_color)
|
||||
valid_mask[valid] = (255, 255, 255)
|
||||
|
||||
stats = {
|
||||
"num_disp": num_disp,
|
||||
"block_size": block_size,
|
||||
"valid_pct": float(np.mean(valid) * 100.0),
|
||||
"disp_p05": float(np.percentile(disp_vis[valid], 5)) if np.count_nonzero(valid) > 20 else 0.0,
|
||||
"disp_p50": float(np.percentile(disp_vis[valid], 50)) if np.count_nonzero(valid) > 20 else 0.0,
|
||||
"disp_p95": float(np.percentile(disp_vis[valid], 95)) if np.count_nonzero(valid) > 20 else 0.0,
|
||||
}
|
||||
|
||||
return disp_raw, disp_color, valid_mask, valid, stats
|
||||
|
||||
|
||||
# ============================================================
|
||||
# View
|
||||
# ============================================================
|
||||
|
||||
def resize_to_height(img, target_h):
|
||||
h, w = img.shape[:2]
|
||||
if h == target_h:
|
||||
return img
|
||||
|
||||
scale = target_h / h
|
||||
new_w = max(1, int(w * scale))
|
||||
return cv2.resize(img, (new_w, target_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
|
||||
def draw_header(canvas, lines):
|
||||
header_h = 24 + 24 * len(lines)
|
||||
cv2.rectangle(canvas, (0, 0), (canvas.shape[1], header_h), (0, 0, 0), -1)
|
||||
|
||||
y = 24
|
||||
for line in lines:
|
||||
cv2.putText(
|
||||
canvas,
|
||||
line,
|
||||
(12, y),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.55,
|
||||
(255, 255, 255),
|
||||
1,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
y += 24
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def make_overlay(left_bgr, disp_color, alpha=0.45):
|
||||
disp_resized = cv2.resize(
|
||||
disp_color,
|
||||
(left_bgr.shape[1], left_bgr.shape[0]),
|
||||
interpolation=cv2.INTER_AREA,
|
||||
)
|
||||
return cv2.addWeighted(left_bgr, 1.0 - alpha, disp_resized, alpha, 0)
|
||||
|
||||
|
||||
def compose_view(left_bgr, right_bgr, disp_color, valid_mask, pair, idx, total, mode, stats, args):
|
||||
if mode == "disp":
|
||||
third = disp_color
|
||||
mode_name = "disparity"
|
||||
elif mode == "mask":
|
||||
third = valid_mask
|
||||
mode_name = "valid mask"
|
||||
else:
|
||||
third = make_overlay(left_bgr, disp_color)
|
||||
mode_name = "overlay"
|
||||
|
||||
left = resize_to_height(left_bgr, args.view_h)
|
||||
right = resize_to_height(right_bgr, args.view_h)
|
||||
third = resize_to_height(third, args.view_h)
|
||||
|
||||
h = min(left.shape[0], right.shape[0], third.shape[0])
|
||||
left = left[:h]
|
||||
right = right[:h]
|
||||
third = third[:h]
|
||||
|
||||
canvas = np.hstack([left, right, third])
|
||||
|
||||
lp, rp = pair
|
||||
|
||||
lines = [
|
||||
f"{idx + 1}/{total} | L={lp.name} | R={rp.name}",
|
||||
f"mode={mode_name} | valid={stats['valid_pct']:.1f}% | disp p05={stats['disp_p05']:.2f} p50={stats['disp_p50']:.2f} p95={stats['disp_p95']:.2f}",
|
||||
f"numDisp={stats['num_disp']} | block={stats['block_size']} | uniqueness={args.uniqueness}",
|
||||
"N/SPACE prox | A ant | M modo | [ ] numDisp | - + block | S salvar | Q sair",
|
||||
]
|
||||
|
||||
return draw_header(canvas, lines)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Main
|
||||
# ============================================================
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument("--root_dir", required=True, help="Pasta raiz onde estão os .bin e meta.json")
|
||||
parser.add_argument("--left_cam", default="CAM_B", help="Câmera esquerda. Ex: CAM_B")
|
||||
parser.add_argument("--right_cam", default="CAM_C", help="Câmera direita. Ex: CAM_C")
|
||||
|
||||
parser.add_argument("--width", type=int, default=-1, help="Largura manual caso não tenha meta.json")
|
||||
parser.add_argument("--height", type=int, default=-1, help="Altura manual caso não tenha meta.json")
|
||||
|
||||
parser.add_argument("--view_h", type=int, default=480)
|
||||
|
||||
parser.add_argument("--num_disp", type=int, default=128)
|
||||
parser.add_argument("--block_size", type=int, default=7)
|
||||
parser.add_argument("--min_disp", type=int, default=0)
|
||||
parser.add_argument("--min_valid_disp", type=float, default=1.0)
|
||||
|
||||
parser.add_argument("--uniqueness", type=int, default=8)
|
||||
parser.add_argument("--speckle_window", type=int, default=80)
|
||||
parser.add_argument("--speckle_range", type=int, default=2)
|
||||
|
||||
parser.add_argument("--no_clahe", action="store_true")
|
||||
parser.add_argument("--swap", action="store_true", help="Inverte left/right depois do pareamento")
|
||||
parser.add_argument("--save_dir", default="stereo_raw10_sgbm_saves")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
root_dir = Path(args.root_dir)
|
||||
save_dir = Path(args.save_dir)
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
pairs, left_bins, right_bins, pair_mode = pair_cam_bins(
|
||||
root_dir=root_dir,
|
||||
left_cam=args.left_cam,
|
||||
right_cam=args.right_cam,
|
||||
)
|
||||
|
||||
if args.swap:
|
||||
pairs = [(r, l) for l, r in pairs]
|
||||
args.left_cam, args.right_cam = args.right_cam, args.left_cam
|
||||
|
||||
if not left_bins:
|
||||
raise RuntimeError(f"Nenhum .bin encontrado com {args.left_cam} em {root_dir}")
|
||||
|
||||
if not right_bins:
|
||||
raise RuntimeError(f"Nenhum .bin encontrado com {args.right_cam} em {root_dir}")
|
||||
|
||||
if not pairs:
|
||||
raise RuntimeError("Não consegui parear CAM_B/C.")
|
||||
|
||||
print(f"[INFO] root_dir: {root_dir}")
|
||||
print(f"[INFO] left_cam: {args.left_cam} | arquivos: {len(left_bins)}")
|
||||
print(f"[INFO] right_cam: {args.right_cam} | arquivos: {len(right_bins)}")
|
||||
print(f"[INFO] pares: {len(pairs)}")
|
||||
print(f"[INFO] pareamento: {pair_mode}")
|
||||
print("[INFO] controles:")
|
||||
print(" N ou SPACE = próxima")
|
||||
print(" A = anterior")
|
||||
print(" M = modo disparity/mask/overlay")
|
||||
print(" [ / ] = diminui/aumenta numDisparities")
|
||||
print(" - / + = diminui/aumenta blockSize")
|
||||
print(" S = salva visual atual")
|
||||
print(" Q ou ESC = sair")
|
||||
|
||||
idx = 0
|
||||
mode = "disp"
|
||||
|
||||
cached_key = None
|
||||
cached_data = None
|
||||
|
||||
cv2.namedWindow("Stereo RAW10 SGBM Viewer", cv2.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
pair = pairs[idx]
|
||||
lp, rp = pair
|
||||
|
||||
key_cache = (
|
||||
str(lp),
|
||||
str(rp),
|
||||
args.num_disp,
|
||||
args.block_size,
|
||||
args.uniqueness,
|
||||
args.speckle_window,
|
||||
args.speckle_range,
|
||||
args.min_disp,
|
||||
args.min_valid_disp,
|
||||
args.no_clahe,
|
||||
args.width,
|
||||
args.height,
|
||||
)
|
||||
|
||||
if key_cache != cached_key:
|
||||
print(f"[RUN] {idx + 1}/{len(pairs)} - L={lp.name} | R={rp.name}")
|
||||
|
||||
try:
|
||||
lw, lh = resolve_width_height(lp, args.left_cam, args)
|
||||
rw, rh = resolve_width_height(rp, args.right_cam, args)
|
||||
|
||||
left_gray, left_bgr = raw10_bin_to_gray_u8(
|
||||
lp,
|
||||
width=lw,
|
||||
height=lh,
|
||||
clahe=not args.no_clahe,
|
||||
)
|
||||
|
||||
right_gray, right_bgr = raw10_bin_to_gray_u8(
|
||||
rp,
|
||||
width=rw,
|
||||
height=rh,
|
||||
clahe=not args.no_clahe,
|
||||
)
|
||||
|
||||
left_gray, right_gray = ensure_same_size(left_gray, right_gray)
|
||||
left_bgr, right_bgr = ensure_same_size(left_bgr, right_bgr)
|
||||
|
||||
_, disp_color, valid_mask, valid, stats = compute_disparity(left_gray, right_gray, args)
|
||||
|
||||
cached_data = (left_bgr, right_bgr, disp_color, valid_mask, stats)
|
||||
cached_key = key_cache
|
||||
|
||||
except Exception as e:
|
||||
print(f"[ERRO] Falha processando par:")
|
||||
print(f" L={lp}")
|
||||
print(f" R={rp}")
|
||||
print(f" erro={e}")
|
||||
|
||||
idx = min(idx + 1, len(pairs) - 1)
|
||||
cached_key = None
|
||||
cached_data = None
|
||||
continue
|
||||
|
||||
else:
|
||||
left_bgr, right_bgr, disp_color, valid_mask, stats = cached_data
|
||||
|
||||
view = compose_view(
|
||||
left_bgr=left_bgr,
|
||||
right_bgr=right_bgr,
|
||||
disp_color=disp_color,
|
||||
valid_mask=valid_mask,
|
||||
pair=pair,
|
||||
idx=idx,
|
||||
total=len(pairs),
|
||||
mode=mode,
|
||||
stats=stats,
|
||||
args=args,
|
||||
)
|
||||
|
||||
cv2.imshow("Stereo RAW10 SGBM Viewer", view)
|
||||
key = cv2.waitKey(0) & 0xFF
|
||||
|
||||
if key in [27, ord("q"), ord("Q")]:
|
||||
break
|
||||
|
||||
elif key in [ord("n"), ord("N"), 32]:
|
||||
idx = min(idx + 1, len(pairs) - 1)
|
||||
cached_key = None
|
||||
|
||||
elif key in [ord("a"), ord("A")]:
|
||||
idx = max(idx - 1, 0)
|
||||
cached_key = None
|
||||
|
||||
elif key in [ord("m"), ord("M")]:
|
||||
if mode == "disp":
|
||||
mode = "mask"
|
||||
elif mode == "mask":
|
||||
mode = "overlay"
|
||||
else:
|
||||
mode = "disp"
|
||||
|
||||
elif key == ord("["):
|
||||
args.num_disp = max(16, args.num_disp - 16)
|
||||
cached_key = None
|
||||
print(f"[PARAM] num_disp={args.num_disp}")
|
||||
|
||||
elif key == ord("]"):
|
||||
args.num_disp = min(512, args.num_disp + 16)
|
||||
cached_key = None
|
||||
print(f"[PARAM] num_disp={args.num_disp}")
|
||||
|
||||
elif key in [ord("-"), ord("_")]:
|
||||
args.block_size = max(3, args.block_size - 2)
|
||||
if args.block_size % 2 == 0:
|
||||
args.block_size -= 1
|
||||
cached_key = None
|
||||
print(f"[PARAM] block_size={args.block_size}")
|
||||
|
||||
elif key in [ord("+"), ord("=")]:
|
||||
args.block_size = min(31, args.block_size + 2)
|
||||
if args.block_size % 2 == 0:
|
||||
args.block_size += 1
|
||||
cached_key = None
|
||||
print(f"[PARAM] block_size={args.block_size}")
|
||||
|
||||
elif key in [ord("s"), ord("S")]:
|
||||
out_path = save_dir / f"stereo_raw10_{idx:04d}_{mode}.png"
|
||||
cv2.imwrite(str(out_path), view)
|
||||
print(f"[SAVE] {out_path}")
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,573 @@
|
|||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
from PIL import Image
|
||||
|
||||
|
||||
IMG_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
|
||||
|
||||
|
||||
def get_providers(force_cpu: bool = False):
|
||||
available = ort.get_available_providers()
|
||||
|
||||
if force_cpu:
|
||||
return ["CPUExecutionProvider"]
|
||||
|
||||
providers = []
|
||||
|
||||
if "CUDAExecutionProvider" in available:
|
||||
providers.append("CUDAExecutionProvider")
|
||||
|
||||
providers.append("CPUExecutionProvider")
|
||||
return providers
|
||||
|
||||
|
||||
def print_model_io(session: ort.InferenceSession):
|
||||
print("\n[MODEL INPUTS]")
|
||||
for i, inp in enumerate(session.get_inputs()):
|
||||
print(f" {i}: name={inp.name} shape={inp.shape} type={inp.type}")
|
||||
|
||||
print("\n[MODEL OUTPUTS]")
|
||||
for i, out in enumerate(session.get_outputs()):
|
||||
print(f" {i}: name={out.name} shape={out.shape} type={out.type}")
|
||||
print("")
|
||||
|
||||
|
||||
def resolve_hw_from_input_shape(shape, fallback_h: int, fallback_w: int):
|
||||
"""
|
||||
Tenta descobrir H/W do input ONNX.
|
||||
Esperado normalmente: [1, 3, H, W] ou [1, H, W, 3].
|
||||
Se for dinâmico, usa fallback.
|
||||
"""
|
||||
if shape is None:
|
||||
return fallback_h, fallback_w
|
||||
|
||||
dims = list(shape)
|
||||
|
||||
def is_int(x):
|
||||
return isinstance(x, int) and x > 0
|
||||
|
||||
# NCHW
|
||||
if len(dims) == 4 and dims[1] == 3:
|
||||
h = dims[2] if is_int(dims[2]) else fallback_h
|
||||
w = dims[3] if is_int(dims[3]) else fallback_w
|
||||
return int(h), int(w)
|
||||
|
||||
# NHWC
|
||||
if len(dims) == 4 and dims[-1] == 3:
|
||||
h = dims[1] if is_int(dims[1]) else fallback_h
|
||||
w = dims[2] if is_int(dims[2]) else fallback_w
|
||||
return int(h), int(w)
|
||||
|
||||
return fallback_h, fallback_w
|
||||
|
||||
|
||||
def is_nchw_input(shape):
|
||||
if shape is None:
|
||||
return True
|
||||
dims = list(shape)
|
||||
return len(dims) == 4 and dims[1] == 3
|
||||
|
||||
|
||||
def preprocess_image(image_path: Path, input_shape, input_h: int, input_w: int, mean, std):
|
||||
img_pil = Image.open(image_path).convert("RGB")
|
||||
rgb_orig = np.array(img_pil)
|
||||
bgr_orig = cv2.cvtColor(rgb_orig, cv2.COLOR_RGB2BGR)
|
||||
|
||||
rgb_resized = cv2.resize(rgb_orig, (input_w, input_h), interpolation=cv2.INTER_AREA)
|
||||
x = rgb_resized.astype(np.float32) / 255.0
|
||||
|
||||
mean_arr = np.array(mean, dtype=np.float32).reshape(1, 1, 3)
|
||||
std_arr = np.array(std, dtype=np.float32).reshape(1, 1, 3)
|
||||
x = (x - mean_arr) / std_arr
|
||||
|
||||
if is_nchw_input(input_shape):
|
||||
x = np.transpose(x, (2, 0, 1)) # CHW
|
||||
|
||||
x = np.expand_dims(x, axis=0).astype(np.float32)
|
||||
|
||||
return bgr_orig, x
|
||||
|
||||
|
||||
def robust_normalize_for_view(depth_m: np.ndarray, invert: bool = False):
|
||||
"""
|
||||
Só para visualização colorida.
|
||||
O modo bandas usa o depth_m direto.
|
||||
"""
|
||||
d = depth_m.astype(np.float32)
|
||||
valid = np.isfinite(d) & (d > 0)
|
||||
|
||||
if np.count_nonzero(valid) < 20:
|
||||
return np.zeros_like(d, dtype=np.float32)
|
||||
|
||||
vals = d[valid]
|
||||
p2 = np.percentile(vals, 2)
|
||||
p98 = np.percentile(vals, 98)
|
||||
|
||||
dn = (d - p2) / (p98 - p2 + 1e-6)
|
||||
dn = np.clip(dn, 0.0, 1.0)
|
||||
|
||||
if invert:
|
||||
dn = 1.0 - dn
|
||||
|
||||
dn[~valid] = 0.0
|
||||
return dn.astype(np.float32)
|
||||
|
||||
|
||||
def depth_to_colormap(depth_norm: np.ndarray):
|
||||
u8 = (depth_norm * 255).astype(np.uint8)
|
||||
return cv2.applyColorMap(u8, cv2.COLORMAP_TURBO)
|
||||
|
||||
|
||||
def make_metric_bands(depth_m: np.ndarray, near_m: float, mid_m: float, far_m: float):
|
||||
"""
|
||||
Bandas baseadas no valor métrico estimado.
|
||||
|
||||
0: < near_m
|
||||
1: near_m até mid_m
|
||||
2: mid_m até far_m
|
||||
3: >= far_m
|
||||
"""
|
||||
d = depth_m.astype(np.float32)
|
||||
valid = np.isfinite(d) & (d > 0)
|
||||
|
||||
bands = np.zeros_like(d, dtype=np.uint8)
|
||||
|
||||
bands[(d >= near_m) & (d < mid_m)] = 1
|
||||
bands[(d >= mid_m) & (d < far_m)] = 2
|
||||
bands[d >= far_m] = 3
|
||||
bands[~valid] = 255
|
||||
|
||||
out = np.zeros((bands.shape[0], bands.shape[1], 3), dtype=np.uint8)
|
||||
|
||||
# BGR
|
||||
out[bands == 0] = (40, 40, 255) # muito perto
|
||||
out[bands == 1] = (40, 180, 255) # perto/médio
|
||||
out[bands == 2] = (40, 255, 120) # médio/longe
|
||||
out[bands == 3] = (255, 180, 40) # longe
|
||||
out[bands == 255] = (0, 0, 0) # inválido
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def extract_depth_from_outputs(outputs, output_names):
|
||||
"""
|
||||
Tenta achar depth em diferentes formatos:
|
||||
- output chamado depth
|
||||
- output chamado points/xyz, usando canal Z
|
||||
- saída única
|
||||
"""
|
||||
name_to_out = {
|
||||
name.lower(): out
|
||||
for name, out in zip(output_names, outputs)
|
||||
}
|
||||
|
||||
# 1) Procura saída com nome depth
|
||||
for name, out in name_to_out.items():
|
||||
if "depth" in name:
|
||||
return squeeze_depth(out)
|
||||
|
||||
# 2) Procura points/xyz e usa Z
|
||||
for name, out in name_to_out.items():
|
||||
if "point" in name or "xyz" in name:
|
||||
arr = np.array(out)
|
||||
return extract_z_from_points(arr)
|
||||
|
||||
# 3) Se tiver só uma saída, usa ela
|
||||
if len(outputs) == 1:
|
||||
arr = np.array(outputs[0])
|
||||
|
||||
# Se parece points, extrai Z
|
||||
if arr.ndim == 4 and (arr.shape[1] == 3 or arr.shape[-1] == 3):
|
||||
return extract_z_from_points(arr)
|
||||
|
||||
return squeeze_depth(arr)
|
||||
|
||||
# 4) Fallback: pega a primeira saída 2D/3D/4D plausível
|
||||
for out in outputs:
|
||||
arr = np.array(out)
|
||||
try:
|
||||
d = squeeze_depth(arr)
|
||||
if d.ndim == 2:
|
||||
return d
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise RuntimeError("Não consegui identificar o mapa de depth nas saídas ONNX.")
|
||||
|
||||
|
||||
def squeeze_depth(arr):
|
||||
arr = np.array(arr)
|
||||
|
||||
# Remove batch/canal unitário
|
||||
arr = np.squeeze(arr)
|
||||
|
||||
if arr.ndim == 2:
|
||||
return arr.astype(np.float32)
|
||||
|
||||
if arr.ndim == 3:
|
||||
# CHW com 1 canal
|
||||
if arr.shape[0] == 1:
|
||||
return arr[0].astype(np.float32)
|
||||
|
||||
# HWC com 1 canal
|
||||
if arr.shape[-1] == 1:
|
||||
return arr[..., 0].astype(np.float32)
|
||||
|
||||
# Se for 3 canais, talvez seja XYZ
|
||||
if arr.shape[0] == 3:
|
||||
return arr[2].astype(np.float32)
|
||||
|
||||
if arr.shape[-1] == 3:
|
||||
return arr[..., 2].astype(np.float32)
|
||||
|
||||
raise RuntimeError(f"Formato de depth não suportado: shape={arr.shape}")
|
||||
|
||||
|
||||
def extract_z_from_points(arr):
|
||||
arr = np.array(arr)
|
||||
|
||||
# NCHW: [1, 3, H, W]
|
||||
if arr.ndim == 4 and arr.shape[1] == 3:
|
||||
return arr[0, 2].astype(np.float32)
|
||||
|
||||
# NHWC: [1, H, W, 3]
|
||||
if arr.ndim == 4 and arr.shape[-1] == 3:
|
||||
return arr[0, ..., 2].astype(np.float32)
|
||||
|
||||
# CHW: [3, H, W]
|
||||
if arr.ndim == 3 and arr.shape[0] == 3:
|
||||
return arr[2].astype(np.float32)
|
||||
|
||||
# HWC: [H, W, 3]
|
||||
if arr.ndim == 3 and arr.shape[-1] == 3:
|
||||
return arr[..., 2].astype(np.float32)
|
||||
|
||||
raise RuntimeError(f"Formato de points/xyz não suportado: shape={arr.shape}")
|
||||
|
||||
|
||||
def maybe_invert_depth_if_needed(depth: np.ndarray, auto_invert: bool):
|
||||
"""
|
||||
Alguns modelos podem devolver inverso/disparity.
|
||||
Aqui deixei opcional. Por padrão não mexe.
|
||||
"""
|
||||
if not auto_invert:
|
||||
return depth
|
||||
|
||||
d = depth.astype(np.float32)
|
||||
valid = np.isfinite(d) & (d > 1e-6)
|
||||
out = np.zeros_like(d, dtype=np.float32)
|
||||
out[valid] = 1.0 / d[valid]
|
||||
return out
|
||||
|
||||
|
||||
def infer_depth(session, image_path: Path, args):
|
||||
input0 = session.get_inputs()[0]
|
||||
input_shape = input0.shape
|
||||
fallback_h, fallback_w = args.input_h, args.input_w
|
||||
input_h, input_w = resolve_hw_from_input_shape(input_shape, fallback_h, fallback_w)
|
||||
|
||||
preview_bgr, x = preprocess_image(
|
||||
image_path=image_path,
|
||||
input_shape=input_shape,
|
||||
input_h=input_h,
|
||||
input_w=input_w,
|
||||
mean=args.mean,
|
||||
std=args.std,
|
||||
)
|
||||
|
||||
feed = {}
|
||||
|
||||
# Alimenta input principal
|
||||
feed[input0.name] = x
|
||||
|
||||
# Alguns modelos têm input extra de intrinsics/K.
|
||||
# Se existir, manda uma matriz aproximada baseada no tamanho de entrada.
|
||||
# Para nosso teste visual, isso é melhor do que travar.
|
||||
for inp in session.get_inputs()[1:]:
|
||||
name = inp.name.lower()
|
||||
shape = inp.shape
|
||||
|
||||
fx = args.fx if args.fx > 0 else input_w * 0.9
|
||||
fy = args.fy if args.fy > 0 else input_w * 0.9
|
||||
cx = args.cx if args.cx >= 0 else input_w / 2.0
|
||||
cy = args.cy if args.cy >= 0 else input_h / 2.0
|
||||
|
||||
K = np.array(
|
||||
[
|
||||
[fx, 0.0, cx],
|
||||
[0.0, fy, cy],
|
||||
[0.0, 0.0, 1.0],
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
if "k" == name or "intr" in name or "camera" in name:
|
||||
if len(shape) == 3:
|
||||
feed[inp.name] = K[None, ...]
|
||||
else:
|
||||
feed[inp.name] = K
|
||||
else:
|
||||
# Fallback para input extra desconhecido
|
||||
# Evita crash, mas imprime para a gente ajustar se precisar.
|
||||
print(f"[WARN] Input extra desconhecido: {inp.name}, shape={inp.shape}. Enviando zeros.")
|
||||
concrete_shape = []
|
||||
for dim in shape:
|
||||
concrete_shape.append(dim if isinstance(dim, int) and dim > 0 else 1)
|
||||
feed[inp.name] = np.zeros(concrete_shape, dtype=np.float32)
|
||||
|
||||
t0 = time.time()
|
||||
outputs = session.run(None, feed)
|
||||
infer_ms = (time.time() - t0) * 1000.0
|
||||
|
||||
output_names = [o.name for o in session.get_outputs()]
|
||||
depth = extract_depth_from_outputs(outputs, output_names)
|
||||
depth = maybe_invert_depth_if_needed(depth, auto_invert=args.inv_depth)
|
||||
|
||||
# Redimensiona para o preview original
|
||||
if depth.shape[:2] != preview_bgr.shape[:2]:
|
||||
depth = cv2.resize(
|
||||
depth.astype(np.float32),
|
||||
(preview_bgr.shape[1], preview_bgr.shape[0]),
|
||||
interpolation=cv2.INTER_CUBIC,
|
||||
)
|
||||
|
||||
# Remove valores absurdos só para visual/estatística
|
||||
depth = depth.astype(np.float32)
|
||||
depth[~np.isfinite(depth)] = 0.0
|
||||
depth[depth < 0] = 0.0
|
||||
|
||||
return preview_bgr, depth, infer_ms
|
||||
|
||||
|
||||
def calc_stats(depth_m: np.ndarray, infer_ms: float):
|
||||
valid = np.isfinite(depth_m) & (depth_m > 0)
|
||||
|
||||
if np.count_nonzero(valid) < 20:
|
||||
return {
|
||||
"p01": 0.0,
|
||||
"p05": 0.0,
|
||||
"p50": 0.0,
|
||||
"p95": 0.0,
|
||||
"p99": 0.0,
|
||||
"mean": 0.0,
|
||||
"std": 0.0,
|
||||
"infer_ms": infer_ms,
|
||||
"valid_pct": 0.0,
|
||||
}
|
||||
|
||||
vals = depth_m[valid]
|
||||
return {
|
||||
"p01": float(np.percentile(vals, 1)),
|
||||
"p05": float(np.percentile(vals, 5)),
|
||||
"p50": float(np.percentile(vals, 50)),
|
||||
"p95": float(np.percentile(vals, 95)),
|
||||
"p99": float(np.percentile(vals, 99)),
|
||||
"mean": float(np.mean(vals)),
|
||||
"std": float(np.std(vals)),
|
||||
"infer_ms": float(infer_ms),
|
||||
"valid_pct": float(np.mean(valid) * 100.0),
|
||||
}
|
||||
|
||||
|
||||
def resize_to_height(img: np.ndarray, target_h: int):
|
||||
h, w = img.shape[:2]
|
||||
if h == target_h:
|
||||
return img
|
||||
|
||||
scale = target_h / h
|
||||
new_w = max(1, int(w * scale))
|
||||
return cv2.resize(img, (new_w, target_h), interpolation=cv2.INTER_AREA)
|
||||
|
||||
|
||||
def draw_header(canvas, lines):
|
||||
header_h = 24 + 24 * len(lines)
|
||||
cv2.rectangle(canvas, (0, 0), (canvas.shape[1], header_h), (0, 0, 0), -1)
|
||||
|
||||
y = 24
|
||||
for line in lines:
|
||||
cv2.putText(
|
||||
canvas,
|
||||
line,
|
||||
(12, y),
|
||||
cv2.FONT_HERSHEY_SIMPLEX,
|
||||
0.55,
|
||||
(255, 255, 255),
|
||||
1,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
y += 24
|
||||
|
||||
return canvas
|
||||
|
||||
|
||||
def compose_view(preview_bgr, depth_m, image_path, idx, total, mode, invert_view, stats, args):
|
||||
depth_norm = robust_normalize_for_view(depth_m, invert=invert_view)
|
||||
depth_color = depth_to_colormap(depth_norm)
|
||||
bands_color = make_metric_bands(depth_m, args.near_m, args.mid_m, args.far_m)
|
||||
|
||||
right_src = depth_color if mode == "depth" else bands_color
|
||||
|
||||
left = resize_to_height(preview_bgr, args.view_h)
|
||||
right = resize_to_height(right_src, args.view_h)
|
||||
|
||||
h = min(left.shape[0], right.shape[0])
|
||||
left = left[:h]
|
||||
right = right[:h]
|
||||
|
||||
canvas = np.hstack([left, right])
|
||||
|
||||
mode_name = "depth colormap" if mode == "depth" else "metric bands"
|
||||
|
||||
lines = [
|
||||
f"{idx + 1}/{total} - {image_path.name}",
|
||||
f"modo={mode_name} | p05={stats['p05']:.2f}m p50={stats['p50']:.2f}m p95={stats['p95']:.2f}m valid={stats['valid_pct']:.1f}% infer={stats['infer_ms']:.1f}ms",
|
||||
f"bands: <{args.near_m:.2f}m | {args.near_m:.2f}-{args.mid_m:.2f}m | {args.mid_m:.2f}-{args.far_m:.2f}m | >{args.far_m:.2f}m",
|
||||
"N/SPACE prox | A ant | M modo | I inverte visual | S salvar | Q sair",
|
||||
]
|
||||
|
||||
return draw_header(canvas, lines)
|
||||
|
||||
|
||||
def parse_mean_std(text):
|
||||
vals = [float(x.strip()) for x in text.split(",")]
|
||||
if len(vals) != 3:
|
||||
raise argparse.ArgumentTypeError("Use 3 valores separados por vírgula. Ex: 0.485,0.456,0.406")
|
||||
return vals
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input_dir", required=True, help="Pasta com previews PNG/JPG")
|
||||
parser.add_argument("--model_path", required=True, help="Caminho do arquivo .onnx")
|
||||
parser.add_argument("--cpu", action="store_true", help="Força CPUExecutionProvider")
|
||||
|
||||
# Usado se o modelo tiver tamanho dinâmico
|
||||
parser.add_argument("--input_h", type=int, default=384)
|
||||
parser.add_argument("--input_w", type=int, default=512)
|
||||
|
||||
# Normalização padrão ImageNet. Se o repo do ONNX pedir outra, ajustamos.
|
||||
parser.add_argument("--mean", type=parse_mean_std, default=[0.485, 0.456, 0.406])
|
||||
parser.add_argument("--std", type=parse_mean_std, default=[0.229, 0.224, 0.225])
|
||||
|
||||
# Intrínsecos aproximados se o ONNX pedir K/intrinsics
|
||||
parser.add_argument("--fx", type=float, default=-1.0)
|
||||
parser.add_argument("--fy", type=float, default=-1.0)
|
||||
parser.add_argument("--cx", type=float, default=-1.0)
|
||||
parser.add_argument("--cy", type=float, default=-1.0)
|
||||
|
||||
# Se o output for inverso/disparity, ativa isto
|
||||
parser.add_argument("--inv_depth", action="store_true")
|
||||
|
||||
# Bandas métricas para visual
|
||||
parser.add_argument("--near_m", type=float, default=0.6)
|
||||
parser.add_argument("--mid_m", type=float, default=1.1)
|
||||
parser.add_argument("--far_m", type=float, default=1.8)
|
||||
|
||||
parser.add_argument("--view_h", type=int, default=560)
|
||||
parser.add_argument("--save_dir", default="unidepth_onnx_viewer_saves")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
input_dir = Path(args.input_dir)
|
||||
model_path = Path(args.model_path)
|
||||
save_dir = Path(args.save_dir)
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if not model_path.exists():
|
||||
raise FileNotFoundError(f"Modelo ONNX não encontrado: {model_path}")
|
||||
|
||||
image_paths = sorted([
|
||||
p for p in input_dir.rglob("*")
|
||||
if p.suffix.lower() in IMG_EXTS
|
||||
])
|
||||
|
||||
if not image_paths:
|
||||
raise RuntimeError(f"Nenhuma imagem encontrada em: {input_dir}")
|
||||
|
||||
providers = get_providers(force_cpu=args.cpu)
|
||||
|
||||
print(f"[INFO] imagens: {len(image_paths)}")
|
||||
print(f"[INFO] model_path: {model_path}")
|
||||
print(f"[INFO] providers: {providers}")
|
||||
print(f"[INFO] ONNX Runtime providers disponiveis: {ort.get_available_providers()}")
|
||||
|
||||
session = ort.InferenceSession(str(model_path), providers=providers)
|
||||
print_model_io(session)
|
||||
|
||||
idx = 0
|
||||
mode = "depth"
|
||||
invert_view = False
|
||||
|
||||
cached_path = None
|
||||
cached_data = None
|
||||
|
||||
cv2.namedWindow("UniDepth ONNX Metric Viewer", cv2.WINDOW_NORMAL)
|
||||
|
||||
while True:
|
||||
image_path = image_paths[idx]
|
||||
|
||||
if cached_path != image_path or cached_data is None:
|
||||
print(f"[RUN] {idx + 1}/{len(image_paths)} - {image_path.name}")
|
||||
|
||||
try:
|
||||
preview_bgr, depth_m, infer_ms = infer_depth(session, image_path, args)
|
||||
stats = calc_stats(depth_m, infer_ms)
|
||||
cached_data = (preview_bgr, depth_m, stats)
|
||||
cached_path = image_path
|
||||
except Exception as e:
|
||||
print(f"[ERRO] Falha em {image_path}: {e}")
|
||||
idx = min(idx + 1, len(image_paths) - 1)
|
||||
cached_path = None
|
||||
cached_data = None
|
||||
continue
|
||||
else:
|
||||
preview_bgr, depth_m, stats = cached_data
|
||||
|
||||
view = compose_view(
|
||||
preview_bgr=preview_bgr,
|
||||
depth_m=depth_m,
|
||||
image_path=image_path,
|
||||
idx=idx,
|
||||
total=len(image_paths),
|
||||
mode=mode,
|
||||
invert_view=invert_view,
|
||||
stats=stats,
|
||||
args=args,
|
||||
)
|
||||
|
||||
cv2.imshow("UniDepth ONNX Metric Viewer", view)
|
||||
key = cv2.waitKey(0) & 0xFF
|
||||
|
||||
if key in [27, ord("q"), ord("Q")]:
|
||||
break
|
||||
|
||||
elif key in [ord("n"), ord("N"), 32]:
|
||||
idx = min(idx + 1, len(image_paths) - 1)
|
||||
cached_path = None
|
||||
|
||||
elif key in [ord("a"), ord("A")]:
|
||||
idx = max(idx - 1, 0)
|
||||
cached_path = None
|
||||
|
||||
elif key in [ord("m"), ord("M")]:
|
||||
mode = "bands" if mode == "depth" else "depth"
|
||||
|
||||
elif key in [ord("i"), ord("I")]:
|
||||
invert_view = not invert_view
|
||||
print(f"[INFO] invert_view={invert_view}")
|
||||
|
||||
elif key in [ord("s"), ord("S")]:
|
||||
out_path = save_dir / f"{image_path.stem}_unidepth_onnx_{mode}.png"
|
||||
cv2.imwrite(str(out_path), view)
|
||||
print(f"[SAVE] {out_path}")
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in New Issue