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:
10
.env.template
Normal file
10
.env.template
Normal file
@@ -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
|
||||||
16
Makefile
Normal file
16
Makefile
Normal file
@@ -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
|
||||||
50
cmd/edge-agent/main.go
Normal file
50
cmd/edge-agent/main.go
Normal file
@@ -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")
|
||||||
|
}
|
||||||
18
config/edge.yaml.template
Normal file
18
config/edge.yaml.template
Normal file
@@ -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
|
||||||
7
go.mod
Normal file
7
go.mod
Normal file
@@ -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
|
||||||
45
internal/config/config.go
Normal file
45
internal/config/config.go
Normal file
@@ -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
|
||||||
|
}
|
||||||
59
internal/control/config_agent.go
Normal file
59
internal/control/config_agent.go
Normal file
@@ -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
|
||||||
|
}
|
||||||
|
}
|
||||||
54
internal/control/heartbeat.go
Normal file
54
internal/control/heartbeat.go
Normal file
@@ -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()
|
||||||
|
}
|
||||||
115
internal/control/ota_agent.go
Normal file
115
internal/control/ota_agent.go
Normal file
@@ -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()
|
||||||
|
}
|
||||||
|
}
|
||||||
13
internal/event/types.go
Normal file
13
internal/event/types.go
Normal file
@@ -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"`
|
||||||
|
}
|
||||||
80
internal/event/uploader.go
Normal file
80
internal/event/uploader.go
Normal file
@@ -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
|
||||||
|
}
|
||||||
209
internal/infer/client.go
Normal file
209
internal/infer/client.go
Normal file
@@ -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:])
|
||||||
|
}
|
||||||
129
internal/stream/ingestor.go
Normal file
129
internal/stream/ingestor.go
Normal file
@@ -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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
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
|
||||||
6
requirements.txt
Normal file
6
requirements.txt
Normal file
@@ -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
|
||||||
11
scripts/build.sh
Normal file
11
scripts/build.sh
Normal file
@@ -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"
|
||||||
27
scripts/install.sh
Normal file
27
scripts/install.sh
Normal file
@@ -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"
|
||||||
19
scripts/package-usb.sh
Normal file
19
scripts/package-usb.sh
Normal file
@@ -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"
|
||||||
10
scripts/uninstall.sh
Normal file
10
scripts/uninstall.sh
Normal file
@@ -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"
|
||||||
13
systemd/edge-agent.service
Normal file
13
systemd/edge-agent.service
Normal file
@@ -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
|
||||||
20
systemd/edge-infer.service
Normal file
20
systemd/edge-infer.service
Normal file
@@ -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
|
||||||
Reference in New Issue
Block a user