281 lines
7.9 KiB
Go
281 lines
7.9 KiB
Go
package embedding
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/binary"
|
|
"encoding/json"
|
|
"fmt"
|
|
"hash/fnv"
|
|
"math"
|
|
"net/http"
|
|
"os"
|
|
"regexp"
|
|
"strings"
|
|
|
|
"materialanalyzer/internal/model"
|
|
)
|
|
|
|
const LocalDimensions = 128
|
|
|
|
type Request struct {
|
|
ImagePath string
|
|
Product model.Product
|
|
Histogram model.ColorHistogram
|
|
Visual model.VisualFeatures
|
|
Texture model.TextureFeatures
|
|
Semantic model.SemanticFeatures
|
|
}
|
|
|
|
type Result struct {
|
|
Vectors [][]float32
|
|
Manifest model.EmbeddingManifest
|
|
}
|
|
|
|
type Provider interface {
|
|
Name() string
|
|
Embed(ctx context.Context, req Request) (Result, error)
|
|
}
|
|
|
|
type LocalProvider struct{}
|
|
|
|
func (LocalProvider) Name() string { return "local" }
|
|
|
|
func (LocalProvider) Embed(ctx context.Context, req Request) (Result, error) {
|
|
select {
|
|
case <-ctx.Done():
|
|
return Result{}, ctx.Err()
|
|
default:
|
|
}
|
|
textureVector := BuildTextureVector(req.Histogram, req.Visual, req.Texture)
|
|
semanticVector := BuildSemanticVector(req.Product, req.Semantic)
|
|
return Result{
|
|
Vectors: [][]float32{textureVector, semanticVector},
|
|
Manifest: model.EmbeddingManifest{
|
|
File: "embedding.bin",
|
|
Format: "float32_le_concatenated",
|
|
Vectors: []model.EmbeddingVector{
|
|
{Name: "dino_v2_texture_local", Model: "local_texture_fingerprint", Dimensions: LocalDimensions, Generator: "local_color_texture_fingerprint_v2", Provider: "local", Normalize: true},
|
|
{Name: "clip_semantic_local", Model: "local_semantic_hash", Dimensions: LocalDimensions, Generator: "local_metadata_semantic_hash_v2", Provider: "local", Normalize: true},
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type HTTPProvider struct {
|
|
BaseURL string
|
|
Client *http.Client
|
|
}
|
|
|
|
func (p HTTPProvider) Name() string { return "http" }
|
|
|
|
func (p HTTPProvider) Embed(ctx context.Context, req Request) (Result, error) {
|
|
if strings.TrimSpace(p.BaseURL) == "" {
|
|
return Result{}, fmt.Errorf("embedding HTTP provider URL is empty")
|
|
}
|
|
client := p.Client
|
|
if client == nil {
|
|
client = http.DefaultClient
|
|
}
|
|
imgData, err := os.ReadFile(req.ImagePath)
|
|
if err != nil {
|
|
return Result{}, fmt.Errorf("read embedding image: %w", err)
|
|
}
|
|
body := httpEmbeddingRequest{
|
|
ImageBase64: base64.StdEncoding.EncodeToString(imgData),
|
|
ImagePath: req.ImagePath,
|
|
Product: req.Product,
|
|
Semantic: req.Semantic,
|
|
}
|
|
payload, err := json.Marshal(body)
|
|
if err != nil {
|
|
return Result{}, err
|
|
}
|
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, embeddingEndpoint(p.BaseURL), bytes.NewReader(payload))
|
|
if err != nil {
|
|
return Result{}, err
|
|
}
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
httpReq.Header.Set("Accept", "application/json")
|
|
|
|
resp, err := client.Do(httpReq)
|
|
if err != nil {
|
|
return Result{}, fmt.Errorf("embedding HTTP request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return Result{}, fmt.Errorf("embedding HTTP status %d", resp.StatusCode)
|
|
}
|
|
var decoded httpEmbeddingResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&decoded); err != nil {
|
|
return Result{}, fmt.Errorf("decode embedding response: %w", err)
|
|
}
|
|
return decoded.toResult()
|
|
}
|
|
|
|
type httpEmbeddingRequest struct {
|
|
ImageBase64 string `json:"image_base64"`
|
|
ImagePath string `json:"image_path"`
|
|
Product model.Product `json:"product"`
|
|
Semantic model.SemanticFeatures `json:"semantic"`
|
|
}
|
|
|
|
type httpEmbeddingResponse struct {
|
|
Vectors []httpEmbeddingVector `json:"vectors"`
|
|
}
|
|
|
|
type httpEmbeddingVector struct {
|
|
Name string `json:"name"`
|
|
Model string `json:"model"`
|
|
Dimension int `json:"dimension"`
|
|
Normalize bool `json:"normalize"`
|
|
Provider string `json:"provider"`
|
|
Generator string `json:"generator"`
|
|
Values []float32 `json:"values"`
|
|
}
|
|
|
|
func (r httpEmbeddingResponse) toResult() (Result, error) {
|
|
if len(r.Vectors) == 0 {
|
|
return Result{}, fmt.Errorf("embedding response has no vectors")
|
|
}
|
|
result := Result{
|
|
Manifest: model.EmbeddingManifest{File: "embedding.bin", Format: "float32_le_concatenated"},
|
|
}
|
|
for _, v := range r.Vectors {
|
|
if v.Dimension <= 0 {
|
|
v.Dimension = len(v.Values)
|
|
}
|
|
if len(v.Values) != v.Dimension {
|
|
return Result{}, fmt.Errorf("embedding %s length = %d, dimension = %d", v.Name, len(v.Values), v.Dimension)
|
|
}
|
|
if v.Name == "" {
|
|
return Result{}, fmt.Errorf("embedding vector name is required")
|
|
}
|
|
result.Vectors = append(result.Vectors, v.Values)
|
|
result.Manifest.Vectors = append(result.Manifest.Vectors, model.EmbeddingVector{
|
|
Name: v.Name, Model: v.Model, Dimensions: v.Dimension,
|
|
Generator: v.Generator, Provider: v.Provider, Normalize: v.Normalize,
|
|
})
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func embeddingEndpoint(base string) string {
|
|
base = strings.TrimRight(base, "/")
|
|
if strings.HasSuffix(base, "/embeddings") {
|
|
return base
|
|
}
|
|
return base + "/embeddings"
|
|
}
|
|
|
|
func BuildTextureVector(hist model.ColorHistogram, visual model.VisualFeatures, texture model.TextureFeatures) []float32 {
|
|
vec := make([]float32, LocalDimensions)
|
|
if len(hist.Values) > 0 {
|
|
target := 96
|
|
for i := 0; i < target; i++ {
|
|
start := i * len(hist.Values) / target
|
|
end := (i + 1) * len(hist.Values) / target
|
|
var sum float64
|
|
for _, v := range hist.Values[start:end] {
|
|
sum += v
|
|
}
|
|
vec[i] = float32(sum)
|
|
}
|
|
}
|
|
stats := []float64{
|
|
visual.Brightness, visual.Contrast, visual.Saturation,
|
|
float64(visual.DominantRGB.R) / 255, float64(visual.DominantRGB.G) / 255, float64(visual.DominantRGB.B) / 255,
|
|
float64(visual.SecondaryRGB.R) / 255, float64(visual.SecondaryRGB.G) / 255, float64(visual.SecondaryRGB.B) / 255,
|
|
texture.Entropy / 8, texture.TextureFrequency, texture.EdgeDensity, texture.OrientationVariance,
|
|
texture.PrimaryOrientationDeg / 180, texture.LBPUniformity,
|
|
texture.GLCM.Contrast / 49, texture.GLCM.Homogeneity, texture.GLCM.Energy, (texture.GLCM.Correlation + 1) / 2,
|
|
}
|
|
for i, v := range stats {
|
|
idx := 96 + i
|
|
if idx >= len(vec) {
|
|
break
|
|
}
|
|
vec[idx] = float32(v)
|
|
}
|
|
normalize(vec)
|
|
return vec
|
|
}
|
|
|
|
func BuildSemanticVector(p model.Product, semantic model.SemanticFeatures) []float32 {
|
|
vec := make([]float32, LocalDimensions)
|
|
text := strings.Join([]string{
|
|
p.Brand, p.GroupName, p.SeriesName, p.StyleName, p.Category, p.Material, p.ColorTone,
|
|
p.Finish, p.Description, semantic.Description, semantic.MaterialType, semantic.VisualStyle,
|
|
semantic.GrainType, semantic.Grain, semantic.Variation, semantic.ColorFamily, semantic.GlossLevel,
|
|
strings.Join(semantic.Tags, " "),
|
|
}, " ")
|
|
for _, token := range tokenize(text) {
|
|
h := fnv.New32a()
|
|
_, _ = h.Write([]byte(token))
|
|
sum := h.Sum32()
|
|
idx := int(sum % LocalDimensions)
|
|
sign := float32(1)
|
|
if sum&0x80 != 0 {
|
|
sign = -1
|
|
}
|
|
vec[idx] += sign
|
|
}
|
|
normalize(vec)
|
|
return vec
|
|
}
|
|
|
|
func WriteBinary(path string, result Result) (model.EmbeddingManifest, error) {
|
|
file, err := os.Create(path)
|
|
if err != nil {
|
|
return model.EmbeddingManifest{}, err
|
|
}
|
|
defer file.Close()
|
|
|
|
manifest := result.Manifest
|
|
if manifest.File == "" {
|
|
manifest.File = "embedding.bin"
|
|
}
|
|
if manifest.Format == "" {
|
|
manifest.Format = "float32_le_concatenated"
|
|
}
|
|
var offset int64
|
|
for i, vec := range result.Vectors {
|
|
if i >= len(manifest.Vectors) {
|
|
manifest.Vectors = append(manifest.Vectors, model.EmbeddingVector{Name: fmt.Sprintf("vector_%d", i), Dimensions: len(vec)})
|
|
}
|
|
manifest.Vectors[i].OffsetBytes = offset
|
|
manifest.Vectors[i].ByteLength = int64(len(vec) * 4)
|
|
if manifest.Vectors[i].Dimensions == 0 {
|
|
manifest.Vectors[i].Dimensions = len(vec)
|
|
}
|
|
for _, v := range vec {
|
|
if err := binary.Write(file, binary.LittleEndian, v); err != nil {
|
|
return model.EmbeddingManifest{}, err
|
|
}
|
|
}
|
|
offset += int64(len(vec) * 4)
|
|
}
|
|
return manifest, nil
|
|
}
|
|
|
|
var tokenRE = regexp.MustCompile(`[a-z0-9]+`)
|
|
|
|
func tokenize(text string) []string {
|
|
return tokenRE.FindAllString(strings.ToLower(text), -1)
|
|
}
|
|
|
|
func normalize(vec []float32) {
|
|
var sum float64
|
|
for _, v := range vec {
|
|
sum += float64(v * v)
|
|
}
|
|
if sum == 0 {
|
|
return
|
|
}
|
|
scale := float32(1 / math.Sqrt(sum))
|
|
for i := range vec {
|
|
vec[i] *= scale
|
|
}
|
|
}
|