348 lines
11 KiB
Go
348 lines
11 KiB
Go
package queue
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"database/sql"
|
||
"encoding/json"
|
||
"floorvisualizer/internal/logger"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
|
||
"floorvisualizer/internal/handler"
|
||
"floorvisualizer/internal/openrouter"
|
||
)
|
||
|
||
// StartWorkers launches n goroutines that consume jobs from the Redis queue.
|
||
func StartWorkers(ctx context.Context, n int, rdb *redis.Client, client *openrouter.Client, db *sql.DB) {
|
||
for i := 0; i < n; i++ {
|
||
go func(workerID int) {
|
||
logger.Info("[worker %d] started", workerID)
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
logger.Info("[worker %d] shutting down", workerID)
|
||
return
|
||
default:
|
||
}
|
||
jobID, err := DequeueJob(ctx, rdb)
|
||
if err != nil {
|
||
if ctx.Err() != nil {
|
||
return
|
||
}
|
||
logger.Error("[worker %d] dequeue error: %v", workerID, err)
|
||
time.Sleep(time.Second)
|
||
continue
|
||
}
|
||
logger.Info("[worker %d] picked job %s", workerID, jobID)
|
||
processJob(ctx, rdb, client, db, jobID)
|
||
}
|
||
}(i + 1)
|
||
}
|
||
}
|
||
|
||
func processJob(ctx context.Context, rdb *redis.Client, client *openrouter.Client, db *sql.DB, jobID string) {
|
||
totalStart := time.Now()
|
||
SetJobStatus(ctx, rdb, jobID, "processing")
|
||
logger.Info("[%s] TIMER job start", jobID)
|
||
|
||
// Load payload
|
||
payload, err := GetJobPayload(ctx, rdb, jobID)
|
||
if err != nil {
|
||
SetJobError(ctx, rdb, jobID, "Failed to load job payload")
|
||
return
|
||
}
|
||
|
||
inputPath := payload["input_path"]
|
||
floorID := payload["floor_id"]
|
||
patternCode := payload["pattern_code"]
|
||
roomCode := payload["room_code"]
|
||
|
||
// Look up floor and pattern from DB
|
||
floor := handler.FindFloorBySKUPublic(db, floorID)
|
||
pattern := handler.FindPatternByCodePublic(patternCode)
|
||
room := handler.FindRoomByCodePublic(roomCode)
|
||
groutLines := payloadBool(payload, "grout_lines")
|
||
whiteBaseboard := payloadBool(payload, "white_baseboard")
|
||
materialSource := payload["material_source"] // "image", "json", or "both"
|
||
if materialSource == "" {
|
||
materialSource = "image"
|
||
}
|
||
|
||
if floor == nil || pattern == nil {
|
||
SetJobError(ctx, rdb, jobID, "Invalid floor or pattern")
|
||
return
|
||
}
|
||
logger.Info("[%s] TIMER options floor_id=%s pattern_code=%s room_code=%s elapsed=%s",
|
||
jobID, floorID, patternCode, roomCode, time.Since(totalStart).Round(time.Millisecond))
|
||
|
||
// Resize input to max 2048px to speed up AI processing
|
||
workPath := filepath.Join("uploads", fmt.Sprintf("input_%s_resized.png", jobID))
|
||
if err := openrouter.ResizeImageToFit(inputPath, workPath); err != nil {
|
||
logger.Warn("[%s] Failed to resize input, using original: %v", jobID, err)
|
||
workPath = inputPath
|
||
}
|
||
|
||
outputPath := filepath.Join("outputs", fmt.Sprintf("result_%s.png", jobID))
|
||
|
||
jobCtx, cancel := context.WithTimeout(ctx, 600*time.Second)
|
||
defer cancel()
|
||
|
||
// Step 0: content check
|
||
SetJobProgress(jobCtx, rdb, jobID, "check", "Checking content...")
|
||
stepStart := time.Now()
|
||
logger.Info("[%s] TIMER step=check start elapsed=%s", jobID, time.Since(totalStart).Round(time.Millisecond))
|
||
passed, score := validateContent(client, jobCtx, workPath)
|
||
logStepDone(jobCtx, rdb, jobID, "check", stepStart, totalStart, "score=%d passed=%t", score, passed)
|
||
if score == 0 {
|
||
logger.Warn("[%s] Content check failed (score=0)", jobID)
|
||
}
|
||
if !passed {
|
||
SetJobError(jobCtx, rdb, jobID, "Content check failed: Please upload an indoor room photo")
|
||
return
|
||
}
|
||
|
||
// Step 1: generate floor (single combined prompt — AI identifies floor + replaces it)
|
||
roomLabel := ""
|
||
if room != nil {
|
||
roomLabel = " · " + room.Name
|
||
}
|
||
SetJobProgress(jobCtx, rdb, jobID, "generate", fmt.Sprintf("Generating %s %s%s...", floor.Name, pattern.Name, roomLabel))
|
||
stepStart = time.Now()
|
||
logger.Info("[%s] TIMER step=generate start floor=%s pattern=%s elapsed=%s",
|
||
jobID, floor.Name, pattern.Name, time.Since(totalStart).Round(time.Millisecond))
|
||
|
||
// Material source: image / json / both
|
||
if materialSource == "json" || materialSource == "both" {
|
||
if desc := handler.LookupMaterialDesc(floor.SKU); desc != "" {
|
||
if floor.Description != "" {
|
||
floor.Description = desc + " | " + floor.Description
|
||
} else {
|
||
floor.Description = desc
|
||
}
|
||
}
|
||
}
|
||
refImagePath := ""
|
||
if materialSource != "json" {
|
||
refImagePath = downloadReferenceImage(jobCtx, floor.ImageURL, jobID)
|
||
}
|
||
|
||
imageSize := floorImageSize()
|
||
logger.Info("[%s] TIMER generate request material=%s image_size=%s input=%s output=%s ref=%s",
|
||
jobID, materialSource, imageSize, workPath, outputPath, refImagePath)
|
||
floorPrompt := handler.BuildCombinedFloorPrompt(*floor, *pattern, room, groutLines, whiteBaseboard)
|
||
_, err = client.GenerateImageToFile(jobCtx, openrouter.ImageGenerationRequest{
|
||
Prompt: floorPrompt, OutputPath: outputPath, InputImagePath: workPath,
|
||
ImageSize: imageSize, ReferenceImagePath: refImagePath,
|
||
})
|
||
if err != nil {
|
||
logStepError(jobCtx, rdb, jobID, "generate", totalStart, stepStart, err)
|
||
SetJobError(jobCtx, rdb, jobID, "Floor generation failed: "+err.Error())
|
||
return
|
||
}
|
||
if err := openrouter.RepairNearBlackBorder(outputPath); err != nil {
|
||
logger.Warn("[%s] Failed to repair output border: %v", jobID, err)
|
||
}
|
||
logStepDone(jobCtx, rdb, jobID, "generate", stepStart, totalStart, "")
|
||
|
||
// Clean up intermediate files
|
||
stepStart = time.Now()
|
||
logger.Info("[%s] TIMER step=cleanup start elapsed=%s", jobID, time.Since(totalStart).Round(time.Millisecond))
|
||
if err := os.Remove(inputPath); err != nil {
|
||
logger.Warn("[%s] Failed to remove input: %v", jobID, err)
|
||
}
|
||
if workPath != inputPath {
|
||
if err := os.Remove(workPath); err != nil {
|
||
logger.Warn("[%s] Failed to remove resized input: %v", jobID, err)
|
||
}
|
||
}
|
||
if refImagePath != "" {
|
||
if err := os.Remove(refImagePath); err != nil {
|
||
logger.Warn("[%s] Failed to remove reference image: %v", jobID, err)
|
||
}
|
||
}
|
||
|
||
totalDuration := time.Since(totalStart)
|
||
_ = SetJobTiming(jobCtx, rdb, jobID, "total", totalDuration, totalDuration)
|
||
SetJobResult(jobCtx, rdb, jobID, "/"+filepath.ToSlash(outputPath))
|
||
logger.Info("[%s] TIMER job done total=%s result=%s", jobID, totalDuration.Round(time.Millisecond), "/"+filepath.ToSlash(outputPath))
|
||
}
|
||
|
||
func logStepDone(ctx context.Context, rdb *redis.Client, jobID, step string, stepStart, totalStart time.Time, format string, args ...any) {
|
||
duration := time.Since(stepStart)
|
||
elapsed := time.Since(totalStart)
|
||
_ = SetJobTiming(ctx, rdb, jobID, step, duration, elapsed)
|
||
if format != "" {
|
||
logger.Info("[%s] TIMER step=%s done duration=%s elapsed=%s %s",
|
||
jobID, step, duration.Round(time.Millisecond), elapsed.Round(time.Millisecond), fmt.Sprintf(format, args...))
|
||
return
|
||
}
|
||
logger.Info("[%s] TIMER step=%s done duration=%s elapsed=%s",
|
||
jobID, step, duration.Round(time.Millisecond), elapsed.Round(time.Millisecond))
|
||
}
|
||
|
||
func logStepError(ctx context.Context, rdb *redis.Client, jobID, step string, totalStart, stepStart time.Time, err error) {
|
||
duration := time.Since(stepStart)
|
||
elapsed := time.Since(totalStart)
|
||
_ = SetJobTiming(ctx, rdb, jobID, step, duration, elapsed)
|
||
logger.Error("[%s] TIMER step=%s error duration=%s elapsed=%s err=%v",
|
||
jobID, step, duration.Round(time.Millisecond), elapsed.Round(time.Millisecond), err)
|
||
}
|
||
|
||
// validateContent checks if the image is an indoor room photo.
|
||
// Uses local CLIP ViT-B/32 first, falls back to Gemini 2.5 Flash if CLIP is unavailable.
|
||
func validateContent(geminiClient *openrouter.Client, ctx context.Context, imagePath string) (bool, int) {
|
||
imgData, err := os.ReadFile(imagePath)
|
||
if err != nil {
|
||
return false, 0
|
||
}
|
||
// Try local CLIP server first
|
||
if score, ok := clipCheck(imgData); ok {
|
||
logger.Info("Content check via local CLIP: score=%d passed=%t", score, score >= 6)
|
||
return score >= 6, score
|
||
}
|
||
// Fallback to Gemini
|
||
checkCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
|
||
defer cancel()
|
||
dataURI := "data:image/png;base64," + base64enc(imgData)
|
||
resp, err := geminiClient.Chat(checkCtx, openrouter.ChatRequest{
|
||
Model: "google/gemini-2.5-flash",
|
||
Messages: []openrouter.Message{{Role: "user", Content: []openrouter.ContentPart{
|
||
{Type: "image_url", ImageURL: &openrouter.ImageURL{URL: dataURI}},
|
||
{Type: "text", Text: `请按 1-10 分评估这张图片是否适合做室内地面替换。
|
||
8-10 分：清晰室内房间，地面可见。
|
||
5-7 分：室内图片，但角度或地面可见度一般。
|
||
1-4 分：室外、宠物特写、物体特写、没有可替换地面，或不适合做地面替换。
|
||
只返回一个数字（1-10）。`},
|
||
}}},
|
||
})
|
||
if err != nil || len(resp.Choices) == 0 {
|
||
return true, 0
|
||
}
|
||
score := parseScore(resp.Choices[0].Message.Content)
|
||
return score >= 6, score
|
||
}
|
||
|
||
// clipCheck sends image bytes to the local CLIP server and returns (score, ok).
|
||
func clipCheck(imgData []byte) (int, bool) {
|
||
req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1:5100/check", bytes.NewReader(imgData))
|
||
if err != nil {
|
||
return 0, false
|
||
}
|
||
req.Header.Set("Content-Type", "application/octet-stream")
|
||
resp, err := http.DefaultClient.Do(req)
|
||
if err != nil {
|
||
return 0, false
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != 200 {
|
||
return 0, false
|
||
}
|
||
var result struct {
|
||
Score int `json:"score"`
|
||
Passed bool `json:"passed"`
|
||
}
|
||
if json.NewDecoder(resp.Body).Decode(&result) != nil {
|
||
return 0, false
|
||
}
|
||
return result.Score, true
|
||
}
|
||
|
||
func parseScore(raw string) int {
|
||
for _, s := range []string{"10", "9", "8", "7", "6", "5", "4", "3", "2", "1"} {
|
||
if len(raw) >= len(s) && raw[:len(s)] == s {
|
||
v := 0
|
||
fmt.Sscanf(s, "%d", &v)
|
||
return v
|
||
}
|
||
}
|
||
return 0
|
||
}
|
||
|
||
func base64enc(data []byte) string {
|
||
const tbl = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
|
||
var b []byte
|
||
for i := 0; i < len(data); i += 3 {
|
||
b0, b1, b2 := data[i], byte(0), byte(0)
|
||
if i+1 < len(data) {
|
||
b1 = data[i+1]
|
||
}
|
||
if i+2 < len(data) {
|
||
b2 = data[i+2]
|
||
}
|
||
b = append(b, tbl[b0>>2])
|
||
b = append(b, tbl[((b0&3)<<4)|(b1>>4)])
|
||
if i+1 < len(data) {
|
||
b = append(b, tbl[((b1&15)<<2)|(b2>>6)])
|
||
} else {
|
||
b = append(b, '=')
|
||
}
|
||
if i+2 < len(data) {
|
||
b = append(b, tbl[b2&63])
|
||
} else {
|
||
b = append(b, '=')
|
||
}
|
||
}
|
||
return string(b)
|
||
}
|
||
|
||
// downloadReferenceImage downloads the floor material image from the given URL
|
||
// and saves it to a local temp file. Returns the local file path, or empty string on failure.
|
||
func downloadReferenceImage(ctx context.Context, imageURL, jobID string) string {
|
||
if imageURL == "" {
|
||
return ""
|
||
}
|
||
refPath := filepath.Join("outputs", fmt.Sprintf("ref_%s.png", jobID))
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil)
|
||
if err != nil {
|
||
logger.Warn("[%s] Failed to create reference image request: %v", jobID, err)
|
||
return ""
|
||
}
|
||
resp, err := http.DefaultClient.Do(req)
|
||
if err != nil {
|
||
logger.Warn("[%s] Failed to download reference image from %s: %v", jobID, imageURL, err)
|
||
return ""
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||
logger.Warn("[%s] Reference image download returned HTTP %d from %s", jobID, resp.StatusCode, imageURL)
|
||
return ""
|
||
}
|
||
f, err := os.Create(refPath)
|
||
if err != nil {
|
||
logger.Warn("[%s] Failed to create reference image file: %v", jobID, err)
|
||
return ""
|
||
}
|
||
defer f.Close()
|
||
if _, err := io.Copy(f, resp.Body); err != nil {
|
||
logger.Warn("[%s] Failed to save reference image: %v", jobID, err)
|
||
return ""
|
||
}
|
||
logger.Info("[%s] Downloaded reference image from %s", jobID, imageURL)
|
||
return refPath
|
||
}
|
||
|
||
func payloadBool(payload map[string]string, key string) bool {
|
||
switch payload[key] {
|
||
case "1", "true", "yes", "y", "on":
|
||
return true
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
func floorImageSize() string {
|
||
value := strings.TrimSpace(os.Getenv("FLOOR_IMAGE_SIZE"))
|
||
if value == "" {
|
||
return "2K"
|
||
}
|
||
return value
|
||
}
|