"""STAGE 2 (standalone, torch): compare ROI strategies for CREStereo dimensioning. full CRE on the whole frame, FastSAM on the whole frame beltroi FastSAM + CRE restricted to the belt region all 3 rigs share objroi belt ROI bounds the segmentation, CRE runs on the object bbox + pad, L and R sharing ONE column window (left-extended by max disparity) objroi_bad same, but the right crop is re-centred on the right image's own bbox - the control that shows what independent centring costs RGB is never masked before CRE: the matcher needs the surround. The mask is applied to the depth, eroded, only when back-projecting. """ import os, sys, json, time, numpy as np, cv2, torch, torch.nn.functional as F CV = "/home/dasha/isaac_assets/cv" CT = "/home/dasha/robozon-sorter/control_test" CFG = sys.argv[1] if len(sys.argv) > 1 else "E60" # cam_configs.DEFAULT; # not imported here because cam_configs needs omni.usd and this stage runs standalone CAP = f"{CT}/captures/{CFG}" VARIANTS = (sys.argv[2].split(",") if len(sys.argv) > 2 else ["full", "beltroi", "objroi", "objroi_bad"]) man = json.load(open(f"{CAP}/manifest.json")) calib = man["calib"] # the calibration this capture was actually shot with TARGET = np.array(man["target"]); BELT_Z = TARGET[2] RIGS = sorted({c["rig"] for c in calib.values()}) dev = "cuda" os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics" sys.path.insert(0, f"{CV}/crestereo") from nets import Model cre = Model(max_disp=256, mixed_precision=False, test_mode=True) cre.load_state_dict(torch.load(f"{CV}/crestereo/models/crestereo_eth3d.pth", map_location="cpu"), strict=True) cre.to(dev).eval() from ultralytics import FastSAM fsam = FastSAM(f"{CV}/FastSAM-s.pt") ZMIN, DPAD, PAD = 0.60, 1.35, 48 VOX = 0.004 def cre_infer(L, R, iters=20): """two-stage CRE on one already-corresponding pair of crops""" h, w = L.shape[:2] Hp, Wp = (h + 7) // 8 * 8, (w + 7) // 8 * 8 Lb = np.zeros((1, 3, Hp, Wp), np.float32); Rb = np.zeros((1, 3, Hp, Wp), np.float32) Lb[0, :, :h, :w] = L.transpose(2, 0, 1); Rb[0, :, :h, :w] = R.transpose(2, 0, 1) iL = torch.from_numpy(Lb).to(dev); iR = torch.from_numpy(Rb).to(dev) dL = F.interpolate(iL, (Hp // 2, Wp // 2), mode="bilinear", align_corners=True) dR = F.interpolate(iR, (Hp // 2, Wp // 2), mode="bilinear", align_corners=True) with torch.inference_mode(): f0 = cre(dL, dR, iters=iters, flow_init=None) f = cre(iL, iR, iters=iters, flow_init=f0) return np.abs(f[0, 0].detach().cpu().numpy())[:h, :w] def belt_roi_px(cam, poly3): """the shared belt region, projected into this camera, as an axis-aligned window""" Minv = np.linalg.inv(np.array(cam["M"])) c = (np.c_[poly3, np.ones(len(poly3))] @ Minv)[:, :3]; z = -c[:, 2] u = c[:, 0] / np.maximum(z, 1e-9) * cam["fx"] + cam["cx"] v = -c[:, 1] / np.maximum(z, 1e-9) * cam["fy"] + cam["cy"] x0 = int(max(0, np.floor(u.min()))); x1 = int(min(cam["width"], np.ceil(u.max()))) y0 = int(max(0, np.floor(v.min()))); y1 = int(min(cam["height"], np.ceil(v.max()))) return x0, y0, x1, y1 def gate_px(cam): Minv = np.linalg.inv(np.array(cam["M"])) c = (np.r_[TARGET, 1.0] @ Minv)[:3]; z = -c[2] return (int(np.clip(c[0] / z * cam["fx"] + cam["cx"], 0, cam["width"] - 1)), int(np.clip(-c[1] / z * cam["fy"] + cam["cy"], 0, cam["height"] - 1))) def segment(img, win, gate): """FastSAM inside `win`; keep the smallest candidate containing the gate pixel. Returns a FULL-FRAME boolean mask.""" H, W = img.shape[:2] x0, y0, x1, y1 = win sub = img[y0:y1, x0:x1] hs, ws = sub.shape[:2]; A = hs * ws res = fsam(sub, imgsz=1024, conf=0.40, iou=0.90, retina_masks=True, max_det=60, verbose=False, device=dev) raw = (res[0].masks.data.cpu().numpy().astype(bool) if res[0].masks is not None else np.zeros((0, hs, ws), bool)) cands = [] for m in raw: if m.shape != (hs, ws): m = cv2.resize(m.astype(np.uint8), (ws, hs)).astype(bool) a = int(m.sum()) if a < 0.0015 * A or a > 0.35 * A: continue ys, xs = np.where(m) bw = xs.max() - xs.min() + 1; bh = ys.max() - ys.min() + 1 if a / (bw * bh) < 0.12: continue cands.append(m) if not cands: return None gx, gy = gate[0] - x0, gate[1] - y0 inside = [m for m in cands if 0 <= gy < hs and 0 <= gx < ws and m[gy, gx]] sel = (min(inside, key=lambda m: int(m.sum())) if inside else min(cands, key=lambda m: (np.where(m)[1].mean() - gx) ** 2 + (np.where(m)[0].mean() - gy) ** 2)) out = np.zeros((H, W), bool); out[y0:y1, x0:x1] = sel return out def backproj(disp, win, mask, cam, xoff_r=0): """disp is defined over `win` of the LEFT image; full-frame pixel coords keep the intrinsics valid, so no cx/cy shift is needed.""" x0, y0, x1, y1 = win d = disp + xoff_r depth = np.where(d > 0.5, cam["fx"] * cam["baseline"] / np.maximum(d, 1e-6), np.nan) mc = cv2.erode(mask[y0:y1, x0:x1].astype(np.uint8), np.ones((3, 3), np.uint8), iterations=2).astype(bool) vs, us = np.mgrid[y0:y1, x0:x1] sel = mc & np.isfinite(depth) & (depth > 1e-3) if sel.sum() < 30: return np.zeros((0, 3)) u, v, z = us[sel], vs[sel], depth[sel] P = np.stack([(u - cam["cx"]) * z / cam["fx"], -(v - cam["cy"]) * z / cam["fy"], -z, np.ones_like(z)], 1) return (P @ np.array(cam["M"]))[:, :3] crop3d = lambda P: P[(np.abs(P[:, 0] - TARGET[0]) < 0.30) & (np.abs(P[:, 1] - TARGET[1]) < 0.30) & (P[:, 2] > BELT_Z + 0.006) & (P[:, 2] < BELT_Z + 0.60)] def voxel(P, v=VOX): k = np.floor(P / v).astype(np.int64) _, i = np.unique(k, axis=0, return_index=True) return P[i] def dims_of(P): if len(P) < 60: return None P = voxel(P) c = P.mean(0); r = np.linalg.norm(P - c, axis=1) P = P[r < np.percentile(r, 94)] if len(P) < 40: return None h = (np.percentile(P[:, 2], 98) - BELT_Z) * 1000.0 rect = cv2.minAreaRect(np.ascontiguousarray(P[:, :2].astype(np.float32))) w, l = sorted([rect[1][0] * 1000.0, rect[1][1] * 1000.0]) return sorted([h, w, l], reverse=True) # ---- the belt region all six cameras see ---- g = 0.005 xs = np.arange(TARGET[0] - 1.6, TARGET[0] + 1.6, g) ys = np.arange(TARGET[1] - 1.6, TARGET[1] + 1.6, g) X, Y = np.meshgrid(xs, ys) G = np.c_[X.ravel(), Y.ravel(), np.full(X.size, BELT_Z)] vis = np.ones(len(G), bool) for cam in calib.values(): Minv = np.linalg.inv(np.array(cam["M"])) c = (np.c_[G, np.ones(len(G))] @ Minv)[:, :3]; z = -c[:, 2] u = c[:, 0] / np.maximum(z, 1e-9) * cam["fx"] + cam["cx"] v = -c[:, 1] / np.maximum(z, 1e-9) * cam["fy"] + cam["cy"] vis &= (z > 1e-3) & (u >= 0) & (u < cam["width"]) & (v >= 0) & (v < cam["height"]) cn, _ = cv2.findContours(vis.reshape(X.shape).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) pol = max(cn, key=cv2.contourArea).reshape(-1, 2) POLY3 = np.c_[xs[pol[:, 0]], ys[pol[:, 1]], np.full(len(pol), BELT_Z)] print(f"общая зона ленты: {vis.sum()*g*g*1e4:.0f} см2") for rig in RIGS: cam = calib[f"{rig}_Left"] x0, y0, x1, y1 = belt_roi_px(cam, POLY3) print(f" {rig:18s} ROI ленты {x1-x0}x{y1-y0} px " f"= {100.0*(x1-x0)*(y1-y0)/(cam['width']*cam['height']):.0f}% кадра") for _ in range(2): cre_infer(np.zeros((160, 224, 3), np.float32), np.zeros((160, 224, 3), np.float32)) results = {v: {} for v in VARIANTS} timing = {v: [] for v in VARIANTS} for name, e in man["items"].items(): gt = sorted(e["gt_scene_mm"], reverse=True) for var in VARIANTS: clouds, per_rig, by_rig = [], {}, {} t0 = time.time() for rig in RIGS: camL = calib[f"{rig}_Left"] IL = cv2.imread(e["files"][f"{rig}_Left"]).astype(np.float32) IR = cv2.imread(e["files"][f"{rig}_Right"]).astype(np.float32) if IL is None or IR is None: continue H, W = IL.shape[:2] full = (0, 0, W, H) broi = belt_roi_px(camL, POLY3) gate = gate_px(camL) segwin = full if var == "full" else broi mask = segment(IL.astype(np.uint8), segwin, gate) if mask is None: continue maxd = int(np.ceil(DPAD * camL["fx"] * camL["baseline"] / ZMIN)) if var in ("full", "beltroi"): x0, y0, x1, y1 = (full if var == "full" else broi) x0 = max(0, x0 - maxd) cL = IL[y0:y1, x0:x1]; cR = IR[y0:y1, x0:x1] disp = cre_infer(cL, cR) P = backproj(disp, (x0, y0, x1, y1), mask, camL) else: ys_, xs_ = np.where(mask) bx0, bx1 = xs_.min(), xs_.max(); by0, by1 = ys_.min(), ys_.max() y0 = max(0, by0 - PAD); y1 = min(H, by1 + PAD) x1 = min(W, bx1 + PAD) x0 = max(0, bx0 - PAD - maxd) # room for the right-image match cL = IL[y0:y1, x0:x1] if var == "objroi": cR = IR[y0:y1, x0:x1] # ONE window, both eyes shift = 0 else: # independently re-centred right crop d0 = int(round(camL["fx"] * camL["baseline"] / max(np.linalg.norm(np.array(camL["pos"]) - TARGET), 1e-6))) rx0 = max(0, x0 - d0); rx1 = min(W, x1 - d0) cR = IR[y0:y1, rx0:rx1] cR = cv2.resize(cR, (cL.shape[1], cL.shape[0])) if cR.shape != cL.shape else cR shift = -d0 disp = cre_infer(cL, cR) P = backproj(disp, (x0, y0, x1, y1), mask, camL, xoff_r=shift) P = crop3d(P) if len(P): clouds.append(P); per_rig[rig] = dims_of(P); by_rig[rig] = P dt = (time.time() - t0) * 1000 timing[var].append(dt) merged = dims_of(np.concatenate(clouds, 0)) if clouds else None wide = [P for r, P in by_rig.items() if r != "Orbbec_Gemini305"] merged2 = dims_of(np.concatenate(wide, 0)) if wide else None results[var][name] = dict(gt=gt, merged=merged, merged_no305=merged2, per_rig=per_rig, ms=round(dt)) mm = (f"{merged[0]:6.0f}{merged[1]:7.0f}{merged[2]:7.0f}" if merged else " нет облака ") mae = (np.mean(np.abs(np.array(merged) - np.array(gt))) if merged else float("nan")) print(f"{name:18s} {var:11s} эталон {gt[0]:5.0f}{gt[1]:6.0f}{gt[2]:6.0f} -> {mm} MAE {mae:6.1f} мм {dt:5.0f} мс") print("\n===== итог =====") print(f"{'вариант':12s} {'MAE, мм':>9s} {'облаков':>9s} {'мс/объект':>11s}") summ = {} for var in VARIANTS: maes = [np.mean(np.abs(np.array(r["merged"]) - np.array(r["gt"]))) for r in results[var].values() if r["merged"]] summ[var] = dict(mae=round(float(np.mean(maes)), 1) if maes else None, n=len(maes), ms=round(float(np.mean(timing[var])))) print(f"{var:12s} {summ[var]['mae'] if maes else '-':>9} " f"{len(maes)}/{len(results[var]):>7} {summ[var]['ms']:>11}") json.dump(dict(summary=summ, results=results), open(f"{CAP}/roi_compare.json", "w"), indent=1) print(f"\n-> {CAP}/roi_compare.json")