Spaces:
Running
Running
File size: 2,754 Bytes
b110593 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
// _ _
// __ _____ __ ___ ___ __ _| |_ ___
// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
// \ V V / __/ (_| |\ V /| | (_| | || __/
// \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
//
// Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
//
// CONTACT: [email protected]
//
package test
import (
"encoding/base64"
"fmt"
"io"
"os"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/weaviate/weaviate/entities/models"
"github.com/weaviate/weaviate/entities/schema"
"github.com/weaviate/weaviate/test/helper"
graphqlhelper "github.com/weaviate/weaviate/test/helper/graphql"
)
func Test_Img2VecNeural(t *testing.T) {
helper.SetupClient(os.Getenv(weaviateEndpoint))
fashionItemClass := &models.Class{
Class: "FashionItem",
Vectorizer: "img2vec-neural",
ModuleConfig: map[string]interface{}{
"img2vec-neural": map[string]interface{}{
"imageFields": []string{"image"},
},
},
Properties: []*models.Property{
{
Name: "image",
DataType: schema.DataTypeBlob.PropString(),
},
{
Name: "description",
DataType: schema.DataTypeText.PropString(),
Tokenization: models.PropertyTokenizationWhitespace,
},
},
}
helper.CreateClass(t, fashionItemClass)
defer helper.DeleteClass(t, fashionItemClass.Class)
getBase64EncodedTestImage := func() string {
image, err := os.Open("./data/pixel.png")
require.Nil(t, err)
require.NotNil(t, image)
content, err := io.ReadAll(image)
require.Nil(t, err)
return base64.StdEncoding.EncodeToString(content)
}
t.Run("import data", func(t *testing.T) {
base64Image := getBase64EncodedTestImage()
obj := &models.Object{
Class: fashionItemClass.Class,
Properties: map[string]interface{}{
"image": base64Image,
"description": "A single black pixel",
},
}
helper.CreateObject(t, obj)
})
t.Run("perform nearImage query", func(t *testing.T) {
queryTemplate := `
{
Get {
FashionItem(
nearImage: {
image: "%s"
}
){
image
description
_additional{vector}
}
}
}`
query := fmt.Sprintf(queryTemplate, getBase64EncodedTestImage())
result := graphqlhelper.AssertGraphQL(t, helper.RootAuth, query)
fashionItems := result.Get("Get", "FashionItem").AsSlice()
require.True(t, len(fashionItems) > 0)
item, ok := fashionItems[0].(map[string]interface{})
require.True(t, ok)
assert.NotNil(t, item["image"])
assert.NotNil(t, item["description"])
vector := item["_additional"].(map[string]interface{})["vector"]
assert.NotNil(t, vector)
})
}
|