agrobot_base/Python/OAK/datasets/oak-fcc-3/depth_homography_viewer.py

706 lines
23 KiB
Python

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()