"""ЭТАП 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")