Сортировочная ячейка Isaac Sim: CV-пайплайн и меши товаров
Замкнутый контур "поток -> CV -> механика": товары идут по конвейеру с шагом 700 мм, класс определяется стереопайплайном во время движения, пушер и плуг реагируют физически. Состав: * control_test/ - ячейка и CV. run_sorting_cv.py + cv_worker.py (два процесса, потому что torch внутри Isaac роняет сцену), cell.py (физика лент, плуга, пушера), measure_plane.py (замер габаритов), README.md и .memory.md с замерами, проблемами и ловушками * robozon_sorter/ - модули симуляции, scripts/ - утилиты, scene/ - сцены * assets/ - меши товаров, плуг, объекты Objaverse Бейзлайн CV: DEFOM-Stereo vitl, вход 480, iters 24, кроп зоны осмотра, без сегментации. На потоке 700 мм - классы 8/9, габариты MAE 32.8 мм, 469 мс на товар при такте 700 мс. Веса моделей (4.5 ГБ) и пропсы конвейера NVIDIA (274 МБ) не включены - источники и команды скачивания в MODELS.md. Выход прогонов (captures/, runtime/) не включён: воспроизводится. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
dasha_f
@@ -0,0 +1,255 @@
|
||||
"""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")
|
||||
Reference in New Issue
Block a user