diff --git a/README.md b/README.md index 3cdf36d..39270dc 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,7 @@ See the [examples](https://github.com/TannerKvarfordt/hfapigo/tree/main/examples - [Fill Mask](https://github.com/TannerKvarfordt/hfapigo/blob/main/examples/fill_mask/main.go) - [Speech Recognition](https://github.com/TannerKvarfordt/hfapigo/blob/main/examples/speech_recognition/main.go) - [Audio Classification](https://github.com/TannerKvarfordt/hfapigo/blob/main/examples/audio_classification/main.go) +- [Object Detection](https://github.com/TannerKvarfordt/hfapigo/blob/main/examples/object_detection) # Resources - [Hugging Face](https://huggingface.co/) diff --git a/examples/object_detection/main.go b/examples/object_detection/main.go new file mode 100644 index 0000000..1f116b6 --- /dev/null +++ b/examples/object_detection/main.go @@ -0,0 +1,141 @@ +package main + +import ( + "fmt" + "image/color" + "image/draw" + "image/png" + "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" + outputImg = "./test-image-output.png" + + detectRetries = 10 + detectBackoff = time.Second * 2 +) + +func main() { + fmt.Printf("Opening image: %s\n", inputImg) + img, err := OpenImg(inputImg) + if err != nil { + fmt.Println("Problem opening image:", err) + return + } + + objects, err := SendRequest(inputImg) + if err != nil { + fmt.Println("Problem during object detection request:", err) + return + } + + // Draw the output file + for _, obj := range objects { + col := color.RGBA{uint8(rand.Intn(266)), uint8(rand.Intn(266)), uint8(rand.Intn(266)), 255} + Rect(obj.Box.XMin, obj.Box.YMin, obj.Box.XMax, obj.Box.YMax, img, col) + } + + outf, err := os.Create(outputImg) + if err != nil { + fmt.Println("Problem creating output file:", err) + return + } + defer outf.Close() + + err = png.Encode(outf, img) + if err != nil { + fmt.Println("Problem encoding output file:", err) + return + } + fmt.Println("Output image written to", outf.Name()) +} + +func OpenImg(imgFile string) (draw.Image, error) { + f, err := os.Open(inputImg) + if err != nil { + return nil, err + } + defer f.Close() + img, err := png.Decode(f) + dimg, ok := img.(draw.Image) + if !ok { + return nil, fmt.Errorf("%T is not a drawable image type", img) + } + return dimg, err +} + +func SendRequest(imgFile string) ([]*hfapigo.ObjectDetectionResponse, error) { + fmt.Printf("Sending object detection request for image (%s)", imgFile) + + type ChanRv struct { + resps []*hfapigo.ObjectDetectionResponse + err error + } + ch := make(chan ChanRv) + + go func() { + objects, err := DetectObjects(inputImg, detectRetries, detectBackoff) + ch <- ChanRv{objects, err} + }() + + for { + select { + case chrv := <-ch: + fmt.Println() + return chrv.resps, chrv.err + default: + fmt.Print(".") + time.Sleep(time.Millisecond * 400) + } + } +} + +func DetectObjects(imgFile string, retries int, backoff time.Duration) ([]*hfapigo.ObjectDetectionResponse, error) { + objects := []*hfapigo.ObjectDetectionResponse{} + var err error + for i := 0; i < retries; i++ { + objects, err = hfapigo.SendObjectDetectionRequest(hfapigo.RecommendedObjectDetectionModel, inputImg) + if err == nil { + break + } + time.Sleep(backoff) + } + return objects, err +} + +// HLine draws a horizontal line +func HLine(x1, y, x2 int, img draw.Image, col color.Color) { + for ; x1 <= x2; x1++ { + img.Set(x1, y, col) + } +} + +// VLine draws a veritcal line +func VLine(x, y1, y2 int, img draw.Image, col color.Color) { + for ; y1 <= y2; y1++ { + img.Set(x, y1, col) + } +} + +// Rect draws a rectangle utilizing HLine() and VLine() +func Rect(x1, y1, x2, y2 int, img draw.Image, col color.Color) { + HLine(x1, y1, x2, img, col) + HLine(x1, y2, x2, img, col) + VLine(x1, y1, y2, img, col) + VLine(x2, y1, y2, img, col) +} diff --git a/examples/object_detection/test-image-output.png b/examples/object_detection/test-image-output.png new file mode 100644 index 0000000..b208088 Binary files /dev/null and b/examples/object_detection/test-image-output.png differ diff --git a/examples/object_detection/test-image.png b/examples/object_detection/test-image.png new file mode 100644 index 0000000..1600265 Binary files /dev/null and b/examples/object_detection/test-image.png differ diff --git a/object_detection_test.go b/object_detection_test.go index 1a180dd..ed1e1b8 100644 --- a/object_detection_test.go +++ b/object_detection_test.go @@ -13,7 +13,7 @@ func TestObjectDetectionRequest(t *testing.T) { resps := []*hfapigo.ObjectDetectionResponse{} var err error for i := 0; i < retries; i++ { - resps, err = hfapigo.SendObjectDetectionRequest(hfapigo.RecommendedObjectDetectionModel, TestFilesDir+"/test-image.jpeg") + resps, err = hfapigo.SendObjectDetectionRequest(hfapigo.RecommendedObjectDetectionModel, TestFilesDir+"/test-image.png") if err == nil { break } else { diff --git a/test_files/test-image.jpeg b/test_files/test-image.jpeg deleted file mode 100644 index 665a81c..0000000 Binary files a/test_files/test-image.jpeg and /dev/null differ diff --git a/test_files/test-image.png b/test_files/test-image.png new file mode 100644 index 0000000..1600265 Binary files /dev/null and b/test_files/test-image.png differ