Files
dasha_f 0d32f32db0 Сортировочная ячейка 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>
2026-08-01 13:07:24 +00:00

432 lines
23 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""ЭТАП 2 потока: батчевый инференс по кадрам real-time прогона, метрики габаритов и классов.
Стенд отдельным процессом, потому что torch внутри Isaac роняет процесс - то же разделение,
что у measure_roi.py.
ЧТО ПРЕДСКАЗЫВАЕТСЯ, А ЧТО ТОЛЬКО СВЕРЯЕТСЯ. Из захвата берутся ТОЛЬКО кадры. Габариты,
округлость k и класс вычисляются из облака точек; разметка из манифеста используется
исключительно для подсчёта ошибки. Окно поиска и "воротный" пиксель выводятся из
калибровки камер и точки осмотра, а не из положения товара.
БАТЧИРОВАНИЕ. Три рига обрабатываются одним вызовом на каждой ступени:
FastSAM получает список из трёх левых кадров, CREStereo - тензор (3, 3, H, W) с кропами,
дополненными до общего размера. Так GPU загружается один раз вместо трёх, и это то, что
позволяет уложиться в такт потока: при шаге 700 мм и 1 м/с на товар отводится 0.70 с.
СХЕМА ROI - objroi, единственная проверенная: общая зона ленты ограничивает область
сегментации, CRE считается по bbox маски с запасом, причём левое и правое окно колонок
ОДНО И ТО ЖЕ, с расширением влево на максимальную диспаратность. Контрольный вариант с
независимым центрированием правого кропа разрушал 8 облаков из 9.
RGB не маскируется до CRE - матчеру нужен фон вокруг предмета; маска применяется к
глубине, с эрозией.
ОКРУГЛОСТЬ k считается из облака, а не читается: берётся сечение на середине высоты,
k = вписанный радиус / описанный. У круга 1.0, у квадрата 0.707, порог класса D - 0.8.
"""
import os, sys, json, time
import numpy as np
import cv2
import torch
import torch.nn.functional as F
CV = "/home/dasha/isaac_assets/cv"
CT = "/home/dasha/robozon-sorter/control_test"
CAP = f"{CT}/captures/flow"
MASK_MODE = os.environ.get("MASK_MODE", "gate") # gate | plane
K_MODE = os.environ.get("K_MODE", "hull") # hull | fit
TAG = os.environ.get("TAG", f"{MASK_MODE}+{K_MODE}")
PITCH_S = 0.70 # такт потока: 700 мм при 1 м/с
dev = "cuda"
sys.path.insert(0, CT)
import classify as CL
man = json.load(open(f"{CAP}/manifest.json"))
calib = man["calib"]
TARGET = np.array(man["target"]); BELT_Z = float(TARGET[2])
RIGS = sorted({c["rig"] for c in calib.values()})
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()
SEG_MODEL = os.environ.get("SEG_MODEL", "fastsam") # fastsam | yolo
if SEG_MODEL == "yolo":
# yolo26n-seg: экземплярная сегментация с классами COCO. На конвейере товары в COCO
# почти не представлены, поэтому классы игнорируются - берутся только маски, а выбор
# кандидата остаётся тем же, что у FastSAM. Локально есть только вариант n (nano);
# yolo26-s-seg по ссылке пришлось бы скачивать.
from ultralytics import YOLO
fsam = YOLO(f"{CV}/yolo26n-seg.pt")
else:
from ultralytics import FastSAM
fsam = FastSAM(f"{CV}/FastSAM-s.pt")
ZMIN, DPAD, PAD, VOX = 0.60, 1.35, 48, 0.004
# Порог отбраковки ракурса. 0.06 м оказался ВРЕДЕН: MAE вырос 49.5 -> 65.4, классы
# 5/9 -> 3/9, потому что правило откидывало два вида из трёх и оставляло одиночный,
# который систематически занижает габарит из-за самозатенения. Центроиды трёх ригов
# расходятся ЗАКОННО: каждый видит свою обращённую к нему поверхность, и у выпуклого тела
# центр видимой части смещён к наблюдателю примерно на 2r/pi - для предмета 300 мм это
# уже 60-95 мм. Поэтому порог грубый: ловим только явные выбросы, когда сегментация взяла
# ленту (у bag было 141 и 230 мм).
REJECT_M = 0.15
def cre_batch(pairs, iters=20):
"""CREStereo одним проходом по нескольким парам кропов.
Кропы у ригов разного размера, поэтому все дополняются нулями до общего (кратного 8)
и складываются в один тензор. Дополнение не влияет на результат: диспаратность
читается только внутри исходных границ каждого кропа.
"""
if not pairs:
return []
hs = [p[0].shape[0] for p in pairs]; ws = [p[0].shape[1] for p in pairs]
Hp = (max(hs) + 7) // 8 * 8; Wp = (max(ws) + 7) // 8 * 8
n = len(pairs)
Lb = np.zeros((n, 3, Hp, Wp), np.float32); Rb = np.zeros((n, 3, Hp, Wp), np.float32)
for i, (L, R) in enumerate(pairs):
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(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)
out = np.abs(f[:, 0].detach().cpu().numpy())
return [out[i][:hs[i], :ws[i]] for i in range(n)]
def segment_batch(imgs, wins, gates):
"""FastSAM одним вызовом по нескольким кадрам; для каждого - самая мелкая маска,
содержащая воротный пиксель."""
subs, offs = [], []
for img, (x0, y0, x1, y1) in zip(imgs, wins):
subs.append(img[y0:y1, x0:x1]); offs.append((x0, y0))
if SEG_MODEL == "yolo":
res = fsam(subs, imgsz=1024, conf=0.10, iou=0.70, retina_masks=True,
device=dev, verbose=False) # низкий conf: классы COCO тут не подходят
else:
res = fsam(subs, imgsz=1024, conf=0.40, iou=0.90, retina_masks=True,
device=dev, verbose=False)
out = []
for r, img, (x0, y0), (gx, gy) in zip(res, imgs, offs, gates):
H, W = img.shape[:2]
if r.masks is None:
out.append(None); continue
raw = r.masks.data.cpu().numpy().astype(bool)
hs, ws = raw.shape[1], raw.shape[2]
gx_, gy_ = gx - x0, gy - y0
best = None
if MASK_MODE == "gate":
for m in raw:
if not (0 <= gy_ < hs and 0 <= gx_ < ws) or not m[gy_, gx_]:
continue
n = int(m.sum())
if n < 200 or n > 0.5 * hs * ws:
continue
if best is None or n < best[0]:
best = (n, m)
else:
cam = CAM_BY_IMG[id(img)]
for m in raw:
n = int(m.sum())
if n < 200 or n > BELT_FRAC * hs * ws:
continue # такой площади бывает только лента
yy, xx = np.nonzero(m)
u = float(xx.mean()) + x0; v = float(yy.mean()) + y0
q = ray_to_belt(cam, u, v)
if q is None:
continue
d = float(np.hypot(q[0] - TARGET[0], q[1] - TARGET[1]))
if d > 0.45:
continue # не у точки осмотра
if best is None or d < best[0]:
best = (d, m)
if best is None:
out.append(None); continue
full = np.zeros((H, W), bool)
full[y0:y0 + hs, x0:x0 + ws] = best[1]
out.append(full)
return out
BELT_FRAC = 0.28 # доля окна: больше - это полотно, а не товар
CAM_BY_IMG = {}
def ray_to_belt(cam, u, v):
"""пиксель -> точка на ПЛОСКОСТИ ЛЕНТЫ по калибровке. Глубина не нужна: плоскость
известна (z = BELT_Z), поэтому луч пересекается с ней аналитически."""
M = np.array(cam["M"])
d_cam = np.array([(u - cam["cx"]) / cam["fx"], -(v - cam["cy"]) / cam["fy"], -1.0, 0.0])
o = (np.array([0.0, 0.0, 0.0, 1.0]) @ M)[:3]
d = (d_cam @ M)[:3]
if abs(d[2]) < 1e-9:
return None
t = (BELT_Z - o[2]) / d[2]
if t <= 0:
return None
return o + t * d
def belt_roi_px(cam, poly3):
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"]
return (int(max(0, np.floor(u.min()))), int(max(0, np.floor(v.min()))),
int(min(cam["width"], np.ceil(u.max()))), int(min(cam["height"], np.ceil(v.max()))))
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 backproj(disp, win, mask, cam):
"""disp задана над окном `win` ЛЕВОГО кадра. Пиксельные координаты берутся ПОЛНОГО
кадра, поэтому интринсики остаются в силе и сдвигать cx/cy не нужно.
Матрица камеры умножается СПРАВА и БЕЗ транспонирования: `P @ M`. Транспонирование
даёт правдоподобные по виду, но неверные координаты - облако уезжает за пределы
отсечения, и на выходе получается ноль точек при исправной сегментации и диспаратности.
"""
x0, y0, x1, y1 = win
depth = np.where(disp > 0.5, cam["fx"] * cam["baseline"] / np.maximum(disp, 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)
mc = mc[:depth.shape[0], :depth.shape[1]]
vs, us = np.mgrid[y0:y0 + depth.shape[0], x0:x0 + depth.shape[1]]
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)
P = (P @ np.array(cam["M"]))[:, :3]
return 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)
_, idx = np.unique(k, axis=0, return_index=True)
return P[idx]
def dims_and_k(P):
"""габариты и округлость из облака. k = вписанный/описанный радиус сечения."""
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])
dims = sorted([h, w, l], reverse=True)
# округлость по сечению на середине высоты
zmid = BELT_Z + (np.percentile(P[:, 2], 98) - BELT_Z) * 0.5
band = P[np.abs(P[:, 2] - zmid) < 0.02][:, :2]
k = float("nan")
if K_MODE == "fit" and len(band) >= 25:
pts = band * 1000.0
x_, y_ = pts[:, 0], pts[:, 1]
A = np.stack([x_, y_, np.ones_like(x_)], 1)
try:
sol, *_ = np.linalg.lstsq(A, x_ ** 2 + y_ ** 2, rcond=None)
cx, cy = sol[0] / 2, sol[1] / 2
for _ in range(40): # геометрическое уточнение центра
dx, dy = x_ - cx, y_ - cy
rr = np.maximum(np.hypot(dx, dy), 1e-9)
R = rr.mean()
J = np.stack([-dx / rr, -dy / rr], 1)
st, *_ = np.linalg.lstsq(J, -(rr - R), rcond=None)
cx += st[0]; cy += st[1]
if abs(st[0]) + abs(st[1]) < 1e-6:
break
r = np.hypot(x_ - cx, y_ - cy)
lo, hi = np.percentile(r, 10), np.percentile(r, 90)
k = float(min(1.0, lo / hi)) if hi > 1e-6 else float("nan")
except np.linalg.LinAlgError:
pass
elif len(band) >= 25:
pts = (band - band.mean(0)) * 1000.0
hull = cv2.convexHull(np.ascontiguousarray(pts.astype(np.float32)))
(_, _), R_out = cv2.minEnclosingCircle(hull)
if R_out > 1e-6:
s = 2.0 # 0.5 мм на пиксель
sh = int(np.ceil((pts.max() - pts.min()) * s)) + 20
img = np.zeros((sh, sh), np.uint8)
poly = np.int32((hull.reshape(-1, 2) - pts.min()) * s + 10)
cv2.fillConvexPoly(img, poly, 255)
dt = cv2.distanceTransform(img, cv2.DIST_L2, 5)
r_in = float(dt.max()) / s
k = min(1.0, r_in / R_out)
return dims, k
# ---- общая зона ленты, видимая всеми шестью камерами ----
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)]
for _ in range(2): # прогрев GPU
cre_batch([(np.zeros((160, 224, 3), np.float32), np.zeros((160, 224, 3), np.float32))])
print(f"[{TAG}] поток: {len(man['items'])} товаров, такт {PITCH_S:.2f} с "
f"(шаг 700 мм при 1 м/с), риги {', '.join(RIGS)}\n")
print(f" {'товар':18s} {'эталон, мм':>18s} {'предсказано, мм':>20s} {'MAE':>6s} "
f"{'k пред':>7s} {'класс':>11s} {'мс':>6s} виды")
print(" " + "-" * 96)
rows, times = [], []
for name, e in man["items"].items():
t0 = time.time()
imgs, wins, gates, cams, rights = [], [], [], [], []
for rig in RIGS:
camL = calib[f"{rig}_Left"]
IL = cv2.imread(e["files"][f"{rig}_Left"])
IR = cv2.imread(e["files"][f"{rig}_Right"])
if IL is None or IR is None:
continue
imgs.append(IL); rights.append(IR); cams.append(camL)
wins.append(belt_roi_px(camL, POLY3)); gates.append(gate_px(camL))
CAM_BY_IMG[id(IL)] = camL
if not imgs:
continue
t_seg0 = time.time()
masks = segment_batch(imgs, wins, gates)
t_seg = (time.time() - t_seg0) * 1000
pairs, metas = [], []
for IL, IR, cam, mask in zip(imgs, rights, cams, masks):
if mask is None:
continue
H, W = IL.shape[:2]
ys_, xs_ = np.where(mask)
maxd = int(np.ceil(DPAD * cam["fx"] * cam["baseline"] / ZMIN))
y0 = max(0, ys_.min() - PAD); y1 = min(H, ys_.max() + PAD)
x1 = min(W, xs_.max() + PAD); x0 = max(0, xs_.min() - PAD - maxd)
pairs.append((IL[y0:y1, x0:x1].astype(np.float32),
IR[y0:y1, x0:x1].astype(np.float32)))
metas.append(((x0, y0, x1, y1), mask, cam))
t_cre0 = time.time()
disps = cre_batch(pairs)
t_cre = (time.time() - t_cre0) * 1000
clouds, cnames = [], []
for disp, (win, mask, cam) in zip(disps, metas):
P = backproj(disp, win, mask, cam)
if P is not None and len(P):
clouds.append(P); cnames.append(cam["rig"])
# ОТБРАКОВКА РАКУРСА ПО 3D-ЦЕНТРОИДУ.
# Слияние идёт по калиброванным экстринсикам, поэтому три облака одного предмета
# обязаны лежать в одном месте. Если сегментация одного рига прихватила ленту, его
# центроид уезжает, и после слияния габарит раздувается: у bag вышло 530 мм при
# истинных 202, у bucket 603 при 287 - и раздутое сечение перестаёт быть круглым,
# из-за чего оба класса D ушли в C.
# Берётся ПОКООРДИНАТНАЯ МЕДИАНА центроидов (устойчива при трёх видах: один выброс
# её не двигает) и выбрасываются виды дальше REJECT_M от неё.
dropped = []
if len(clouds) >= 2:
cent = np.array([c.mean(0) for c in clouds])
med = np.median(cent, axis=0)
off = np.linalg.norm(cent - med, axis=1)
keep = off <= REJECT_M
# никогда не опускаться ниже двух видов: слияние двух ракурсов лучше одиночного
if keep.sum() < 2 and len(clouds) >= 2:
keep = np.zeros(len(clouds), bool)
keep[np.argsort(off)[:2]] = True
if 2 <= keep.sum() < len(clouds):
dropped = [(cnames[i], round(float(off[i]) * 1000)) for i in range(len(clouds))
if not keep[i]]
clouds = [c for c, k in zip(clouds, keep) if k]
cnames = [n for n, k in zip(cnames, keep) if k]
dt = (time.time() - t0) * 1000
times.append(dt)
gt = e["gt"]
gt_dims = sorted(gt["dims_mm"], reverse=True)
if not clouds:
nm = sum(1 for m in masks if m is not None)
print(f" {name:18s} {str([round(v) for v in gt_dims]):>18s} "
f" масок {nm}/{len(imgs)}, пар в CRE {len(pairs)}, облаков 0")
rows.append(dict(name=name, gt=gt_dims, gt_cls=gt["zone_scene"],
gt_label=gt["zone_label"], pred=None))
continue
out = dims_and_k(np.vstack(clouds))
if out is None:
print(f" {name:18s} {'габариты не взялись':>40s}"); continue
dims, k = out
mae = float(np.mean(np.abs(np.array(dims) - np.array(gt_dims))))
pred_cls = CL.classify(dims, 0.0 if np.isnan(k) else k)
ok = "верно" if pred_cls == gt["zone_scene"] else f"ОШИБКА({gt['zone_scene']})"
drp = ("" if not dropped else
" откинут " + ", ".join(f"{n} ({d} мм)" for n, d in dropped))
print(f" {name:18s} {str([round(v) for v in gt_dims]):>18s} "
f"{str([round(v) for v in dims]):>20s} {mae:6.1f} {k:7.2f} "
f"{pred_cls + ' ' + ok:>11s} {dt:6.0f} {len(clouds)}в{drp}")
rows.append(dict(name=name, gt=gt_dims, gt_cls=gt["zone_scene"], gt_label=gt["zone_label"], gt_k=gt.get("k"),
pred=[round(v, 1) for v in dims], pred_k=round(float(k), 3),
pred_cls=pred_cls, mae=round(mae, 1), ms=round(dt),
views=len(clouds), dropped=dropped,
t_seg=round(t_seg), t_cre=round(t_cre)))
got = [r for r in rows if r.get("pred")]
print("\n === МЕТРИКИ ===")
if got:
maes = [r["mae"] for r in got]
print(f" габариты: MAE медиана {np.median(maes):.1f} мм, среднее {np.mean(maes):.1f}, "
f"худший {max(maes):.1f} ({max(got, key=lambda r: r['mae'])['name']})")
acc = sum(1 for r in got if r["pred_cls"] == r["gt_cls"])
accl = sum(1 for r in got if r["pred_cls"] == r.get("gt_label"))
print(f" классы против геометрии сцены: {acc}/{len(got)} = {100.0*acc/len(got):.0f}%")
print(f" классы против паспортной метки: {accl}/{len(got)} = {100.0*accl/len(got):.0f}%"
f" (расходятся там, где метка не соответствует мешу)")
lab = ["B", "C", "D"]
print(" матрица (строки - истина, столбцы - предсказание):")
print(" " + "".join(f"{c:>5s}" for c in lab))
for a in lab:
row = [sum(1 for r in got if r["gt_cls"] == a and r["pred_cls"] == b) for b in lab]
print(f" {a:3s} " + "".join(f"{v:5d}" for v in row))
print(f" облако собрано: {len(got)}/{len(rows)}")
if times:
print(f"\n время на товар: медиана {np.median(times):.0f} мс, худшее {max(times):.0f} мс")
print(f" такт потока {PITCH_S*1000:.0f} мс -> "
f"{'УКЛАДЫВАЕТСЯ' if np.median(times) < PITCH_S*1000 else 'НЕ УКЛАДЫВАЕТСЯ'} "
f"(запас {PITCH_S*1000 - np.median(times):+.0f} мс)")
if got:
print(f" из них сегментация {np.median([r['t_seg'] for r in got]):.0f} мс, "
f"CREStereo {np.median([r['t_cre'] for r in got]):.0f} мс")
json.dump(rows, open(f"{CAP}/flow_metrics_{TAG}.json", "w"), indent=1, ensure_ascii=False)
print(f"\n -> {CAP}/flow_metrics.json")