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

565 lines
17 KiB
Python
Raw Normal View History

2026-05-27 19:34:48 +00:00
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()