init: tianyan-edge edge agent with Ascend NPU inference backend
- Go main process: stream ingestor, infer client, event uploader, heartbeat, config agent, OTA agent - Python inference server rewritten to use CANN ACL (acl module) replacing ultralytics/PyTorch; supports YOLOv8 raw and YOLOv10 NMS-free output formats via OUTPUT_FORMAT env var - systemd services for edge-agent and edge-infer - build/install/package scripts targeting arm64 Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
305
python/infer_server.py
Normal file
305
python/infer_server.py
Normal file
@@ -0,0 +1,305 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Edge inference server — Ascend NPU (CANN ACL) backend.
|
||||
Replaces ultralytics/PyTorch with ACL for Atlas 200I DK2.
|
||||
Socket protocol unchanged: 4-byte big-endian length prefix + JSON body.
|
||||
|
||||
Output formats (OUTPUT_FORMAT env):
|
||||
raw — YOLOv8 style [1, 4+nc, anchors], NMS applied in Python
|
||||
nms_free — YOLOv10 style [1, topk, 6], already NMS'd by model
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import struct
|
||||
import socket
|
||||
import base64
|
||||
import logging
|
||||
import numpy as np
|
||||
import cv2
|
||||
import acl
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
SOCK_PATH = os.getenv("INFER_SOCKET", "/tmp/edge-infer.sock")
|
||||
MODEL_PATH = os.getenv("OM_MODEL", "model.om")
|
||||
CONF_TH = float(os.getenv("CONF_THRESHOLD", "0.5"))
|
||||
IOU_TH = float(os.getenv("IOU_THRESHOLD", "0.45"))
|
||||
DEVICE_ID = int(os.getenv("DEVICE_ID", "0"))
|
||||
NAMES_FILE = os.getenv("NAMES_FILE", "")
|
||||
OUTPUT_FMT = os.getenv("OUTPUT_FORMAT", "raw") # raw | nms_free
|
||||
|
||||
ACL_MEM_MALLOC_NORMAL_ONLY = 0
|
||||
ACL_MEMCPY_HOST_TO_DEVICE = 1
|
||||
ACL_MEMCPY_DEVICE_TO_HOST = 2
|
||||
|
||||
|
||||
def load_names(path):
|
||||
if path and os.path.exists(path):
|
||||
with open(path) as f:
|
||||
return [l.strip() for l in f if l.strip()]
|
||||
return [str(i) for i in range(1000)]
|
||||
|
||||
|
||||
class AclModel:
|
||||
def __init__(self, model_path, device_id):
|
||||
self.device_id = device_id
|
||||
self._init_acl()
|
||||
self._load_model(model_path)
|
||||
self._alloc_outputs()
|
||||
log.info("model loaded path=%s input=%s outputs=%d",
|
||||
model_path, self.input_shape, self.output_num)
|
||||
|
||||
def _init_acl(self):
|
||||
ret = acl.init()
|
||||
assert ret == 0, f"acl.init failed ret={ret}"
|
||||
ret = acl.rt.set_device(self.device_id)
|
||||
assert ret == 0, f"set_device failed ret={ret}"
|
||||
self.context, ret = acl.rt.create_context(self.device_id)
|
||||
assert ret == 0, f"create_context failed ret={ret}"
|
||||
|
||||
def _load_model(self, path):
|
||||
self.model_id, ret = acl.mdl.load_from_file(path)
|
||||
assert ret == 0, f"load_from_file failed ret={ret}"
|
||||
|
||||
self.desc = acl.mdl.create_desc()
|
||||
ret = acl.mdl.get_desc(self.desc, self.model_id)
|
||||
assert ret == 0
|
||||
|
||||
self.input_num = acl.mdl.get_num_inputs(self.desc)
|
||||
self.output_num = acl.mdl.get_num_outputs(self.desc)
|
||||
|
||||
dims, ret = acl.mdl.get_input_dims(self.desc, 0)
|
||||
assert ret == 0
|
||||
self.input_shape = list(dims["dims"]) # [1, 3, H, W]
|
||||
self.input_h = self.input_shape[2]
|
||||
self.input_w = self.input_shape[3]
|
||||
self.input_size = acl.mdl.get_input_size_by_index(self.desc, 0)
|
||||
|
||||
self.output_shapes = []
|
||||
for i in range(self.output_num):
|
||||
d, ret = acl.mdl.get_output_dims(self.desc, i)
|
||||
assert ret == 0
|
||||
self.output_shapes.append(list(d["dims"]))
|
||||
|
||||
def _alloc_outputs(self):
|
||||
self.out_bufs = []
|
||||
self.out_sizes = []
|
||||
for i in range(self.output_num):
|
||||
sz = acl.mdl.get_output_size_by_index(self.desc, i)
|
||||
buf, ret = acl.rt.malloc(sz, ACL_MEM_MALLOC_NORMAL_ONLY)
|
||||
assert ret == 0
|
||||
self.out_bufs.append(buf)
|
||||
self.out_sizes.append(sz)
|
||||
|
||||
def preprocess(self, img_bgr):
|
||||
"""Letterbox → RGB → NCHW float32 [0,1]. Returns blob, scale, pad_top, pad_left."""
|
||||
h0, w0 = img_bgr.shape[:2]
|
||||
scale = min(self.input_h / h0, self.input_w / w0)
|
||||
nh, nw = int(h0 * scale), int(w0 * scale)
|
||||
resized = cv2.resize(img_bgr, (nw, nh), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
canvas = np.full((self.input_h, self.input_w, 3), 114, dtype=np.uint8)
|
||||
pad_top = (self.input_h - nh) // 2
|
||||
pad_left = (self.input_w - nw) // 2
|
||||
canvas[pad_top:pad_top + nh, pad_left:pad_left + nw] = resized
|
||||
|
||||
rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
|
||||
blob = rgb.astype(np.float32) / 255.0
|
||||
blob = np.ascontiguousarray(blob.transpose(2, 0, 1)[np.newaxis]) # NCHW
|
||||
return blob, scale, pad_top, pad_left
|
||||
|
||||
def run(self, blob):
|
||||
"""Push blob to NPU, execute, pull outputs back as numpy arrays."""
|
||||
in_ds = acl.mdl.create_dataset()
|
||||
in_buf, ret = acl.rt.malloc(self.input_size, ACL_MEM_MALLOC_NORMAL_ONLY)
|
||||
assert ret == 0
|
||||
ret = acl.rt.memcpy(in_buf, self.input_size,
|
||||
blob.tobytes(), self.input_size,
|
||||
ACL_MEMCPY_HOST_TO_DEVICE)
|
||||
assert ret == 0
|
||||
db = acl.create_data_buffer(in_buf, self.input_size)
|
||||
_, ret = acl.mdl.add_dataset_buffer(in_ds, db)
|
||||
assert ret == 0
|
||||
|
||||
out_ds = acl.mdl.create_dataset()
|
||||
for i in range(self.output_num):
|
||||
db = acl.create_data_buffer(self.out_bufs[i], self.out_sizes[i])
|
||||
_, ret = acl.mdl.add_dataset_buffer(out_ds, db)
|
||||
assert ret == 0
|
||||
|
||||
ret = acl.mdl.execute(self.model_id, in_ds, out_ds)
|
||||
assert ret == 0, f"mdl.execute failed ret={ret}"
|
||||
|
||||
outputs = []
|
||||
for i in range(self.output_num):
|
||||
sz = self.out_sizes[i]
|
||||
host = np.zeros(sz, dtype=np.uint8)
|
||||
ret = acl.rt.memcpy(host.ctypes.data, sz,
|
||||
self.out_bufs[i], sz,
|
||||
ACL_MEMCPY_DEVICE_TO_HOST)
|
||||
assert ret == 0
|
||||
outputs.append(host.view(np.float32).reshape(self.output_shapes[i]))
|
||||
|
||||
acl.rt.free(in_buf)
|
||||
acl.mdl.destroy_dataset(in_ds)
|
||||
acl.mdl.destroy_dataset(out_ds)
|
||||
return outputs
|
||||
|
||||
def destroy(self):
|
||||
for buf in self.out_bufs:
|
||||
acl.rt.free(buf)
|
||||
acl.mdl.unload(self.model_id)
|
||||
acl.mdl.destroy_desc(self.desc)
|
||||
acl.rt.destroy_context(self.context)
|
||||
acl.rt.reset_device(self.device_id)
|
||||
acl.finalize()
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Post-processing
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _xywh2xyxy(boxes):
|
||||
out = np.empty_like(boxes)
|
||||
out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
|
||||
out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
|
||||
out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
|
||||
out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
|
||||
return out
|
||||
|
||||
|
||||
def _unpad(x1, y1, x2, y2, scale, pad_top, pad_left, orig_h, orig_w):
|
||||
x1 = max(0.0, (x1 - pad_left) / scale)
|
||||
y1 = max(0.0, (y1 - pad_top) / scale)
|
||||
x2 = min(float(orig_w), (x2 - pad_left) / scale)
|
||||
y2 = min(float(orig_h), (y2 - pad_top) / scale)
|
||||
return x1, y1, x2, y2
|
||||
|
||||
|
||||
def postprocess_raw(output, conf_th, iou_th, scale, pad_top, pad_left, orig_h, orig_w):
|
||||
"""YOLOv8 raw output [1, 4+nc, anchors] → detections list."""
|
||||
pred = output[0].T # [anchors, 4+nc]
|
||||
boxes = pred[:, :4] # cx,cy,w,h in input coords
|
||||
scores = pred[:, 4:]
|
||||
|
||||
cls_ids = scores.argmax(axis=1)
|
||||
confs = scores[np.arange(len(scores)), cls_ids]
|
||||
|
||||
mask = confs >= conf_th
|
||||
boxes, confs, cls_ids = boxes[mask], confs[mask], cls_ids[mask]
|
||||
if len(boxes) == 0:
|
||||
return []
|
||||
|
||||
xyxy = _xywh2xyxy(boxes)
|
||||
keep = cv2.dnn.NMSBoxes(xyxy.tolist(), confs.tolist(), conf_th, iou_th)
|
||||
if len(keep) == 0:
|
||||
return []
|
||||
|
||||
dets = []
|
||||
for idx in np.array(keep).flatten():
|
||||
x1, y1, x2, y2 = _unpad(*xyxy[idx], scale, pad_top, pad_left, orig_h, orig_w)
|
||||
dets.append({"class_id": int(cls_ids[idx]),
|
||||
"conf": float(confs[idx]),
|
||||
"bbox": [x1, y1, x2, y2]})
|
||||
return dets
|
||||
|
||||
|
||||
def postprocess_nms_free(output, conf_th, scale, pad_top, pad_left, orig_h, orig_w):
|
||||
"""YOLOv10 NMS-free output [1, topk, 6] (x1,y1,x2,y2,conf,cls) → detections list."""
|
||||
dets = []
|
||||
for row in output[0]:
|
||||
x1, y1, x2, y2, conf, cls_id = row
|
||||
if conf < conf_th:
|
||||
continue
|
||||
x1, y1, x2, y2 = _unpad(x1, y1, x2, y2, scale, pad_top, pad_left, orig_h, orig_w)
|
||||
dets.append({"class_id": int(cls_id),
|
||||
"conf": float(conf),
|
||||
"bbox": [x1, y1, x2, y2]})
|
||||
return dets
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Socket helpers
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def recv_msg(conn):
|
||||
hdr = conn.recv(4)
|
||||
if not hdr:
|
||||
return None
|
||||
length = struct.unpack(">I", hdr)[0]
|
||||
data = b""
|
||||
while len(data) < length:
|
||||
chunk = conn.recv(length - len(data))
|
||||
if not chunk:
|
||||
return None
|
||||
data += chunk
|
||||
return data
|
||||
|
||||
|
||||
def send_msg(conn, payload):
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
conn.sendall(struct.pack(">I", len(data)) + data)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Inference entry
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def infer_one(model, names, msg):
|
||||
jpg = base64.b64decode(msg["jpeg_b64"])
|
||||
arr = np.frombuffer(jpg, dtype=np.uint8)
|
||||
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
|
||||
orig_h, orig_w = img.shape[:2]
|
||||
|
||||
blob, scale, pad_top, pad_left = model.preprocess(img)
|
||||
outputs = model.run(blob)
|
||||
|
||||
if OUTPUT_FMT == "nms_free":
|
||||
raw_dets = postprocess_nms_free(
|
||||
outputs[0], CONF_TH, scale, pad_top, pad_left, orig_h, orig_w)
|
||||
else:
|
||||
raw_dets = postprocess_raw(
|
||||
outputs[0], CONF_TH, IOU_TH, scale, pad_top, pad_left, orig_h, orig_w)
|
||||
|
||||
dets = [{"class": names[d["class_id"]] if d["class_id"] < len(names) else str(d["class_id"]),
|
||||
"conf": d["conf"],
|
||||
"bbox": d["bbox"]} for d in raw_dets]
|
||||
|
||||
return {"stream_id": msg["stream_id"],
|
||||
"device_id": msg["device_id"],
|
||||
"ts": msg["ts"],
|
||||
"detections": dets}
|
||||
|
||||
|
||||
def main():
|
||||
names = load_names(NAMES_FILE)
|
||||
model = AclModel(MODEL_PATH, DEVICE_ID)
|
||||
|
||||
if os.path.exists(SOCK_PATH):
|
||||
os.remove(SOCK_PATH)
|
||||
srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
srv.bind(SOCK_PATH)
|
||||
srv.listen(4)
|
||||
log.info("infer server ready socket=%s device=%d fmt=%s",
|
||||
SOCK_PATH, DEVICE_ID, OUTPUT_FMT)
|
||||
|
||||
try:
|
||||
while True:
|
||||
conn, _ = srv.accept()
|
||||
try:
|
||||
while True:
|
||||
data = recv_msg(conn)
|
||||
if data is None:
|
||||
break
|
||||
msg = json.loads(data.decode("utf-8"))
|
||||
out = infer_one(model, names, msg)
|
||||
send_msg(conn, out)
|
||||
finally:
|
||||
conn.close()
|
||||
finally:
|
||||
model.destroy()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
3
python/requirements.txt
Normal file
3
python/requirements.txt
Normal file
@@ -0,0 +1,3 @@
|
||||
# acl is provided by the CANN SDK on-device, not installed via pip
|
||||
opencv-python-headless>=4.9.0
|
||||
numpy>=1.24.0
|
||||
Reference in New Issue
Block a user