Сортировочная ячейка 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 @@
|
||||
from .pipeline import CreRoiV2b, classify, roundness # noqa: F401
|
||||
@@ -0,0 +1,355 @@
|
||||
"""CRE-ROI v2b: FastSAM segment-everything -> cross-view common object -> fixed-320 ROI
|
||||
crop -> ONE batched CRE-Stereo pass over the 3 camera pairs -> 3-view fusion -> B/C/D.
|
||||
|
||||
Why each piece is there
|
||||
ROI crop the item covers a small part of a 1280-wide frame; cropping spends the
|
||||
network's resolution on the cargo instead of on belt.
|
||||
same L/R window disparity is invariant to an equal column shift, so the crop must use
|
||||
identical [x0,x1] in both eyes and be left-padded by the max disparity,
|
||||
or the right-hand counterpart falls outside the crop.
|
||||
batching the 3 crops go through CRE as one forward pass; that is the "v2b" part.
|
||||
gate pixel the inspection point is projected into every camera, and the blob covering
|
||||
it is kept - by construction all views then measure the SAME physical item.
|
||||
consistency a view whose 3D centroid disagrees with the median is dropped; that is the
|
||||
cross-view rule enforced again in 3D, where it is unambiguous.
|
||||
|
||||
The stereo rig must be RECTIFIED (parallel optical axes). Verged pairs break
|
||||
depth = fx*b/disp and the reconstruction lands metres away.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .. import config as C
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ geometry helpers
|
||||
def _taubin(xy):
|
||||
"""algebraic circle fit; returns (cx, cy, R, mean relative residual)"""
|
||||
x = xy[:, 0].astype(np.float64)
|
||||
y = xy[:, 1].astype(np.float64)
|
||||
xm, ym = x.mean(), y.mean()
|
||||
u, v = x - xm, y - ym
|
||||
Suu, Svv, Suv = (u * u).sum(), (v * v).sum(), (u * v).sum()
|
||||
Suuu, Svvv = (u ** 3).sum(), (v ** 3).sum()
|
||||
Suvv, Svuu = (u * v * v).sum(), (v * u * u).sum()
|
||||
try:
|
||||
uc, vc = np.linalg.solve(np.array([[Suu, Suv], [Suv, Svv]]),
|
||||
0.5 * np.array([Suuu + Suvv, Svvv + Svuu]))
|
||||
except np.linalg.LinAlgError:
|
||||
return None
|
||||
cx, cy = uc + xm, vc + ym
|
||||
R = np.sqrt(max(uc * uc + vc * vc + (Suu + Svv) / len(x), 1e-12))
|
||||
if not np.isfinite(R) or R < 1e-6:
|
||||
return None
|
||||
return cx, cy, R, float(np.abs(np.hypot(x - cx, y - cy) - R).mean() / R)
|
||||
|
||||
|
||||
def _rin_rout(xy):
|
||||
"""K = inscribed / circumscribed radius of the convex hull; 1.0 for a perfect circle"""
|
||||
pts = np.ascontiguousarray(xy.astype(np.float32))
|
||||
if len(pts) < 3:
|
||||
return 0.0
|
||||
_, r_out = cv2.minEnclosingCircle(pts)
|
||||
if r_out < 1e-6:
|
||||
return 0.0
|
||||
xmin, ymin = xy.min(0)
|
||||
w, h = float(np.ptp(xy[:, 0])), float(np.ptp(xy[:, 1]))
|
||||
sc = 180.0 / max(w, h, 1e-6)
|
||||
img = np.zeros((int(h * sc) + 10, int(w * sc) + 10), np.uint8)
|
||||
try:
|
||||
hull = cv2.convexHull(pts).reshape(-1, 2)
|
||||
except cv2.error:
|
||||
return 0.0
|
||||
cv2.fillConvexPoly(img, ((hull - [xmin, ymin]) * sc + 5).astype(np.int32), 255)
|
||||
return float(cv2.distanceTransform(img, cv2.DIST_L2, 5).max()) / sc / r_out
|
||||
|
||||
|
||||
def _section_k(xy, res_tol=0.06, min_span=120.0):
|
||||
if len(xy) < 15:
|
||||
return 0.0
|
||||
fit = _taubin(xy)
|
||||
if fit is not None:
|
||||
cx, cy, R, res = fit
|
||||
ang = np.arctan2(xy[:, 1] - cy, xy[:, 0] - cx)
|
||||
span = np.unique((((ang + np.pi) / (2 * np.pi)) * 48).astype(int) % 48).size / 48.0 * 360.0
|
||||
ext = max(np.ptp(xy[:, 0]), np.ptp(xy[:, 1])) + 1e-9
|
||||
if res < res_tol and span >= min_span and 0.35 * ext < R < 2.0 * ext:
|
||||
return 1.0
|
||||
return _rin_rout(xy)
|
||||
|
||||
|
||||
def voxel(P, v=0.004):
|
||||
key = np.floor(P / v).astype(np.int64)
|
||||
_, idx = np.unique(key, axis=0, return_index=True)
|
||||
return P[idx]
|
||||
|
||||
|
||||
def roundness(P):
|
||||
"""max circular-section K over the top-down section and one belt-aligned cross section"""
|
||||
if len(P) < 60:
|
||||
return 0.0
|
||||
P = voxel(P)
|
||||
ks = [_section_k(P[:, :2])]
|
||||
xy = P[:, :2] - P[:, :2].mean(0)
|
||||
try:
|
||||
_, _, V = np.linalg.svd(xy, full_matrices=False)
|
||||
ax = V[0]
|
||||
except np.linalg.LinAlgError:
|
||||
ax = np.array([1.0, 0.0])
|
||||
along = xy @ ax
|
||||
perp = xy @ np.array([-ax[1], ax[0]])
|
||||
mid = np.abs(along - np.median(along)) < 0.15 * (np.ptp(along) + 1e-9)
|
||||
if mid.sum() > 15:
|
||||
ks.append(_section_k(np.c_[perp[mid], P[mid, 2]]))
|
||||
return max(ks)
|
||||
|
||||
|
||||
def classify(P, belt_z: float):
|
||||
"""belt-plane OBB -> dims in real millimetres -> B / C / D"""
|
||||
if len(P) < 60:
|
||||
return "?", [0, 0, 0], 0.0
|
||||
P = voxel(P)
|
||||
centre = P.mean(0)
|
||||
radial = np.linalg.norm(P - centre, axis=1)
|
||||
P = P[radial < np.percentile(radial, 94)] # trim segmentation fringe
|
||||
if len(P) < 40:
|
||||
return "?", [0, 0, 0], 0.0
|
||||
height = (np.percentile(P[:, 2], 98) - belt_z) / C.DIM_SCALE * 1000.0
|
||||
(_, _), (rw, rh), _ = cv2.minAreaRect(np.ascontiguousarray(P[:, :2].astype(np.float32)))
|
||||
foot = sorted([rw / C.DIM_SCALE * 1000.0, rh / C.DIM_SCALE * 1000.0])
|
||||
dims = sorted([height, foot[0], foot[1]], reverse=True)
|
||||
k = roundness(P)
|
||||
a, b, c = C.OVERSIZE_MAX
|
||||
if dims[0] > a or dims[1] > b or dims[2] > c or dims[2] < C.MIN_DIM:
|
||||
return "C", [round(x) for x in dims], k
|
||||
return ("D" if k > C.ROUND_K else "B"), [round(x) for x in dims], k
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ the pipeline
|
||||
class CreRoiV2b:
|
||||
"""Holds the models, the render products and the per-frame decision."""
|
||||
|
||||
def __init__(self, calib_path: Path | str = None, device="cuda"):
|
||||
self.device = device
|
||||
calib_path = Path(calib_path or C.CONFIG / "calib.json")
|
||||
self.calib = json.loads(calib_path.read_text())
|
||||
self.belt_z = self.calib.get("belt_top", C.BELT_Z)
|
||||
self.aim = np.array(self.calib["center"], dtype=float)
|
||||
if not self.calib.get("rectified"):
|
||||
raise ValueError(
|
||||
f"{calib_path} is not marked rectified. depth = fx*b/disp assumes parallel "
|
||||
"optical axes; rebuild the rig with sim.scene.build_camera_rig()."
|
||||
)
|
||||
self._load_models()
|
||||
self.cams = {}
|
||||
|
||||
# -- models -------------------------------------------------------------
|
||||
def _load_models(self):
|
||||
import sys
|
||||
cre_dir = C.MODELS / "crestereo"
|
||||
if str(cre_dir) not in sys.path:
|
||||
sys.path.insert(0, str(cre_dir))
|
||||
from nets import Model # noqa: E402 (vendored CRE-Stereo)
|
||||
|
||||
self.cre = Model(max_disp=256, mixed_precision=False, test_mode=True)
|
||||
weights = C.MODELS / "crestereo_eth3d.pth"
|
||||
self.cre.load_state_dict(torch.load(weights, map_location="cpu"), strict=True)
|
||||
self.cre.to(self.device).eval()
|
||||
|
||||
from ultralytics import FastSAM # noqa: E402
|
||||
self.fsam = FastSAM(str(C.MODELS / "FastSAM-s.pt"))
|
||||
|
||||
# -- render products ----------------------------------------------------
|
||||
def attach_cameras(self):
|
||||
"""one RGB annotator per eye; call once, after the stage is populated"""
|
||||
import omni.replicator.core as rep
|
||||
|
||||
for name, cc in self.calib["cameras"].items():
|
||||
K = cc["intrinsics"]
|
||||
res = (K["width"], K["height"])
|
||||
ann = {}
|
||||
for side, path in (("L", cc["left_path"]), ("R", cc["right_path"])):
|
||||
rp = rep.create.render_product(path, res)
|
||||
a = rep.AnnotatorRegistry.get_annotator("rgb")
|
||||
a.attach(rp)
|
||||
ann[side] = a
|
||||
self.cams[name] = dict(K=K, b=cc["baseline_m"],
|
||||
LW=np.array(cc["left_world"], dtype=float), **ann)
|
||||
self.gate_px = {n: self._gate_px(n) for n in self.cams}
|
||||
return self.cams
|
||||
|
||||
async def warmup(self, steps=120, tries=4):
|
||||
"""Step the app until the eyes actually return an image.
|
||||
|
||||
`attach_cameras` creates the render products and returns at once, but RTX yields
|
||||
nothing for a fresh render product for a long while, so the first `get_data()` comes
|
||||
back PURE BLACK - RGB max 0 with alpha ~255, meaning geometry is being hit and is
|
||||
simply unresolved, not that the view is empty. FastSAM then segments nothing and
|
||||
every early item is logged unclassified, which reads as "the pipeline is broken".
|
||||
|
||||
Measured on this cell: 60 steps still black; 120 steps plus a second settle gives
|
||||
RGB max 92 / mean 21 on the D435 left eye. So this waits and CHECKS rather than
|
||||
trusting a fixed count - it returns the darkest eye it saw, for the run log.
|
||||
"""
|
||||
import isaacsim.core.experimental.utils.app as app_utils
|
||||
worst = 0
|
||||
for attempt in range(tries):
|
||||
await app_utils.update_app_async(steps=steps)
|
||||
worst = 255
|
||||
for c in self.cams.values():
|
||||
for side in ("L", "R"):
|
||||
a = np.asarray(c[side].get_data())
|
||||
worst = min(worst, int(a[..., :3].max()) if a.size else 0)
|
||||
if worst > 3:
|
||||
return dict(ok=True, attempts=attempt + 1, darkest_eye_max=worst)
|
||||
return dict(ok=False, attempts=tries, darkest_eye_max=worst)
|
||||
|
||||
def _gate_px(self, name):
|
||||
c = self.cams[name]
|
||||
K = c["K"]
|
||||
p = np.append(self.aim, 1.0) @ np.linalg.inv(c["LW"])
|
||||
d = max(-p[2], 1e-6)
|
||||
u = K["cx"] + K["fx"] * p[0] / d
|
||||
v = K["cy"] - K["fy"] * p[1] / d
|
||||
return int(np.clip(u, 0, K["width"] - 1)), int(np.clip(v, 0, K["height"] - 1))
|
||||
|
||||
# -- stages -------------------------------------------------------------
|
||||
def segment(self, rgb, gate):
|
||||
"""smallest plausible blob covering the gate pixel"""
|
||||
t0 = time.time()
|
||||
res = self.fsam(rgb[..., ::-1], device=self.device, retina_masks=True,
|
||||
imgsz=1024, conf=0.35, iou=0.9, verbose=False)
|
||||
ms = (time.time() - t0) * 1000
|
||||
if not res or res[0].masks is None:
|
||||
return None, ms
|
||||
gx, gy = gate
|
||||
best, best_area = None, np.inf
|
||||
for m in res[0].masks.data.cpu().numpy():
|
||||
mb = m > 0.5
|
||||
area = int(mb.sum())
|
||||
if area < 150 or area > C.MAX_MASK_FRAC * mb.size:
|
||||
continue
|
||||
if mb[gy, gx] and area < best_area:
|
||||
best, best_area = mb, area
|
||||
return best, ms
|
||||
|
||||
@staticmethod
|
||||
def crop(left, right, mask, K, baseline):
|
||||
"""identical column window in both eyes, left-padded by the max disparity"""
|
||||
H, W = mask.shape
|
||||
ys, xs = np.where(mask)
|
||||
max_disp = int(np.ceil(C.DISP_PAD * K["fx"] * baseline / C.Z_MIN))
|
||||
x0 = max(0, xs.min() - max_disp - C.ROI_PAD_X)
|
||||
x1 = min(W, xs.max() + C.ROI_PAD_X)
|
||||
y0 = max(0, ys.min() - C.ROI_PAD_V)
|
||||
y1 = min(H, ys.max() + C.ROI_PAD_V)
|
||||
cl, cr = left[y0:y1, x0:x1], right[y0:y1, x0:x1]
|
||||
h, w = cl.shape[:2]
|
||||
s = min(C.ROI_FIXED / max(h, w, 1), 2.5)
|
||||
if abs(s - 1) > 0.02:
|
||||
cl = cv2.resize(cl, (max(int(round(w * s)), 8), max(int(round(h * s)), 8)))
|
||||
cr = cv2.resize(cr, (cl.shape[1], cl.shape[0]))
|
||||
return np.ascontiguousarray(cl), np.ascontiguousarray(cr), (x0, y0, x1, y1), s, (h, w)
|
||||
|
||||
def infer_batch(self, crops, iters=20):
|
||||
"""all camera crops in ONE two-stage CRE pass, zero-padded to a common /8 canvas"""
|
||||
sizes = [c[0].shape[:2] for c in crops]
|
||||
Hp = (max(h for h, _ in sizes) + 7) // 8 * 8
|
||||
Wp = (max(w for _, w in sizes) + 7) // 8 * 8
|
||||
lb = np.zeros((len(crops), 3, Hp, Wp), np.float32)
|
||||
rb = np.zeros_like(lb)
|
||||
for i, (l, r) in enumerate(crops):
|
||||
h, w = l.shape[:2]
|
||||
lb[i, :, :h, :w] = l.transpose(2, 0, 1)
|
||||
rb[i, :, :h, :w] = r.transpose(2, 0, 1)
|
||||
il = torch.from_numpy(lb).to(self.device)
|
||||
ir = torch.from_numpy(rb).to(self.device)
|
||||
ild = F.interpolate(il, (Hp // 2, Wp // 2), mode="bilinear", align_corners=True)
|
||||
ird = F.interpolate(ir, (Hp // 2, Wp // 2), mode="bilinear", align_corners=True)
|
||||
with torch.inference_mode():
|
||||
init = self.cre(ild, ird, iters=iters, flow_init=None)
|
||||
flow = self.cre(il, ir, iters=iters, flow_init=init)
|
||||
disp = np.abs(flow[:, 0].detach().cpu().numpy())
|
||||
return [disp[i, :h, :w] for i, (h, w) in enumerate(sizes)]
|
||||
|
||||
def backproject(self, disp_small, mask, box, s, orig, K, baseline, LW):
|
||||
h, w = orig
|
||||
disp = cv2.resize(disp_small, (w, h)) / s
|
||||
depth = np.where(disp > 0.5, K["fx"] * baseline / np.maximum(disp, 1e-6), np.nan)
|
||||
x0, y0, x1, y1 = box
|
||||
# erode: silhouette pixels straddle the depth discontinuity and smear the cloud
|
||||
m = 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 = m & np.isfinite(depth) & (depth > 1e-3)
|
||||
u, v, d = us[sel], vs[sel], depth[sel]
|
||||
cam = np.stack([(u - K["cx"]) * d / K["fx"],
|
||||
-(v - K["cy"]) * d / K["fy"],
|
||||
-d, np.ones_like(d)], 1)
|
||||
return (cam @ LW)[:, :3]
|
||||
|
||||
def _in_workspace(self, P):
|
||||
return P[(np.abs(P[:, 0] - self.aim[0]) < 0.45) & (np.abs(P[:, 1]) < 0.35) &
|
||||
(P[:, 2] > self.belt_z + 0.012) & (P[:, 2] < self.belt_z + 0.80)]
|
||||
|
||||
# -- one decision -------------------------------------------------------
|
||||
def measure(self):
|
||||
"""capture -> segment -> crop -> batched CRE -> fuse -> classify"""
|
||||
t0 = time.time()
|
||||
crops, meta, seg_ms = [], [], []
|
||||
for name, c in self.cams.items():
|
||||
left = c["L"].get_data()[..., :3]
|
||||
right = c["R"].get_data()[..., :3]
|
||||
mask, ms = self.segment(left, self.gate_px[name])
|
||||
seg_ms.append(round(ms))
|
||||
if mask is None:
|
||||
continue
|
||||
cl, cr, box, s, orig = self.crop(left.astype(np.float32) / 255.0,
|
||||
right.astype(np.float32) / 255.0,
|
||||
mask, c["K"], c["b"])
|
||||
crops.append((cl, cr))
|
||||
meta.append((name, mask, box, s, orig))
|
||||
|
||||
if not crops:
|
||||
return dict(cls="?", dims=[0, 0, 0], k=0.0, views=0, pts=0,
|
||||
seg_ms=seg_ms, cre_ms=0, total_ms=round((time.time() - t0) * 1000))
|
||||
|
||||
t_cre = time.time()
|
||||
disps = self.infer_batch(crops)
|
||||
cre_ms = round((time.time() - t_cre) * 1000)
|
||||
|
||||
clouds = []
|
||||
for (name, mask, box, s, orig), disp in zip(meta, disps):
|
||||
c = self.cams[name]
|
||||
P = self._in_workspace(
|
||||
self.backproject(disp, mask, box, s, orig, c["K"], c["b"], c["LW"]))
|
||||
if len(P) >= 40:
|
||||
clouds.append((name, P))
|
||||
|
||||
dropped = []
|
||||
if len(clouds) > 1:
|
||||
cent = np.array([p.mean(0) for _, p in clouds])
|
||||
med = np.median(cent, 0)
|
||||
keep = [i for i in range(len(clouds))
|
||||
if np.linalg.norm(cent[i] - med) < C.VIEW_CONSISTENCY]
|
||||
dropped = [clouds[i][0] for i in range(len(clouds)) if i not in keep]
|
||||
if keep:
|
||||
clouds = [clouds[i] for i in keep]
|
||||
|
||||
if not clouds:
|
||||
return dict(cls="?", dims=[0, 0, 0], k=0.0, views=0, pts=0, seg_ms=seg_ms,
|
||||
cre_ms=cre_ms, total_ms=round((time.time() - t0) * 1000))
|
||||
|
||||
P = np.vstack([p for _, p in clouds])
|
||||
cls, dims, k = classify(P, self.belt_z)
|
||||
return dict(cls=cls, dims=dims, k=round(float(k), 3), views=len(clouds),
|
||||
dropped=dropped, pts=len(P), seg_ms=seg_ms, cre_ms=cre_ms,
|
||||
total_ms=round((time.time() - t0) * 1000))
|
||||
Reference in New Issue
Block a user