Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
84 lines (69 sloc) 2.73 KB
// Tensorflow Serving Go client for the inception model
// First of all compile the proto files:
// git clone --recursive https://github.com/tensorflow/serving.git
// protoc -I=serving -I serving/tensorflow --go_out=plugins=grpc:$GOPATH/src serving/tensorflow_serving/apis/*.proto
// protoc -I=serving/tensorflow --go_out=plugins=grpc:$GOPATH/src serving/tensorflow/tensorflow/core/framework/*.proto
// protoc -I=serving/tensorflow --go_out=plugins=grpc:$GOPATH/src serving/tensorflow/tensorflow/core/protobuf/{saver,meta_graph}.proto
// protoc -I=serving/tensorflow --go_out=plugins=grpc:$GOPATH/src serving/tensorflow/tensorflow/core/example/*.proto
package main
import (
"context"
"log"
pb "github.com/jnummelin/go-inception-client/tensorflow_serving/apis"
tf_core_framework "github.com/jnummelin/go-inception-client/tensorflow/core/framework"
google_protobuf "github.com/golang/protobuf/ptypes/wrappers"
tf "github.com/tensorflow/tensorflow/tensorflow/go"
"google.golang.org/grpc"
)
type InceptionClient struct {
ServingAddress string
}
// Predict
func (i InceptionClient) Predict(imageBytes []byte) (*pb.PredictResponse, error) {
tensor, err := tf.NewTensor(string(imageBytes))
if err != nil {
log.Fatalln("Cannot read image file")
return nil, err
}
tensorString, ok := tensor.Value().(string)
if !ok {
log.Fatalln("Cannot type assert tensor value to string")
return nil, err
}
request := &pb.PredictRequest{
ModelSpec: &pb.ModelSpec{
Name: "inception",
SignatureName: "predict_images",
Version: &google_protobuf.Int64Value{
Value: int64(1),
},
},
Inputs: map[string]*tf_core_framework.TensorProto{
"images": &tf_core_framework.TensorProto{
Dtype: tf_core_framework.DataType_DT_STRING,
TensorShape: &tf_core_framework.TensorShapeProto{
Dim: []*tf_core_framework.TensorShapeProto_Dim{
&tf_core_framework.TensorShapeProto_Dim{
Size: int64(1),
},
},
},
StringVal: [][]byte{[]byte(tensorString)},
},
},
}
log.Printf("serving address: %s", i.ServingAddress)
conn, err := grpc.Dial(i.ServingAddress, grpc.WithInsecure())
if err != nil {
log.Fatalf("Cannot connect to the grpc server: %v\n", err)
return nil, err
}
defer conn.Close()
client := pb.NewPredictionServiceClient(conn)
resp, err := client.Predict(context.Background(), request)
if err != nil {
log.Fatalln(err)
return nil, err
}
return resp, nil
}