fyah
f6a85d7dfc
- 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>
306 lines
10 KiB
Python
306 lines
10 KiB
Python
#!/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()
|