89 lines
1.6 KiB
Go
89 lines
1.6 KiB
Go
package logic
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"github.com/go-resty/resty/v2"
|
|
"io"
|
|
"time"
|
|
|
|
"gitlink.org.cn/JointCloud/pcm-octopus/internal/svc"
|
|
"gitlink.org.cn/JointCloud/pcm-octopus/octopus"
|
|
|
|
"github.com/zeromicro/go-zero/core/logx"
|
|
)
|
|
|
|
type GetInferResultLogic struct {
|
|
ctx context.Context
|
|
svcCtx *svc.ServiceContext
|
|
logx.Logger
|
|
}
|
|
|
|
func NewGetInferResultLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GetInferResultLogic {
|
|
return &GetInferResultLogic{
|
|
ctx: ctx,
|
|
svcCtx: svcCtx,
|
|
Logger: logx.WithContext(ctx),
|
|
}
|
|
}
|
|
|
|
func (l *GetInferResultLogic) GetInferResult(stream octopus.Octopus_GetInferResultServer) error {
|
|
resp := &octopus.InferResultResp{}
|
|
|
|
imageData := bytes.Buffer{}
|
|
imageSize := 0
|
|
var url string
|
|
var filename string
|
|
|
|
for {
|
|
req, err := stream.Recv()
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
|
|
url = req.InferUrl
|
|
filename = req.FileName
|
|
chunk := req.FileBytes
|
|
size := len(chunk)
|
|
|
|
imageSize += size
|
|
_, err = imageData.Write(chunk)
|
|
if err != nil {
|
|
return errors.New("cannot write chunk data")
|
|
}
|
|
}
|
|
|
|
result, err := callInferUrl(url, filename, imageData.Bytes())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
resp.Result = result
|
|
|
|
err = stream.SendAndClose(resp)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func callInferUrl(url string, fileName string, fileBytes []byte) (string, error) {
|
|
var res Res
|
|
client := resty.New().SetTimeout(time.Duration(20) * time.Second)
|
|
req := client.R()
|
|
_, err := req.
|
|
SetFileReader("file", fileName, bytes.NewReader(fileBytes)).
|
|
SetResult(&res).
|
|
Post(url)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return res.Result, nil
|
|
}
|
|
|
|
type Res struct {
|
|
Result string `json:"result"`
|
|
}
|