# -*- coding: utf-8 -*-
“”"
Pure vision pipeline extracted from robot_fsm_ticup_float.py
Camera → BGR→NV12 → BPU forward → DFL decode → NMS → mask processing
No FSM, no serial, no web.
“”"
import cv2
import numpy as np
from hobot_dnn import pyeasy_dnn as dnn
import time
# ── Model config ──
MODEL_PATH = “/home/sunrise/yolov5_best_quant(kl.bin”
INPUT_SIZE = 640
REG_MAX = 16
MASK_COEFFS = 32
CONF_THRES = 0.25
NMS_THRES = 0.30
FRAME_W, FRAME_H = 640, 480
CLASSES = [‘blackball’, ‘blueball’, ‘bluezone’,
'redball', 'redzone', 'yellowball'\]
NUM_CLS = len(CLASSES)
BLACK_BALL = 0
BLUE_BALL = 1
BLUE_ZONE = 2
RED_BALL = 3
RED_ZONE = 4
YELLOW_BALL = 5
OWN_ZONE = RED_ZONE # red team
OPPONENT_ZONE = BLUE_ZONE
CLASS_CONF_BIAS = np.zeros(NUM_CLS, dtype=np.float32)
CLASS_CONF_BIAS[BLUE_BALL] = -0.05
CLASS_CONF_BIAS[OWN_ZONE] = 0.10
CLASS_CONF_BIAS[OPPONENT_ZONE] = 0.05
CLASS_CONF_BIAS[YELLOW_BALL] = 0.10
CLASS_CONF_BIAS[BLACK_BALL] = 0.10
ZONE_MASK_EXPAND = 30 # own zone: expand upward
ZONE_BOTTOM_SHRINK = 10 # own zone: shrink downward
OPPONENT_ZONE_EXPAND = 20 # opponent zone: expand all 4 directions
def sigmoid(x):
return 1 / (1 + np.exp(-x))
def bgr2nv12_opencv(image):
"""BGR (640,480,3) uint8 → NV12 (460800,) uint8 for BPU input"""
height, width = image.shape\[0\], image.shape\[1\]
area = height \* width
yuv420p = cv2.cvtColor(image, cv2.COLOR_BGR2YUV_I420).reshape((area \* 3 // 2,))
y = yuv420p\[:area\]
uv_planar = yuv420p\[area:area + area // 2\]
u = uv_planar\[:area // 4\]
v = uv_planar\[area // 4:\]
uv_interleaved = np.zeros(area // 2, dtype=np.uint8)
uv_interleaved\[0::2\] = u
uv_interleaved\[1::2\] = v
nv12 = np.concatenate((y, uv_interleaved))
return nv12
# ── DFL decode ──
_DFL_ARANGE = np.arange(REG_MAX, dtype=np.float32)
def _dfl_decode(bd_raw):
"""
Input: (4\*REG_MAX, h, w) float32 raw box distribution
Output: (4, h, w) float32 dl, dt, dr, db in stride units
"""
h, w = bd_raw.shape\[1\], bd_raw.shape\[2\]
bd = bd_raw.reshape(4, REG_MAX, h \* w)
bd_max = bd.max(axis=1, keepdims=True)
bd_exp = np.exp(bd - bd_max)
bd = bd_exp / bd_exp.sum(axis=1, keepdims=True)
return (bd \* \_DFL_ARANGE.reshape(1, REG_MAX, 1)).sum(axis=1).reshape(4, h, w)
# ── Precomputed grids (3 scale levels) ──
SCALES = [(8, 80), (16, 40), (32, 20)] # (stride, grid_height)
_GRID_GX, _GRID_GY = {}, {}
for _s, _h in SCALES:
gx = (np.arange(\_h, dtype=np.float32) + 0.5) \* \_s
gy = (np.arange(\_h, dtype=np.float32) + 0.5) \* \_s
\_GRID_GX\[(\_s, \_h)\] = gx
\_GRID_GY\[(\_s, \_h)\] = gy
MAX_DETS = 500
def _crop_mask(mask_full, x1, y1, x2, y2,
input_size, proto_h, proto_w, img_w, img_h):
"""
Crop a 160×160 proto mask to bbox region, resize to img coordinates.
mask_full: (160, 160) float32 sigmoid-activated proto segment
Returns: (img_h, img_w) uint8 binary mask
"""
scale = input_size / proto_h
px1 = max(0, int(x1 / scale))
py1 = max(0, int(y1 / scale))
px2 = min(proto_w, int(x2 / scale) + 1)
py2 = min(proto_h, int(y2 / scale) + 1)
crop = mask_full\[py1:py2, px1:px2\]
if crop.size == 0:
return None
dx1 = int(px1 \* img_w / proto_w)
dy1 = int(py1 \* img_h / proto_h)
dx2 = int(px2 \* img_w / proto_w)
dy2 = int(py2 \* img_h / proto_h)
bw = max(1, dx2 - dx1)
bh = max(1, dy2 - dy1)
crop = cv2.resize(crop, (bw, bh))
crop = (crop > 0.5).astype(np.uint8)
full_mask = np.zeros((img_h, img_w), dtype=np.uint8)
src_y1 = max(0, -dy1)
src_x1 = max(0, -dx1)
src_y2 = min(bh, img_h - dy1)
src_x2 = min(bw, img_w - dx1)
if src_y2 > src_y1 and src_x2 > src_x1:
full_mask\[dy1+src_y1:dy1+src_y2, dx1+src_x1:dx1+src_x2\] = crop\[src_y1:src_y2, src_x1:src_x2\]
return full_mask
def process_output_seg(outputs, img_w, img_h, need_mask=True):
"""
Parse 10-output float model → list of dicts {cls, cx, cy, w, h, conf, mask, mask_dilated, filter_contour, ...}
outputs layout (10 tensors, NCHW, float32 buf):
\[0\] cls (NUM_CLS, 80, 80) ← stride 8, 6400 cells
\[1\] box (4\*REG_MAX, 80, 80) ← stride 8
\[2\] mc (MASK_COEFFS, 80, 80) ← stride 8
\[3\] cls (NUM_CLS, 40, 40) ← stride 16, 1600 cells
\[4\] box (4\*REG_MAX, 40, 40) ← stride 16
\[5\] mc (MASK_COEFFS, 40, 40) ← stride 16
\[6\] cls (NUM_CLS, 20, 20) ← stride 32, 400 cells
\[7\] box (4\*REG_MAX, 20, 20) ← stride 32
\[8\] mc (MASK_COEFFS, 20, 20) ← stride 32
\[9\] proto (MASK_COEFFS, 160, 160) ← mask prototype
Data flow per scale:
outputs\[off\].buffer (NCHW) → reshape(NUM_CLS, h, h) → sigmoid → per-class threshold
outputs\[off+1\].buffer → \_dfl_decode → dl,dt,dr,db
outputs\[off+2\].buffer → reshape(MASK_COEFFS, h, h) → extract per-det coeffs
"""
rw = img_w / INPUT_SIZE
rh = img_h / INPUT_SIZE
xyxy_buf = np.zeros((MAX_DETS, 4), dtype=np.float32)
scores_buf = np.zeros(MAX_DETS, dtype=np.float32)
cls_buf = np.zeros(MAX_DETS, dtype=np.uint8)
mask_buf = np.zeros((MAX_DETS, MASK_COEFFS), dtype=np.float32)
det_idx = 0
\# --- Proto (last output) ---
proto = None
proto_h = proto_w = 0
if need_mask:
try:
proto = outputs\[9\].buffer.astype(np.float32).reshape(1, MASK_COEFFS, 160, 160)\[0\]
proto_h, proto_w = 160, 160
except:
pass
for si, (stride, h) in enumerate(SCALES):
off = si \* 3
cd = outputs\[off\].buffer.astype(np.float32).reshape(NUM_CLS, h, h)
bd = \_dfl_decode(outputs\[off + 1\].buffer.astype(np.float32).reshape(4 \* REG_MAX, h, h))
cprob = sigmoid(cd)
mc = outputs\[off + 2\].buffer.astype(np.float32).reshape(MASK_COEFFS, h, h)
dl_all = bd\[0\] \* stride; dt_all = bd\[1\] \* stride
dr_all = bd\[2\] \* stride; db_all = bd\[3\] \* stride
gx_grid = \_GRID_GX\[(stride, h)\]; gy_grid = \_GRID_GY\[(stride, h)\]
for c in range(NUM_CLS):
thresh = CONF_THRES - CLASS_CONF_BIAS\[c\]
above = cprob\[c\] > thresh
if not above.any():
continue
n = int(above.sum())
if det_idx + n > MAX_DETS:
n = MAX_DETS - det_idx
if n <= 0:
continue
ys, xs = np.where(above)
ys = ys\[:n\]; xs = xs\[:n\]
scores_buf\[det_idx:det_idx + n\] = cprob\[c, ys, xs\]
cls_buf\[det_idx:det_idx + n\] = c
\# yolo bbox decode: grid center ± distribution offset
xyxy_buf\[det_idx:det_idx + n, 0\] = gx_grid\[xs\] - dl_all\[ys, xs\]
xyxy_buf\[det_idx:det_idx + n, 1\] = gy_grid\[ys\] - dt_all\[ys, xs\]
xyxy_buf\[det_idx:det_idx + n, 2\] = gx_grid\[xs\] + dr_all\[ys, xs\]
xyxy_buf\[det_idx:det_idx + n, 3\] = gy_grid\[ys\] + db_all\[ys, xs\]
if need_mask and c in (OWN_ZONE, OPPONENT_ZONE):
mask_buf\[det_idx:det_idx + n\] = mc\[:, ys, xs\].transpose()
det_idx += n
if det_idx == 0:
return \[\]
xyxy = xyxy_buf\[:det_idx\]; scores = scores_buf\[:det_idx\]
cls = cls_buf\[:det_idx\]; masks_raw = mask_buf\[:det_idx\]
\# Clip to model input size
xyxy\[:, 0\] = np.clip(xyxy\[:, 0\], 0, INPUT_SIZE); xyxy\[:, 1\] = np.clip(xyxy\[:, 1\], 0, INPUT_SIZE)
xyxy\[:, 2\] = np.clip(xyxy\[:, 2\], 0, INPUT_SIZE); xyxy\[:, 3\] = np.clip(xyxy\[:, 3\], 0, INPUT_SIZE)
valid = (xyxy\[:, 2\] > xyxy\[:, 0\]) & (xyxy\[:, 3\] > xyxy\[:, 1\])
xyxy = xyxy\[valid\]; scores = scores\[valid\]; cls = cls\[valid\]; masks_raw = masks_raw\[valid\]
if len(xyxy) == 0:
return \[\]
\# NMS (class-agnostic)
idx_nms = cv2.dnn.NMSBoxes(xyxy.tolist(), scores.tolist(), CONF_THRES, NMS_THRES)
if len(idx_nms) == 0:
return \[\]
keep = idx_nms.flatten()
\# Scale bbox to image coordinates
x1_k = xyxy\[keep, 0\] \* rw; y1_k = xyxy\[keep, 1\] \* rh
x2_k = xyxy\[keep, 2\] \* rw; y2_k = xyxy\[keep, 3\] \* rh
sc_k = scores\[keep\]; cl_k = cls\[keep\]; mk_k = masks_raw\[keep\]
\# Batch mask decode: coeffs @ proto → sigmoid → (N, 160, 160)
zone_keep = \[i for i in range(len(keep)) if int(cl_k\[i\]) in (OWN_ZONE, OPPONENT_ZONE)\]
batch_masks = None
if proto is not None and zone_keep:
try:
z_idx = np.array(zone_keep)
masks_flat = mk_k\[z_idx\].astype(np.float32) @ proto.reshape(MASK_COEFFS, -1)
batch_masks = sigmoid(masks_flat).reshape(len(zone_keep), proto_h, proto_w)
except:
pass
final_dets = \[\]
for i in range(len(keep)):
cid = int(cl_k\[i\])
det = {
'cls': cid,
'cx': float((x1_k\[i\] + x2_k\[i\]) \* 0.5),
'cy': float((y1_k\[i\] + y2_k\[i\]) \* 0.5),
'w': float(x2_k\[i\] - x1_k\[i\]),
'h': float(y2_k\[i\] - y1_k\[i\]),
'conf': float(sc_k\[i\]),
'mask': None, 'mask_proto': None,
'mask_coeffs': None,
}
\# Zone mask processing
if cid in (OWN_ZONE, OPPONENT_ZONE) and batch_masks is not None:
zi = zone_keep.index(i) if i in zone_keep else -1
if zi >= 0:
det\['mask_proto'\] = batch_masks\[zi\]
try:
x1_m = float(xyxy\[keep\[i\], 0\]); y1_m = float(xyxy\[keep\[i\], 1\])
x2_m = float(xyxy\[keep\[i\], 2\]); y2_m = float(xyxy\[keep\[i\], 3\])
det\['mask'\] = \_crop_mask(batch_masks\[zi\], x1_m, y1_m, x2_m, y2_m,
INPUT_SIZE, proto_h, proto_w, img_w, img_h)
if det\['mask'\] is not None and det\['mask'\].any():
contours, \_ = cv2.findContours(det\['mask'\], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours:
cnt = max(contours, key=cv2.contourArea)
M = cv2.moments(cnt)
if M\['m00'\] > 0:
\# Mask centroid overrides bbox center
det\['cx'\] = M\['m10'\] / M\['m00'\]
det\['cy'\] = M\['m01'\] / M\['m00'\]
epsilon = 0.005 \* cv2.arcLength(cnt, True)
approx = cv2.approxPolyDP(cnt, epsilon, True)
pts = approx.reshape(-1, 2)
det\['mask_contour'\] = pts
det\['mask_x1'\] = int(pts\[:, 0\].min())
det\['mask_y1'\] = int(pts\[:, 1\].min())
det\['mask_x2'\] = int(pts\[:, 0\].max())
det\['mask_y2'\] = int(pts\[:, 1\].max())
dm = det\['mask'\].copy()
\# Opponent zone: dilate all 4 directions
if cid == OPPONENT_ZONE and OPPONENT_ZONE_EXPAND > 0:
n = OPPONENT_ZONE_EXPAND
up = np.roll(dm, -n, axis=0); up\[-n:, :\] = 0
down = np.roll(dm, n, axis=0); down\[:n, :\] = 0
left = np.roll(dm, -n, axis=1); left\[:, -n:\] = 0
right = np.roll(dm, n, axis=1); right\[:, :n\] = 0
dm = np.maximum.reduce(\[dm, up, down, left, right\])
elif cid == OWN_ZONE:
\# Own zone: expand upward, shrink downward
if ZONE_MASK_EXPAND > 0:
n = ZONE_MASK_EXPAND
up = np.roll(dm, -n, axis=0)
up\[-n:, :\] = 0
dm = np.maximum(dm, up)
if ZONE_BOTTOM_SHRINK > 0:
b = ZONE_BOTTOM_SHRINK
down = np.roll(dm, b, axis=0)
down\[:b, :\] = 0
dm = np.minimum(dm, down)
det\['mask_dilated'\] = dm
dc, \_ = cv2.findContours(dm, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if dc:
det\['filter_contour'\] = max(dc, key=cv2.contourArea)
else:
det\['mask'\] = None
else:
det\['mask'\] = None
except:
pass
final_dets.append(det)
return final_dets
# ── Main loop (camera → inference → dets, no FSM) ──
def main():
print(f"Loading model {MODEL_PATH}...")
models = dnn.load(MODEL_PATH)
print("Model loaded.")
cap = cv2.VideoCapture(1)
if not cap.isOpened():
cap.release(); cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("ERROR: Cannot open camera")
return
cap.set(3, FRAME_W); cap.set(4, FRAME_H)
cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, 0)
cap.set(cv2.CAP_PROP_EXPOSURE, 0.3)
\# Warmup
for \_ in range(5):
success, frame = cap.read()
if success:
img = cv2.resize(frame, (INPUT_SIZE, INPUT_SIZE))
nv12 = bgr2nv12_opencv(img)
outputs = models\[0\].forward(\[nv12\])
process_output_seg(outputs, frame.shape\[1\], frame.shape\[0\], need_mask=False)
print("Ready.")
fps_t0 = time.time()
fps_count = 0
frame_fail = 0
try:
while True:
success, frame = cap.read()
if not success:
frame_fail += 1
if frame_fail <= 3 or frame_fail % 100 == 0:
print(f"WARN: frame read fail (total {frame_fail})")
time.sleep(0.01)
if frame_fail == 10:
cap.release()
import os
os.system("sudo rmmod uvcvideo 2>/dev/null; sudo modprobe uvcvideo 2>/dev/null; sleep 0.5")
cap = cv2.VideoCapture(1)
if not cap.isOpened():
cap.release(); cap = cv2.VideoCapture(0)
cap.set(3, FRAME_W); cap.set(4, FRAME_H)
cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, 0); cap.set(cv2.CAP_PROP_EXPOSURE, 0.3)
print("RECONNECT: reopened camera")
frame_fail = 0
continue
\# ── Data flow ──
img = cv2.resize(frame, (INPUT_SIZE, INPUT_SIZE)) # (640,480,3) → (640,640,3)
nv12 = bgr2nv12_opencv(img) # BGR → NV12 (460800,)
outputs = models\[0\].forward(\[nv12\]) # BPU forward, returns 10 tensors
dets = process_output_seg(outputs, frame.shape\[1\], frame.shape\[0\],
need_mask=True) # parse outputs → det dicts
\# FPS
fps_count += 1
if fps_count >= 30:
fps_val = 30 / (time.time() - fps_t0)
fps_t0 = time.time()
fps_count = 0
print(f"FPS: {fps_val:.1f} dets: {len(dets)}")
except KeyboardInterrupt:
pass
finally:
cap.release()
print("Done.")
if _name_ == “_main_”:
main()