FloorMaterialAnalyzer/internal/embedding/embedding.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
	}
}