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:
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
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user