zt-closet-backend/pkg/classification/classification.go

// Package classification 封装衣物分类服务调用，将衣物图片上传到 Fashion-CLIP 分类服务，
// 返回结构化的衣物属性分类结果。
package classification

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"time"

	"git.lzcwy.cn/zhengchiliu/zt-closet-backend/global"
)

const (
	classifyPath = "/classify"
	healthzPath  = "/healthz"
)

// DetailItem 分类明细项。
type DetailItem struct {
	Label      string  `json:"label"`
	Confidence float64 `json:"confidence"`
}

// Details 各维度 Top-5 明细。
type Details struct {
	CategoryMain []DetailItem `json:"category_main"`
	CategorySub  []DetailItem `json:"category_sub"`
	Style        []DetailItem `json:"style"`
	ColorMain    []DetailItem `json:"color_main"`
	Thickness    []DetailItem `json:"thickness"`
	Version      []DetailItem `json:"version"`
	Material     []DetailItem `json:"material"`
	Pattern      []DetailItem `json:"pattern"`
	SleeveType   []DetailItem `json:"sleeve_type"`
	ClothesType  []DetailItem `json:"clothes_type"`
	NecklineType []DetailItem `json:"neckline_type"`
	Formality    []DetailItem `json:"formality"`
	Occasions    []DetailItem `json:"occasions"`
	VisualImpact []DetailItem `json:"visual_impact"`
}

// ClassifyResult 衣物分类结果。
type ClassifyResult struct {
	CategoryMain           string  `json:"category_main"`
	CategoryMainConfidence float64 `json:"category_main_confidence"`
	CategorySub            string  `json:"category_sub"`
	CategorySubConfidence  float64 `json:"category_sub_confidence"`
	Style                  string  `json:"style"`
	StyleConfidence        float64 `json:"style_confidence"`
	ColorMain              string  `json:"color_main"`
	ColorMainConfidence    float64 `json:"color_main_confidence"`
	ColorSub               string  `json:"color_sub"`
	ColorSubConfidence     float64 `json:"color_sub_confidence"`
	Thickness              string  `json:"thickness"`
	ThicknessConfidence    float64 `json:"thickness_confidence"`
	Version                string  `json:"version"`
	VersionConfidence      float64 `json:"version_confidence"`
	Material               string  `json:"material"`
	MaterialConfidence     float64 `json:"material_confidence"`
	Pattern                string  `json:"pattern"`
	PatternConfidence      float64 `json:"pattern_confidence"`
	SleeveType             string  `json:"sleeve_type"`
	SleeveTypeConfidence   float64 `json:"sleeve_type_confidence"`
	ClothesType            string  `json:"clothes_type"`
	ClothesTypeConfidence  float64 `json:"clothes_type_confidence"`
	NecklineType           string  `json:"neckline_type"`
	NecklineTypeConfidence float64 `json:"neckline_type_confidence"`
	Formality              string  `json:"formality"`
	FormalityLabel         string  `json:"formality_label"`
	FormalityConfidence    float64 `json:"formality_confidence"`
	Occasions              string  `json:"occasions"`
	VisualImpact           string  `json:"visual_impact"`
	VisualImpactLabel      string  `json:"visual_impact_label"`
	VisualImpactConfidence float64 `json:"visual_impact_confidence"`
	Details                Details `json:"details"`
}

// Client 分类服务客户端。
type Client struct {
	baseURL    string
	httpClient *http.Client
}

// NewClient 从全局配置创建客户端。
func NewClient() *Client {
	return NewClientWithAddr(global.ClassificationModelSetting.Host, global.ClassificationModelSetting.Port)
}

// NewClientWithAddr 使用指定 Host/Port 创建客户端，便于测试与解耦。
func NewClientWithAddr(host string, port int) *Client {
	return &Client{
		baseURL: fmt.Sprintf("http://%s:%d", host, port),
		httpClient: &http.Client{
			Timeout: 120 * time.Second,
		},
	}
}

// Classify 调用分类服务，上传图片并返回结构化的衣物分类结果。
func (c *Client) Classify(ctx context.Context, imageData []byte) (*ClassifyResult, error) {
	var body bytes.Buffer
	writer := multipart.NewWriter(&body)

	part, err := writer.CreateFormFile("file", "image.png")
	if err != nil {
		return nil, fmt.Errorf("classify create form file: %w", err)
	}
	if _, err := part.Write(imageData); err != nil {
		return nil, fmt.Errorf("classify write image: %w", err)
	}
	if err := writer.Close(); err != nil {
		return nil, fmt.Errorf("classify close multipart: %w", err)
	}

	request, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+classifyPath, &body)
	if err != nil {
		return nil, fmt.Errorf("classify create request: %w", err)
	}
	request.Header.Set("Content-Type", writer.FormDataContentType())

	response, err := c.httpClient.Do(request)
	if err != nil {
		return nil, fmt.Errorf("classify do request: %w", err)
	}
	defer response.Body.Close()

	data, err := io.ReadAll(response.Body)
	if err != nil {
		return nil, fmt.Errorf("classify read response: %w", err)
	}

	if response.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("classify api error: status=%d, body=%s", response.StatusCode, summarize(data))
	}

	var result ClassifyResult
	if err := json.Unmarshal(data, &result); err != nil {
		return nil, fmt.Errorf("classify decode response: %w", err)
	}
	return &result, nil
}

// HealthCheck 检查分类服务是否就绪。
func (c *Client) HealthCheck(ctx context.Context) error {
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+healthzPath, nil)
	if err != nil {
		return fmt.Errorf("classify create healthz request: %w", err)
	}

	response, err := c.httpClient.Do(request)
	if err != nil {
		return fmt.Errorf("classify healthz request: %w", err)
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("classify healthz status: %d", response.StatusCode)
	}
	return nil
}

// summarize 截断响应体摘要，避免错误信息打印大体积 JSON。
func summarize(data []byte) string {
	const maxLen = 512
	if len(data) > maxLen {
		return string(data[:maxLen]) + "..."
	}
	return string(data)
}