package main
|
|
import (
|
"context"
|
"io/ioutil"
|
"os"
|
"strings"
|
"time"
|
"unsafe"
|
|
"reid/common"
|
|
"basic.com/valib/gogpu.git"
|
|
"basic.com/pubsub/protomsg.git"
|
|
"github.com/gogo/protobuf/proto"
|
|
"nanomsg.org/go-mangos"
|
"nanomsg.org/go-mangos/protocol/rep"
|
"nanomsg.org/go-mangos/transport/ipc"
|
"nanomsg.org/go-mangos/transport/tcp"
|
)
|
|
type reid struct {
|
handle *ReID
|
fnLogger func(...interface{})
|
|
gpu int
|
ipc string
|
}
|
|
// Create Reid
|
func Create(config string, typ, id string, gpu int, shm bool, ipc2Rule string, ruleMaxSize int, fn func(...interface{}), reserved map[string]interface{}) interface{} {
|
|
cfg, err := common.ReadConfig(config)
|
if err != nil {
|
fn("Reid SDK Create Error When Read Config: ", err)
|
return nil
|
}
|
|
sModel, sGPU, sIPC :=
|
"reid-model",
|
"gpu-memory",
|
"ipc-url"
|
|
params := []string{sModel, sGPU, sIPC}
|
|
for _, v := range params {
|
if _, ok := cfg.Param[v]; !ok {
|
fn("Reid SDK Create Error Because of Param Not Found: ", v)
|
return nil
|
}
|
}
|
|
gpuM := common.Atoi(cfg.Param[sGPU])
|
|
rGPU := gpu
|
|
if rGPU == -1 {
|
rGPU = gogpu.ValidGPU(gpuM + 512)
|
}
|
|
handle := NewSDK(rGPU, cfg.Param[sModel])
|
if handle == nil {
|
fn("Reid SDK Create Error When New SDK")
|
return nil
|
}
|
|
return &reid{
|
handle: handle,
|
fnLogger: fn,
|
|
gpu: rGPU,
|
ipc: cfg.Param[sIPC],
|
}
|
}
|
|
// Run run
|
func Run(ctx context.Context, i interface{}) {
|
s := i.(*reid)
|
|
var sock mangos.Socket
|
var err error
|
var msg []byte
|
|
for {
|
if sock, err = rep.NewSocket(); err != nil {
|
s.fnLogger("can't get new rep socket: ", err)
|
time.Sleep(5 * time.Millisecond)
|
} else {
|
break
|
}
|
}
|
|
sock.AddTransport(ipc.NewTransport())
|
sock.AddTransport(tcp.NewTransport())
|
|
for {
|
if err = sock.Listen(s.ipc); err != nil {
|
suf := "ipc://"
|
p := strings.Index(s.ipc, suf)
|
if p >= 0 {
|
file := s.ipc[p+len(string(suf)):]
|
os.Remove(file)
|
s.fnLogger("remove:", file)
|
}
|
s.fnLogger("can't listen on rep socket: ", err)
|
time.Sleep(5 * time.Millisecond)
|
} else {
|
break
|
}
|
}
|
|
for {
|
select {
|
case <-ctx.Done():
|
sock.Close()
|
return
|
default:
|
|
if msg, err = sock.Recv(); err != nil {
|
s.fnLogger("REID~~~~~~Recv From HumanTrack error: ", err)
|
continue
|
}
|
|
i := &protomsg.Image{}
|
err := proto.Unmarshal(msg, i)
|
if err != nil {
|
s.fnLogger("REID~~~~~~protobuf decode CameraImage error: ", err)
|
continue
|
}
|
if i.Data == nil {
|
s.fnLogger("REID~~~~~~protomsg.Image data null")
|
continue
|
}
|
s.fnLogger("REID~~~~~~Recv Image:", len(i.Data))
|
|
feat := s.handle.Extract2(unsafe.Pointer(&i.Data[0]), int(i.Width), int(i.Height), 3)
|
if feat == nil {
|
// feat = make([]float32, 1)
|
} else {
|
for k := 0; k < 3; k++ {
|
s.fnLogger("REID~~~~~~extractor---human_feats------%f", feat[k+2000])
|
}
|
s.fnLogger("REID~~~~~~Run Reid Use GPU: ", s.gpu)
|
}
|
buf := float32SliceAsByteSlice(feat)
|
ioutil.WriteFile("./reid-feat-byte.txt", buf, 0644)
|
|
if err = sock.Send(buf); err != nil {
|
s.fnLogger("REID~~~~~~Send HumanTrack error: ", err)
|
}
|
|
}
|
|
}
|
}
|
|
func float32SliceAsByteSlice(src []float32) []byte {
|
if len(src) == 0 {
|
return nil
|
}
|
|
l := len(src) * 4
|
ptr := unsafe.Pointer(&src[0])
|
return (*[1 << 26]byte)((*[1 << 26]byte)(ptr))[:l:l]
|
}
|
