Image to text

.gitignore

@@ -35,4 +35,5 @@ examples/token_classification/token_classification
 examples/translation/translation
 examples/zeroshot/zeroshot
 examples/text_to_image/*
-!examples/text_to_image/*.go
+!examples/text_to_image/*.go
+examples/image_to_text/image_to_text

README.md

@@ -19,6 +19,7 @@ See the [examples](./examples) directory.
 - [Fill Mask](./examples/fill_mask/main.go)
 - [Image Classification](./examples/image_classification/main.go)
 - [Image Segmentation](./examples/image_segmentation/main.go)
+- [Image-To-Text](./examples/image_to_text/main.go)
 - [Object Detection](./examples/object_detection/main.go)
 - [Question Answering](./examples/question_answering/main.go)
 - [Sentence Similarity](./examples/sentence_similarity/main.go)
@@ -39,4 +40,6 @@ See the [examples](./examples) directory.
   - [Datasets](https://huggingface.co/datasets)
   - [Hugging Face Inference API](https://api-inference.huggingface.co/docs/python/html/index.html) (HF API)
   - [HF on GitHub](https://github.com/huggingface)
+  - [Diffuser Docs](https://huggingface.co/docs/diffusers/using-diffusers/pipeline_overview)
   - Official [Python bindings](https://github.com/huggingface/hfapi) for the HF API
+  - Official [JavaScript bindings](https://github.com/huggingface/huggingface.js) for the HF Inference API

examples/image_to_text/main.go

@@ -0,0 +1,55 @@
+package main
+
+import (
+	"fmt"
+	"math/rand"
+	"os"
+	"time"
+
+	"github.com/TannerKvarfordt/hfapigo"
+)
+
+const HuggingFaceTokenEnv = "HUGGING_FACE_TOKEN"
+
+func init() {
+	rand.Seed(time.Now().UnixNano())
+	key := os.Getenv(HuggingFaceTokenEnv)
+	if key != "" {
+		hfapigo.SetAPIKey(key)
+	}
+}
+
+const inputImg = "./test-image.png"
+
+func main() {
+	fmt.Printf("Sending image to text request for image %s", inputImg)
+
+	type ChanRv struct {
+		resps []*hfapigo.ImageToTextResponse
+		err   error
+	}
+	ch := make(chan ChanRv)
+
+	go func() {
+		resps, err := hfapigo.SendImageToTextRequest(hfapigo.RecommendedImageToTextModel, inputImg)
+		ch <- ChanRv{resps: resps, err: err}
+	}()
+
+	for {
+		select {
+		case chrv := <-ch:
+			fmt.Println()
+			if chrv.err != nil {
+				fmt.Println(chrv.err)
+				return
+			}
+			for _, r := range chrv.resps {
+				fmt.Printf("Caption: %s\n", r.GeneratedText)
+			}
+			return
+		default:
+			fmt.Print(".")
+			time.Sleep(time.Millisecond * 400)
+		}
+	}
+}

image_to_text.go

@@ -0,0 +1,25 @@
+package hfapigo
+
+import "encoding/json"
+
+const RecommendedImageToTextModel = "nlpconnect/vit-gpt2-image-captioning"
+
+type ImageToTextResponse struct {
+	// The generated caption
+	GeneratedText string `json:"generated_text"`
+}
+
+func SendImageToTextRequest(model, imageFile string) ([]*ImageToTextResponse, error) {
+	respBody, err := MakeHFAPIRequestWithMedia(model, imageFile)
+	if err != nil {
+		return nil, err
+	}
+
+	resps := []*ImageToTextResponse{}
+	err = json.Unmarshal(respBody, &resps)
+	if err != nil {
+		return nil, err
+	}
+
+	return resps, nil
+}

image_to_text_test.go

@@ -0,0 +1,38 @@
+package hfapigo_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/TannerKvarfordt/hfapigo"
+)
+
+func TestImageToText(t *testing.T) {
+	const retries = 10
+
+	resps := []*hfapigo.ImageToTextResponse{}
+	var err error
+	for i := 0; i < retries; i++ {
+		resps, err = hfapigo.SendImageToTextRequest(hfapigo.RecommendedImageToTextModel, TestFilesDir+"/test-image.png")
+		if err == nil {
+			break
+		} else {
+			time.Sleep(time.Second * 5)
+		}
+	}
+
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(resps) == 0 {
+		t.Fatal("Expected non-empty response")
+	}
+	for _, resp := range resps {
+		if resp == nil {
+			t.Fatal("nil response received")
+		}
+		if resp.GeneratedText == "" {
+			t.Fatal("Expected non-empty caption")
+		}
+	}
+}