tensorflow3.go

tensorflow3.go

const (
	graphFile = "/model/imagenet_comp_graph_label_strings.txt"
	labelsFile = "/model/imagenet_comp_graph_label_strings.txt"
)

graph, labels, err := loadModel()
if err != nil {
	log.Fatalf("unable to load model: %v", err)
}

func loadModel() (*tf.Graph, []string, error) {
	// Load inception model
	model, err := ioutil.ReadFile(graphFile)
	if err != nil {
		return nil, nil, err
	}
	graph := tf.NewGraph()
	if err := graph.Import(model, ""); err != nil {
		return nil, nil, err
	}

	// Load labels
	labelsFile, err := os.Open()
	if err != nil {
		return nil, nil, err
	}
	defer labelsFile.Close()
	scanner := bufio.NewScanner(labelsFile)
	var labels []string
	for scanner.Scan() {
		labels = append(labels, scanner.Text())
	}

	return graph, labels, scanner.Err()
}