commit f6a85d7dfcd4400bf8c9bfd177724aab15cf9593 Author: fyah <727475183@qq.com> Date: Wed May 6 21:49:04 2026 +0800 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) diff --git a/.env.template b/.env.template new file mode 100644 index 0000000..0ea9efb --- /dev/null +++ b/.env.template @@ -0,0 +1,10 @@ +EDGE_ID=edge-demo-001 +CLOUD_URL=http://localhost:8004 +EDGE_TOKEN= + +# 多路用逗号分隔，留空则使用合成演示流 +RTSP_URLS= + +YOLO_MODEL=yolov8n.pt +INFER_FPS=5 +CONF_THRESHOLD=0.5 diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..1a80bb8 --- /dev/null +++ b/Makefile @@ -0,0 +1,16 @@ +PYTHON=python3 + +build: + bash scripts/build.sh + +install: + bash scripts/install.sh + +package: + bash scripts/package-usb.sh + +uninstall: + bash scripts/uninstall.sh + +chmod: + chmod +x scripts/*.sh diff --git a/cmd/edge-agent/main.go b/cmd/edge-agent/main.go new file mode 100644 index 0000000..664342d --- /dev/null +++ b/cmd/edge-agent/main.go @@ -0,0 +1,50 @@ +package main + +import ( + "context" + "flag" + "log" + "os" + "os/signal" + "sync" + "syscall" + + "tianyan-edge/internal/config" + "tianyan-edge/internal/control" + "tianyan-edge/internal/event" + "tianyan-edge/internal/infer" + "tianyan-edge/internal/stream" +) + +func main() { + cfgPath := flag.String("config", "/opt/tianyan-edge/config/edge.yaml", "config file") + flag.Parse() + + cfg := config.Load(*cfgPath) + log.Printf("edge-agent start id=%s cloud=%s", cfg.EdgeID, cfg.CloudURL) + + ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer cancel() + + frames := make(chan stream.Frame, 50) + events := make(chan event.SuspectedEvent, 200) + + var wg sync.WaitGroup + run := func(fn func(context.Context)) { + wg.Add(1) + go func() { + defer wg.Done() + fn(ctx) + }() + } + + run(stream.NewIngestor(cfg, frames).Run) + run(infer.NewClient(cfg, frames, events).Run) + run(event.NewUploader(cfg, events).Run) + run(control.NewHeartbeat(cfg).Run) + run(control.NewConfigAgent(cfg).Run) + run(control.NewOTAAgent(cfg).Run) + + wg.Wait() + log.Println("edge-agent stop") +} diff --git a/config/edge.yaml.template b/config/edge.yaml.template new file mode 100644 index 0000000..bcdaa42 --- /dev/null +++ b/config/edge.yaml.template @@ -0,0 +1,18 @@ +edge_id: edge-demo-001 +cloud_url: http://localhost:8004 +edge_token: "" + +# 5 路示例 +rtsp_urls: + - rtsp://user:pass@camera1/stream1 + - rtsp://user:pass@camera2/stream1 + - rtsp://user:pass@camera3/stream1 + - rtsp://user:pass@camera4/stream1 + - rtsp://user:pass@camera5/stream1 + +infer_socket: /tmp/edge-infer.sock +infer_fps: 5 +infer_workers: 3 +conf_threshold: 0.5 +dedup_window_sec: 30 +version: 1.0.0 diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..4bb9890 --- /dev/null +++ b/go.mod @@ -0,0 +1,7 @@ +module tianyan-edge + +go 1.22 + +require gopkg.in/yaml.v3 v3.0.1 + +require gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect diff --git a/internal/config/config.go b/internal/config/config.go new file mode 100644 index 0000000..ced1260 --- /dev/null +++ b/internal/config/config.go @@ -0,0 +1,45 @@ +package config + +import ( + "os" + + "gopkg.in/yaml.v3" +) + +type Config struct { + EdgeID string `yaml:"edge_id"` + CloudURL string `yaml:"cloud_url"` + EdgeToken string `yaml:"edge_token"` + RTSPURLs []string `yaml:"rtsp_urls"` + InferSocket string `yaml:"infer_socket"` + InferFPS int `yaml:"infer_fps"` + InferWorkers int `yaml:"infer_workers"` + ConfThreshold float64 `yaml:"conf_threshold"` + DedupWindowSec int `yaml:"dedup_window_sec"` + Version string `yaml:"version"` +} + +func Load(path string) *Config { + cfg := &Config{ + EdgeID: "edge-demo-001", + CloudURL: "http://localhost:8004", + InferSocket: "/tmp/edge-infer.sock", + InferFPS: 5, + InferWorkers: 2, + ConfThreshold: 0.5, + DedupWindowSec: 30, + Version: "1.0.0", + } + data, err := os.ReadFile(path) + if err != nil { + return cfg + } + _ = yaml.Unmarshal(data, cfg) + if cfg.InferWorkers <= 0 { + cfg.InferWorkers = 1 + } + if cfg.DedupWindowSec <= 0 { + cfg.DedupWindowSec = 30 + } + return cfg +} diff --git a/internal/control/config_agent.go b/internal/control/config_agent.go new file mode 100644 index 0000000..88ee804 --- /dev/null +++ b/internal/control/config_agent.go @@ -0,0 +1,59 @@ +package control + +import ( + "context" + "encoding/json" + "fmt" + "log" + "net/http" + "time" + + "tianyan-edge/internal/config" +) + +type ConfigAgent struct { + cfg *config.Config +} + +func NewConfigAgent(cfg *config.Config) *ConfigAgent { + return &ConfigAgent{cfg: cfg} +} + +func (a *ConfigAgent) Run(ctx context.Context) { + client := &http.Client{Timeout: 5 * time.Second} + ticker := time.NewTicker(60 * time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + a.pull(client) + } + } +} + +func (a *ConfigAgent) pull(client *http.Client) { + url := fmt.Sprintf("%s/api/v1/edge/config?edge_id=%s", a.cfg.CloudURL, a.cfg.EdgeID) + req, _ := http.NewRequest(http.MethodGet, url, nil) + req.Header.Set("Authorization", "Bearer "+a.cfg.EdgeToken) + resp, err := client.Do(req) + if err != nil { + log.Printf("config pull: %v", err) + return + } + defer resp.Body.Close() + if resp.StatusCode != 200 { + return + } + var data map[string]any + if err := json.NewDecoder(resp.Body).Decode(&data); err != nil { + return + } + if v, ok := data["infer_fps"].(float64); ok && int(v) > 0 { + a.cfg.InferFPS = int(v) + } + if v, ok := data["conf_threshold"].(float64); ok { + a.cfg.ConfThreshold = v + } +} diff --git a/internal/control/heartbeat.go b/internal/control/heartbeat.go new file mode 100644 index 0000000..42d3aef --- /dev/null +++ b/internal/control/heartbeat.go @@ -0,0 +1,54 @@ +package control + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "log" + "net/http" + "time" + + "tianyan-edge/internal/config" +) + +type Heartbeat struct { + cfg *config.Config +} + +func NewHeartbeat(cfg *config.Config) *Heartbeat { + return &Heartbeat{cfg: cfg} +} + +func (h *Heartbeat) Run(ctx context.Context) { + client := &http.Client{Timeout: 5 * time.Second} + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + h.send(client) + } + } +} + +func (h *Heartbeat) send(client *http.Client) { + payload := map[string]any{ + "edge_id": h.cfg.EdgeID, + "version": h.cfg.Version, + "ts": float64(time.Now().UnixMilli()) / 1000.0, + } + body, _ := json.Marshal(payload) + url := fmt.Sprintf("%s/api/v1/edge/heartbeat", h.cfg.CloudURL) + req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+h.cfg.EdgeToken) + resp, err := client.Do(req) + if err != nil { + log.Printf("heartbeat: %v", err) + return + } + resp.Body.Close() +} diff --git a/internal/control/ota_agent.go b/internal/control/ota_agent.go new file mode 100644 index 0000000..d9194f1 --- /dev/null +++ b/internal/control/ota_agent.go @@ -0,0 +1,115 @@ +package control + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "os" + "path/filepath" + "time" + + "tianyan-edge/internal/config" +) + +type OTAAgent struct { + cfg *config.Config +} + +func NewOTAAgent(cfg *config.Config) *OTAAgent { + return &OTAAgent{cfg: cfg} +} + +func (o *OTAAgent) Run(ctx context.Context) { + client := &http.Client{Timeout: 30 * time.Second} + ticker := time.NewTicker(10 * time.Minute) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + if inMaintenanceWindow() { + o.check(client) + } + } + } +} + +func inMaintenanceWindow() bool { + h := time.Now().Hour() + return h >= 22 || h < 6 +} + +func (o *OTAAgent) check(client *http.Client) { + url := fmt.Sprintf("%s/api/v1/edge/update/manifest?edge_id=%s¤t_version=%s", + o.cfg.CloudURL, o.cfg.EdgeID, o.cfg.Version) + req, _ := http.NewRequest(http.MethodGet, url, nil) + req.Header.Set("Authorization", "Bearer "+o.cfg.EdgeToken) + resp, err := client.Do(req) + if err != nil { + log.Printf("ota manifest: %v", err) + return + } + defer resp.Body.Close() + if resp.StatusCode != 200 { + return + } + var mf map[string]string + if err := json.NewDecoder(resp.Body).Decode(&mf); err != nil { + return + } + target := mf["target_version"] + if target == "" || target == o.cfg.Version { + return + } + if err := o.downloadAndVerify(client, mf); err != nil { + o.report(client, target, "failed") + return + } + o.report(client, target, "verified") +} + +func (o *OTAAgent) downloadAndVerify(client *http.Client, mf map[string]string) error { + url, target := mf["package_url"], mf["target_version"] + expected := mf["sha256"] + if url == "" || expected == "" { + return fmt.Errorf("invalid manifest") + } + resp, err := client.Get(url) + if err != nil { + return err + } + defer resp.Body.Close() + data, err := io.ReadAll(resp.Body) + if err != nil { + return err + } + sum := sha256.Sum256(data) + actual := hex.EncodeToString(sum[:]) + if actual != expected { + return fmt.Errorf("checksum mismatch") + } + _ = os.MkdirAll("/opt/tianyan-edge/staging", 0o755) + path := filepath.Join("/opt/tianyan-edge/staging", target+".tar.gz") + return os.WriteFile(path, data, 0o644) +} + +func (o *OTAAgent) report(client *http.Client, version, status string) { + payload := map[string]any{"edge_id": o.cfg.EdgeID, "version": version, + "status": status, "ts": float64(time.Now().UnixMilli()) / 1000.0} + body, _ := json.Marshal(payload) + url := fmt.Sprintf("%s/api/v1/edge/update/report", o.cfg.CloudURL) + req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+o.cfg.EdgeToken) + resp, err := client.Do(req) + if err == nil { + resp.Body.Close() + } +} diff --git a/internal/event/types.go b/internal/event/types.go new file mode 100644 index 0000000..43044ee --- /dev/null +++ b/internal/event/types.go @@ -0,0 +1,13 @@ +package event + +type SuspectedEvent struct { + EventID string `json:"event_id"` + EdgeID string `json:"edge_id"` + DeviceID string `json:"device_id"` + StreamURL string `json:"stream_url"` + TS float64 `json:"ts"` + Class string `json:"class"` + Conf float64 `json:"conf"` + BBox []float64 `json:"bbox"` + ImageB64 string `json:"image_b64"` +} diff --git a/internal/event/uploader.go b/internal/event/uploader.go new file mode 100644 index 0000000..70f3872 --- /dev/null +++ b/internal/event/uploader.go @@ -0,0 +1,80 @@ +package event + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "log" + "net/http" + "time" + + "tianyan-edge/internal/config" +) + +const maxBuffer = 500 + +type Uploader struct { + cfg *config.Config + events <-chan SuspectedEvent + buf []SuspectedEvent +} + +func NewUploader(cfg *config.Config, events <-chan SuspectedEvent) *Uploader { + return &Uploader{cfg: cfg, events: events} +} + +func (u *Uploader) Run(ctx context.Context) { + client := &http.Client{Timeout: 10 * time.Second} + ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case ev := <-u.events: + if !u.upload(client, ev) { + u.buffer(ev) + } + case <-ticker.C: + u.flushBuffer(client) + } + } +} + +func (u *Uploader) buffer(ev SuspectedEvent) { + if len(u.buf) < maxBuffer { + u.buf = append(u.buf, ev) + } else { + log.Println("uploader: offline buffer full, dropping event") + } +} + +func (u *Uploader) flushBuffer(client *http.Client) { + remaining := u.buf[:0] + for _, ev := range u.buf { + if !u.upload(client, ev) { + remaining = append(remaining, ev) + } + } + u.buf = remaining +} + +func (u *Uploader) upload(client *http.Client, ev SuspectedEvent) bool { + url := fmt.Sprintf("%s/api/v1/edge/events/suspected", u.cfg.CloudURL) + body, _ := json.Marshal(ev) + req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+u.cfg.EdgeToken) + resp, err := client.Do(req) + if err != nil { + log.Printf("uploader: upload error: %v", err) + return false + } + defer resp.Body.Close() + ok := resp.StatusCode == 200 || resp.StatusCode == 201 + if !ok { + log.Printf("uploader: upload failed status=%d", resp.StatusCode) + } + return ok +} diff --git a/internal/infer/client.go b/internal/infer/client.go new file mode 100644 index 0000000..97b61a3 --- /dev/null +++ b/internal/infer/client.go @@ -0,0 +1,209 @@ +package infer + +import ( + "context" + "crypto/rand" + "encoding/base64" + "encoding/binary" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "log" + "net" + "sync" + "time" + + "tianyan-edge/internal/config" + "tianyan-edge/internal/event" + "tianyan-edge/internal/stream" +) + +type frameMsg struct { + StreamID int `json:"stream_id"` + DeviceID string `json:"device_id"` + URL string `json:"url"` + TS float64 `json:"ts"` + JPEGB64 string `json:"jpeg_b64"` +} + +type detection struct { + Class string `json:"class"` + Conf float64 `json:"conf"` + BBox []float64 `json:"bbox"` +} + +type resultMsg struct { + StreamID int `json:"stream_id"` + DeviceID string `json:"device_id"` + TS float64 `json:"ts"` + Detections []detection `json:"detections"` +} + +type dedupState struct { + mu sync.Mutex + seen map[string]float64 +} + +type Client struct { + cfg *config.Config + frames <-chan stream.Frame + events chan<- event.SuspectedEvent + dedup *dedupState +} + +func NewClient(cfg *config.Config, frames <-chan stream.Frame, events chan<- event.SuspectedEvent) *Client { + return &Client{ + cfg: cfg, + frames: frames, + events: events, + dedup: &dedupState{seen: make(map[string]float64)}, + } +} + +func (c *Client) Run(ctx context.Context) { + var wg sync.WaitGroup + for i := 0; i < c.cfg.InferWorkers; i++ { + wg.Add(1) + go func(workerID int) { + defer wg.Done() + c.workerLoop(ctx, workerID) + }(i) + } + wg.Wait() +} + +func (c *Client) workerLoop(ctx context.Context, workerID int) { + for { + select { + case <-ctx.Done(): + return + default: + } + conn, err := net.Dial("unix", c.cfg.InferSocket) + if err != nil { + log.Printf("infer[w%d]: socket connect failed: %v, retry 2s", workerID, err) + time.Sleep(2 * time.Second) + continue + } + log.Printf("infer[w%d]: socket connected", workerID) + if !c.loop(ctx, conn, workerID) { + conn.Close() + return + } + conn.Close() + time.Sleep(500 * time.Millisecond) + } +} + +func (c *Client) loop(ctx context.Context, conn net.Conn, workerID int) bool { + for { + select { + case <-ctx.Done(): + return false + case f := <-c.frames: + if err := c.send(conn, f); err != nil { + log.Printf("infer[w%d]: send error: %v", workerID, err) + return true + } + res, err := c.recv(conn) + if err != nil { + log.Printf("infer[w%d]: recv error: %v", workerID, err) + return true + } + c.emitEvents(f, res) + } + } +} + +func (c *Client) emitEvents(f stream.Frame, res *resultMsg) { + for _, d := range res.Detections { + if c.isDuplicate(f.DeviceID, d.Class, f.TS) { + continue + } + ev := event.SuspectedEvent{ + EventID: mustUUID(), + EdgeID: c.cfg.EdgeID, + DeviceID: f.DeviceID, + StreamURL: f.URL, + TS: f.TS, + Class: d.Class, + Conf: d.Conf, + BBox: d.BBox, + ImageB64: base64.StdEncoding.EncodeToString(f.JPEG), + } + select { + case c.events <- ev: + default: + log.Println("infer: event queue full, dropping") + } + } +} + +func (c *Client) isDuplicate(deviceID, className string, ts float64) bool { + window := float64(c.cfg.DedupWindowSec) + key := fmt.Sprintf("%s|%s", deviceID, className) + c.dedup.mu.Lock() + defer c.dedup.mu.Unlock() + last, ok := c.dedup.seen[key] + if ok && ts-last < window { + return true + } + c.dedup.seen[key] = ts + if len(c.dedup.seen) > 5000 { + for k, v := range c.dedup.seen { + if ts-v > 2*window { + delete(c.dedup.seen, k) + } + } + } + return false +} + +func (c *Client) send(conn net.Conn, f stream.Frame) error { + msg := frameMsg{StreamID: f.StreamID, DeviceID: f.DeviceID, URL: f.URL, + TS: f.TS, JPEGB64: base64.StdEncoding.EncodeToString(f.JPEG)} + data, err := json.Marshal(msg) + if err != nil { + return err + } + return writeMsg(conn, data) +} + +func (c *Client) recv(conn net.Conn) (*resultMsg, error) { + data, err := readMsg(conn) + if err != nil { + return nil, err + } + var r resultMsg + return &r, json.Unmarshal(data, &r) +} + +func writeMsg(conn net.Conn, data []byte) error { + hdr := make([]byte, 4) + binary.BigEndian.PutUint32(hdr, uint32(len(data))) + if _, err := conn.Write(hdr); err != nil { + return err + } + _, err := conn.Write(data) + return err +} + +func readMsg(conn net.Conn) ([]byte, error) { + hdr := make([]byte, 4) + if _, err := io.ReadFull(conn, hdr); err != nil { + return nil, err + } + size := binary.BigEndian.Uint32(hdr) + buf := make([]byte, size) + _, err := io.ReadFull(conn, buf) + return buf, err +} + +func mustUUID() string { + b := make([]byte, 16) + _, _ = rand.Read(b) + return hex.EncodeToString(b[:4]) + "-" + hex.EncodeToString(b[4:6]) + "-" + + hex.EncodeToString(b[6:8]) + "-" + hex.EncodeToString(b[8:10]) + "-" + + hex.EncodeToString(b[10:]) +} diff --git a/internal/stream/ingestor.go b/internal/stream/ingestor.go new file mode 100644 index 0000000..38825a1 --- /dev/null +++ b/internal/stream/ingestor.go @@ -0,0 +1,129 @@ +package stream + +import ( + "bytes" + "context" + "fmt" + "io" + "log" + "os/exec" + "sync" + "time" + + "tianyan-edge/internal/config" +) + +type Frame struct { + StreamID int + DeviceID string + URL string + JPEG []byte + TS float64 +} + +type Ingestor struct { + cfg *config.Config + frames chan<- Frame +} + +func NewIngestor(cfg *config.Config, frames chan<- Frame) *Ingestor { + return &Ingestor{cfg: cfg, frames: frames} +} + +func (ing *Ingestor) Run(ctx context.Context) { + if len(ing.cfg.RTSPURLs) == 0 { + log.Println("stream: no RTSP URLs configured, idle") + <-ctx.Done() + return + } + var wg sync.WaitGroup + for i, url := range ing.cfg.RTSPURLs { + wg.Add(1) + go func(idx int, u string) { + defer wg.Done() + ing.streamLoop(ctx, idx, u) + }(i, url) + } + wg.Wait() +} + +func (ing *Ingestor) streamLoop(ctx context.Context, idx int, url string) { + deviceID := fmt.Sprintf("cam-%03d", idx) + fps := ing.cfg.InferFPS + if fps <= 0 { + fps = 5 + } + for { + select { + case <-ctx.Done(): + return + default: + } + cmd := exec.CommandContext(ctx, "ffmpeg", + "-rtsp_transport", "tcp", "-i", url, + "-vf", fmt.Sprintf("fps=%d", fps), + "-f", "image2pipe", "-vcodec", "mjpeg", "-q:v", "5", "-", + ) + stdout, err := cmd.StdoutPipe() + if err != nil { + log.Printf("stream[%d] stdout pipe failed: %v, retry 10s", idx, err) + time.Sleep(10 * time.Second) + continue + } + if err := cmd.Start(); err != nil { + log.Printf("stream[%d] ffmpeg start failed: %v, retry 10s", idx, err) + time.Sleep(10 * time.Second) + continue + } + log.Printf("stream[%d] connected url=%s", idx, url) + ing.readFrames(ctx, stdout, idx, deviceID, url) + cmd.Wait() + select { + case <-ctx.Done(): + return + default: + log.Printf("stream[%d] lost, retry 5s", idx) + time.Sleep(5 * time.Second) + } + } +} + +func (ing *Ingestor) readFrames(ctx context.Context, r io.Reader, idx int, deviceID, url string) { + soi, eoi := []byte{0xFF, 0xD8}, []byte{0xFF, 0xD9} + buf, tmp := make([]byte, 0, 1<<20), make([]byte, 32768) + for { + select { + case <-ctx.Done(): + return + default: + } + n, err := r.Read(tmp) + if n > 0 { + buf = append(buf, tmp[:n]...) + for { + s := bytes.Index(buf, soi) + if s < 0 { + buf = buf[:0] + break + } + buf = buf[s:] + e := bytes.Index(buf[2:], eoi) + if e < 0 { + break + } + end := e + 4 + frame := make([]byte, end) + copy(frame, buf[:end]) + buf = buf[end:] + select { + case ing.frames <- Frame{StreamID: idx, DeviceID: deviceID, URL: url, + JPEG: frame, TS: float64(time.Now().UnixMilli()) / 1000.0}: + default: + } + } + } + if err != nil { + return + } + } +} diff --git a/python/infer_server.py b/python/infer_server.py new file mode 100644 index 0000000..3dedb51 --- /dev/null +++ b/python/infer_server.py @@ -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() diff --git a/python/requirements.txt b/python/requirements.txt new file mode 100644 index 0000000..137d90c --- /dev/null +++ b/python/requirements.txt @@ -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 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..b33a96c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +aiohttp>=3.9.0 +opencv-python-headless>=4.9.0 +ultralytics>=8.2.0 +psutil>=5.9.0 +python-dotenv>=1.0.0 +numpy>=1.24.0 diff --git a/scripts/build.sh b/scripts/build.sh new file mode 100644 index 0000000..7730a78 --- /dev/null +++ b/scripts/build.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "$0")/.." && pwd)" +BUILD_DIR="$ROOT_DIR/build" +mkdir -p "$BUILD_DIR" + +CGO_ENABLED=0 GOOS=linux GOARCH=arm64 \ + go build -o "$BUILD_DIR/edge-agent" "$ROOT_DIR/cmd/edge-agent" + +echo "built: $BUILD_DIR/edge-agent" diff --git a/scripts/install.sh b/scripts/install.sh new file mode 100644 index 0000000..21944eb --- /dev/null +++ b/scripts/install.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +set -euo pipefail + +PREFIX=/opt/tianyan-edge +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" +PYTHON_BIN="${PYTHON_BIN:-python3}" + +sudo mkdir -p "$PREFIX"/{bin,python,config,systemd,staging,model} +sudo install -m 755 "$ROOT_DIR/build/edge-agent" "$PREFIX/bin/edge-agent" +sudo install -m 644 "$ROOT_DIR/python/infer_server.py" "$PREFIX/python/infer_server.py" +sudo install -m 644 "$ROOT_DIR/python/requirements.txt" "$PREFIX/python/requirements.txt" + +if [ ! -f "$PREFIX/config/edge.yaml" ]; then + sudo install -m 644 "$ROOT_DIR/config/edge.yaml.template" "$PREFIX/config/edge.yaml" +fi + +sudo install -m 644 "$ROOT_DIR/systemd/edge-agent.service" /etc/systemd/system/edge-agent.service +sudo install -m 644 "$ROOT_DIR/systemd/edge-infer.service" /etc/systemd/system/edge-infer.service + +sudo "$PYTHON_BIN" -m pip install -r "$PREFIX/python/requirements.txt" + +sudo systemctl daemon-reload +sudo systemctl enable edge-infer.service edge-agent.service +sudo systemctl restart edge-infer.service edge-agent.service + +echo "installed to $PREFIX" \ No newline at end of file diff --git a/scripts/package-usb.sh b/scripts/package-usb.sh new file mode 100644 index 0000000..cce9148 --- /dev/null +++ b/scripts/package-usb.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "$0")/.." && pwd)" +DIST_DIR="$ROOT_DIR/dist" +PKG_NAME="tianyan-edge-$(date +%Y%m%d-%H%M%S)" +PKG_DIR="$DIST_DIR/$PKG_NAME" + +mkdir -p "$PKG_DIR" +cp -r "$ROOT_DIR/build" "$PKG_DIR/" +cp -r "$ROOT_DIR/python" "$PKG_DIR/" +cp -r "$ROOT_DIR/systemd" "$PKG_DIR/" +cp -r "$ROOT_DIR/config" "$PKG_DIR/" +cp -r "$ROOT_DIR/scripts/install.sh" "$ROOT_DIR/scripts/uninstall.sh" "$PKG_DIR/" + +cd "$DIST_DIR" +tar -czf "$PKG_NAME.tar.gz" "$PKG_NAME" + +echo "usb package ready: $DIST_DIR/$PKG_NAME.tar.gz" diff --git a/scripts/uninstall.sh b/scripts/uninstall.sh new file mode 100644 index 0000000..f7fcf3d --- /dev/null +++ b/scripts/uninstall.sh @@ -0,0 +1,10 @@ +#!/usr/bin/env bash +set -euo pipefail + +sudo systemctl stop edge-agent.service edge-infer.service || true +sudo systemctl disable edge-agent.service edge-infer.service || true +sudo rm -f /etc/systemd/system/edge-agent.service /etc/systemd/system/edge-infer.service +sudo systemctl daemon-reload +sudo rm -rf /opt/tianyan-edge + +echo "uninstalled /opt/tianyan-edge" diff --git a/systemd/edge-agent.service b/systemd/edge-agent.service new file mode 100644 index 0000000..4abfb1c --- /dev/null +++ b/systemd/edge-agent.service @@ -0,0 +1,13 @@ +[Unit] +Description=Tianyan Edge Agent +After=network-online.target edge-infer.service +Wants=network-online.target + +[Service] +Type=simple +ExecStart=/opt/tianyan-edge/bin/edge-agent -config /opt/tianyan-edge/config/edge.yaml +Restart=always +RestartSec=3 + +[Install] +WantedBy=multi-user.target diff --git a/systemd/edge-infer.service b/systemd/edge-infer.service new file mode 100644 index 0000000..4ea295f --- /dev/null +++ b/systemd/edge-infer.service @@ -0,0 +1,20 @@ +[Unit] +Description=Tianyan Edge Inference Server (Ascend NPU) +After=network.target + +[Service] +Type=simple +Environment=INFER_SOCKET=/tmp/edge-infer.sock +Environment=OM_MODEL=/opt/tianyan-edge/model/model.om +Environment=CONF_THRESHOLD=0.5 +Environment=IOU_THRESHOLD=0.45 +Environment=DEVICE_ID=0 +Environment=NAMES_FILE=/opt/tianyan-edge/model/names.txt +# raw: YOLOv8 output [1,4+nc,anchors] nms_free: YOLOv10 output [1,topk,6] +Environment=OUTPUT_FORMAT=raw +ExecStart=/usr/bin/python3 /opt/tianyan-edge/python/infer_server.py +Restart=always +RestartSec=3 + +[Install] +WantedBy=multi-user.target