From 4b5f5eb736b2313a6e94d11129e408839f85711f Mon Sep 17 00:00:00 2001 From: dindang Date: Mon, 27 Jul 2026 14:39:45 +0800 Subject: [PATCH] v2 --- README.md | 4 + docker-compose.yml | 6 +- internal/handler/floor_handler.go | 447 +++++++++++++++--------------- internal/openrouter/image.go | 127 ++++++++- internal/queue/worker.go | 49 +++- main.go | 54 ++++ 6 files changed, 456 insertions(+), 231 deletions(-) diff --git a/README.md b/README.md index cba7e37..3efa71e 100644 --- a/README.md +++ b/README.md @@ -199,3 +199,7 @@ AI 图片生成使用 Redis List 做消息队列，Worker 池消费： - Worker 从 BRPOP 出队 → 调用 OpenRouter → 写入结果 - 并发数通过 `WORKER_COUNT` 控制，默认 5 - 进度查询: `GET /floor/status?job_id=xxx` + "description": "Light to medium golden oak-brown wood tone with warm +- yellow undertones and subtle tonal variation, featuring fine to moderately +- pronounced straight and gentle cathedral grain, a smooth lightly textured surface, +- low satin sheen, and a clean, natural, uniform visual character." diff --git a/docker-compose.yml b/docker-compose.yml index 5ff4426..b362d6d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -7,16 +7,14 @@ services: - "8099:8099" volumes: - floor_outputs:/app/outputs - - floor_outputs:/app/outputs - - floor_logs: - floor_uploads:/app/uploads + - floor_uploads:/app/uploads - floor_logs:/app/logs environment: - PORT=8099 - DATABASE_URL=postgres://postgres:floorviz123@db:5432/floorvisualizer?sslmode=disable - REDIS_ADDR=redis:6379 - JWT_SECRET=change-this-to-a-strong-secret-in-production - - OPENROUTER_API_KEY=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 + - OPENROUTER_API_KEY=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 - WORKER_COUNT=5 - LOG_DIR=/app/logs - LOG_LEVEL=INFO diff --git a/internal/handler/floor_handler.go b/internal/handler/floor_handler.go index bc4f4ce..33a64b2 100644 --- a/internal/handler/floor_handler.go +++ b/internal/handler/floor_handler.go @@ -3,6 +3,7 @@ package handler import ( "context" "database/sql" + "encoding/json" "fmt" "io" "net/http" @@ -15,22 +16,23 @@ import ( "floorvisualizer/internal/repository" ) -// ── Floor options ────────────────────────────────────────── +// 閳光偓閳光偓 Floor options 閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓 type FloorOption struct { - ID string `json:"id"` - Name string `json:"name"` - Category string `json:"category"` - Material string `json:"material"` - ColorTone string `json:"color_tone"` - Finish string `json:"finish"` - SizeLabel string `json:"size_label"` - WidthIn float64 `json:"width_in"` - LengthIn float64 `json:"length_in"` - ImageURL string `json:"image_url"` - SKU string `json:"sku"` - Description string `json:"description"` - Variants []FloorOption `json:"variants,omitempty"` + ID string `json:"id"` + Name string `json:"name"` + Category string `json:"category"` + Material string `json:"material"` + ColorTone string `json:"color_tone"` + Finish string `json:"finish"` + SizeLabel string `json:"size_label"` + WidthIn float64 `json:"width_in"` + LengthIn float64 `json:"length_in"` + ImageURL string `json:"image_url"` + SKU string `json:"sku"` + Description string `json:"description"` + FloorDescription string `json:"floor_description,omitempty"` + Variants []FloorOption `json:"variants,omitempty"` } type PatternOption struct { @@ -48,8 +50,8 @@ type RoomOption struct { var woodPatterns = []PatternOption{ {Code: 1, Name: "直铺（工字拼）Straight Lay", Category: "wood", Pattern: "planks running from bottom-left toward top-right of the image, joints staggered randomly, at least 6 inch offset between rows"}, {Code: 2, Name: "横铺 Horizontal Lay", Category: "wood", Pattern: "planks running horizontally left to right, parallel to the bottom edge of the image, joints staggered randomly, at least 6 inch offset"}, - {Code: 3, Name: "人字拼 Herringbone", Category: "wood", Pattern: "planks at 90° forming continuous zigzag V-lines from bottom-left to top-right of the image"}, - {Code: 4, Name: "鱼骨拼 Chevron", Category: "wood", Pattern: "planks with 45° mitered ends forming continuous V-points aligned from bottom-left to top-right of the image"}, + {Code: 3, Name: "人字铺 Herringbone", Category: "wood", Pattern: "ALL planks must be IDENTICAL width - no piece may appear narrower. Planks at 90 degrees forming continuous zigzag V-lines from bottom-left to top-right of the image. At far distance, planks remain distinct with visible seams - NOT merging into a solid mass."}, + {Code: 4, Name: "鱼骨铺 Chevron", Category: "wood", Pattern: "planks with 45 degree mitered ends forming continuous V-points aligned from bottom-left to top-right of the image"}, } var tilePatterns = []PatternOption{ @@ -57,7 +59,7 @@ var tilePatterns = []PatternOption{ {Code: 6, Name: "工字铺（1/2错缝）Running Bond", Category: "tile", Pattern: "rows parallel to bottom edge of image, each row offset 50% from previous, horizontal grout lines continuous, vertical grout lines staggered"}, {Code: 7, Name: "三七错铺 1/3 Offset", Category: "tile", Pattern: "rows parallel to bottom edge of image, each row offset exactly 33% of tile length, every third row aligns"}, {Code: 8, Name: "六边形 Hexagonal", Category: "tile", Pattern: "six-sided tiles interlocked in honeycomb mesh, flat edges horizontal (parallel to bottom edge of image), grout follows hexagonal edges"}, - {Code: 9, Name: "斜铺（60度对角）Diagonal", Category: "tile", Pattern: "tiles rotated 60° to the image frame, grout lines run diagonally at 60° and 150° to the bottom edge, no lines parallel to image edges"}, + {Code: 9, Name: "斜铺（60度对角）Diagonal", Category: "tile", Pattern: "tiles rotated 60 degrees to the image frame, grout lines run diagonally at 60 and 150 degrees to the bottom edge, no lines parallel to image edges"}, } var roomOptions = []RoomOption{ @@ -101,7 +103,7 @@ func loadFloorOptions(db *sql.DB) ([]FloorOption, error) { is_official_price, main_image_url, room_image_url, source_url, description, status, COALESCE(coverage_sqft_per_box,0), COALESCE(wood_species,''), size_label, COALESCE(width_in,0), COALESCE(length_in,0) - FROM products ORDER BY brand, style_name, id LIMIT 3000`) + FROM products ORDER BY brand, style_name, id LIMIT 5000`) if err != nil { return nil, err } @@ -111,7 +113,7 @@ func loadFloorOptions(db *sql.DB) ([]FloorOption, error) { return nil, err } type key struct{ brand, style string } - seen := map[key]int{} // key → index in out slice + seen := map[key]int{} // key 閳?index in out slice var out []FloorOption for _, p := range prods { k := key{p.Brand, p.StyleName} @@ -189,7 +191,7 @@ func findFloorBySKU(db *sql.DB, sku string) *FloorOption { if (canonicalWidth > 0 && canonicalLength > 0) && (f.WidthIn > 0 && f.LengthIn > 0) { if canonicalWidth != f.WidthIn || canonicalLength != f.LengthIn { f.Description = fmt.Sprintf( - `SIZE MAPPING: The reference image shows %.1f"×%.1f" planks. Generate %.1f"×%.1f" planks — same exact material, only scale the board size proportionally. | %s`, + `尺寸映射：参考图显示 %.1f" x %.1f" 的板/砖；本次生成 %.1f" x %.1f" 的板/砖。材质必须完全相同，只按比例调整板/砖尺寸。 %s`, canonicalWidth, canonicalLength, f.WidthIn, f.LengthIn, f.Description, ) } @@ -197,7 +199,7 @@ func findFloorBySKU(db *sql.DB, sku string) *FloorOption { return f } -// ── Public accessors for worker ──────────────────────────── +// 閳光偓閳光偓 Public accessors for worker 閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓 func FindFloorBySKUPublic(db *sql.DB, sku string) *FloorOption { return findFloorBySKU(db, sku) } func FindPatternByCodePublic(code string) *PatternOption { @@ -212,58 +214,42 @@ func FindRoomByCodePublic(code string) *RoomOption { } func BuildMaskPromptPublic(room *RoomOption) string { return buildMaskPrompt(room) } func BuildFloorPromptPublic(f FloorOption, p PatternOption, room *RoomOption) string { - return buildFloorPrompt(f, p, room) + return buildCompactFloorPrompt(f, p, room) } // BuildCombinedFloorPrompt builds a single prompt that identifies the floor area AND replaces it. -// No separate mask generation step — the AI does both in one pass. -func BuildCombinedFloorPrompt(f FloorOption, p PatternOption, room *RoomOption) string { - floorPrompt := buildFloorPrompt(f, p, room) - roomCtx := "" - if room != nil && room.Code != 0 { - roomCtx = fmt.Sprintf(" This is a %s.", room.Name) - } - - return fmt.Sprintf( - `FIRST — identify the floor: Look at this room photo.%s -The floor is the flat horizontal walking surface — NOT walls, NOT baseboards, NOT furniture, NOT stairs. -You will replace ONLY the floor area. Every other pixel must remain bit-identical to the input photo. - -%s - -Return ONLY the completed image. No text.`, - roomCtx, - floorPrompt, - ) +// No separate mask generation step 閳?the AI does both in one pass. +func BuildCombinedFloorPrompt(f FloorOption, p PatternOption, room *RoomOption, groutLines, whiteBaseboard bool) string { + return buildCompactCombinedFloorPrompt(f, p, room, groutLines, whiteBaseboard) } // buildPatternConstraint returns installation rules specific to each pattern code. func buildPatternConstraint(code int) string { switch code { case 1: - return `PATTERN RULES — Straight Lay (wood): Planks run from the bottom-left of the image toward the top-right, at roughly 45° to the image frame. Long edges of each plank point toward the upper-right corner. Joints staggered 6" minimum — random pattern. 1/8" expansion gap.` + return "铺法规则：木地板直铺，板材连续铺设，接缝随机错开至少 6 英寸。" case 2: - return `PATTERN RULES — Horizontal Lay (wood): Planks run horizontally — left to right, parallel to the bottom edge of the image. Long edges of each plank are horizontal across the photo. Joints staggered 6" minimum — random pattern. 1/8" expansion gap.` + return "铺法规则：木地板横铺，板材沿房间主方向铺设，接缝随机错开至少 6 英寸。" case 3: - return `PATTERN RULES — Herringbone (wood): Planks at 90° forming continuous zigzag. Each plank end meets neighbor at 90°. V-lines run from bottom-left to top-right of the image. Joints tight — continuous unbroken zigzag.` + return "铺法规则：人字拼，板材 90 度相接形成连续折线，每片板宽一致，接缝清楚。" case 4: - return `PATTERN RULES — Chevron (wood): Planks with 45° mitered ends form continuous V. V-points aligned in a straight line running bottom-left to top-right across the image. Joint lines are straight and continuous. NOT herringbone — mitered joints, not staggered.` + return "铺法规则：鱼骨拼，板材 45 度斜切端形成连续 V 形，必须是斜切对接。" case 5: - return `PATTERN RULES — Straight Grid (tile): Tiles aligned in rows and columns. Grout lines run horizontally (parallel to bottom edge of image) and vertically (perpendicular to bottom edge). NO offset between rows. Grout ~1/8".` + return "铺法规则：瓷砖直铺网格对缝，砖块成行成列对齐，无错缝，砖缝约 1/8 英寸。" case 6: - return `PATTERN RULES — Running Bond (tile): Each row offset 50% from previous — classic brickwork. Grout lines parallel to bottom edge of image are continuous; grout lines perpendicular to bottom edge are staggered. Grout ~1/8".` + return "铺法规则：瓷砖工字铺，相邻行错开砖长的 50%，行向砖缝连续，竖向接缝错开。" case 7: - return `PATTERN RULES — 1/3 Offset (tile): Each row offset exactly 33% of tile length. Every third row aligns. Rows run parallel to bottom edge of image. Grout ~1/8". NOT 50% bond — offset is exactly 1/3.` + return "铺法规则：瓷砖三七错铺，相邻行错开砖长的 33%，每三行重新对齐，不是 50% 工字铺。" case 8: - return `PATTERN RULES — Hexagonal (tile): Six-sided tiles interlocked in honeycomb mesh. Each tile touches 6 neighbors. Flat edges of hexagons are horizontal (parallel to bottom edge of image). Grout follows hexagonal edges — network pattern, not grid.` + return "铺法规则：六边形瓷砖，六边形砖互锁成蜂窝网，砖缝沿六边形边缘，不是方格网。" case 9: - return `PATTERN RULES — Diagonal (tile): Tiles rotated 60° relative to the image frame — all tiles at a steep diagonal, diamond orientation. Grout lines run diagonally across the photo at 60° and 150° to the bottom edge. None parallel or perpendicular to image edges. Triangle cuts at all walls. Grout ~1/8".` + return "铺法规则：瓷砖斜铺，砖块相对墙边形成菱形或三角切砖效果，砖缝约 1/8 英寸。" default: return "" } } -// ── Floor Generate Handler (Redis queue) ────────────────── +// 閳光偓閳光偓 Floor Generate Handler (Redis queue) 閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓 func SetQueue(enqueue func(ctx context.Context, jobID string, payload map[string]string) error) { generateQueue = enqueue @@ -334,10 +320,13 @@ func FloorGenerate(client *openrouter.Client, db *sql.DB) http.HandlerFunc { } // Store job payload then enqueue payload := map[string]string{ - "input_path": inputPath, - "floor_id": r.FormValue("floor_id"), - "pattern_code": r.FormValue("pattern_code"), - "room_code": r.FormValue("room_code"), + "input_path": inputPath, + "floor_id": r.FormValue("floor_id"), + "pattern_code": r.FormValue("pattern_code"), + "room_code": r.FormValue("room_code"), + "grout_lines": boolFormValue(r, "grout_lines"), + "white_baseboard": boolFormValue(r, "white_baseboard"), + "material_source": r.FormValue("material_source"), } if err := generateQueue(r.Context(), jobID, payload); err != nil { writeErr(w, 500, "Failed to enqueue job") @@ -348,9 +337,9 @@ func FloorGenerate(client *openrouter.Client, db *sql.DB) http.HandlerFunc { } } -// ── Progress ─────────────────────────────────────────────── +// 閳光偓閳光偓 Progress 閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓 -// ── Floor Status Query ───────────────────────────────── +// 閳光偓閳光偓 Floor Status Query 閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓 func FloorStatus(w http.ResponseWriter, r *http.Request) { jobID := r.URL.Query().Get("job_id") @@ -370,7 +359,7 @@ func FloorStatus(w http.ResponseWriter, r *http.Request) { writeJSON(w, 200, js) } -// ── Options ──────────────────────────────────────────────── +// 閳光偓閳光偓 Options 閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓 func FloorOptions(db *sql.DB) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { @@ -387,185 +376,201 @@ func FloorOptions(db *sql.DB) http.HandlerFunc { } } + // Load special test products with floor descriptions + specialFloors := loadSpecialFloors(db) + writeJSON(w, 200, map[string]any{ - "wood_floors": allWood, - "wood_patterns": woodPatterns, - "tile_floors": allTile, - "tile_patterns": tilePatterns, - "vinyl_floors": vinylFloors, - "rooms": roomOptions, + "wood_floors": allWood, + "wood_patterns": woodPatterns, + "tile_floors": allTile, + "tile_patterns": tilePatterns, + "vinyl_floors": vinylFloors, + "special_floors": specialFloors, + "rooms": roomOptions, }) } } -// ── Prompts ──────────────────────────────────────────────── +// materialIndex caches SKU 閳?material description from MaterialAssets JSON files. +var materialIndex map[string]string + +// LoadMaterialIndex loads the material description index from a JSON file. +func LoadMaterialIndex(path string) error { + data, err := os.ReadFile(path) + if err != nil { + return err + } + return json.Unmarshal(data, &materialIndex) +} + +// LookupMaterialDesc returns the material description for a given SKU from the loaded index. +func LookupMaterialDesc(sku string) string { + if materialIndex == nil { + return "" + } + return materialIndex[sku] +} + +// loadSpecialFloors loads the 7 special test products with pre-generated floor descriptions. +func loadSpecialFloors(db *sql.DB) []FloorOption { + data, err := os.ReadFile("outputs/special_7_products.json") + if err != nil { + return nil + } + var raw []struct { + Brand string `json:"brand"` + Style string `json:"style"` + Material string `json:"material"` + Description string `json:"description"` + ImageURL string `json:"image_url"` + } + if json.Unmarshal(data, &raw) != nil { + return nil + } + var out []FloorOption + for _, r := range raw { + // Find any real SKU for this style to get DB variants and sizes + var sku string + db.QueryRow(`SELECT sku FROM products WHERE style_name=$1 AND brand=$2 LIMIT 1`, + r.Style, r.Brand).Scan(&sku) + if sku == "" { + continue + } + f := findFloorBySKU(db, sku) + if f == nil { + continue + } + f.FloorDescription = r.Description + // Keep f.ImageURL from DB (canonical smallest) 閳?do NOT override with local path + out = append(out, *f) + } + return out +} + +// 閳光偓閳光偓 Prompts 閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓閳光偓 func buildMaskPrompt(room *RoomOption) string { - base := `You are a precise floor segmentation tool. Create a pure black-and-white mask for this room photo. - -STRICT DEFINITIONS: -- WHITE (#FFFFFF): ONLY the flat, horizontal walking surface (the floor). Include the entire visible floor area up to the exact edge where it meets walls, baseboards, cabinets, or stairs. -- BLACK (#000000): EVERYTHING that is not flat horizontal floor — walls, baseboards/skirting boards, stairs/stair risers (they are NOT flat floor), furniture and furniture legs, doors and door frames, windows, cabinets, appliances, rugs/carpets, ceiling, decorations, people, pets. - -EDGE RULES: -- At the wall-floor boundary: cut exactly along the junction. The baseboard/skirting board is BLACK, the floor is WHITE. -- Furniture legs resting on the floor: the legs are BLACK. The floor visible BETWEEN and AROUND legs is WHITE. -- Stairs: each step's horizontal tread AND vertical riser are BLACK (stairs are not a continuous flat plane). -- Rugs/carpets on the floor: BLACK (they are not the permanent floor). - -SELF-CHECK before output: -- Did you include any wall area? → fix it. -- Did you miss any floor area between furniture legs? → fill it WHITE. -- Did you paint any stair surface WHITE? → make it BLACK. -- Are all edges sharp and precise at boundaries? → verify. - -Output ONLY the mask image. No text, no explanation.` + base := "你是精确的地面分割工具。请为这张室内照片生成纯黑白遮罩图。\n\n" + + "严格定义：\n" + + "- 白色（#FFFFFF）：只表示平整、水平、可行走的地面。包含所有可见地面，一直到墙面、踢脚线、柜体或楼梯边缘。\n" + + "- 黑色（#000000）：所有非水平地面的内容，包括墙面、踢脚线、楼梯、家具、家具腿、门窗、柜体、电器、地毯、天花板、人物、宠物。\n\n" + + "边缘规则：\n" + + "- 墙地交界处必须沿交界线精确切分；踢脚线为黑色，地面为白色。\n" + + "- 家具腿和椅轮为黑色；家具腿之间和周围真实可见的地面为白色。\n" + + "- 楼梯和地毯都标为黑色。\n\n" + + "只输出遮罩图片，不要文字说明。" if room != nil && room.Code != 0 { - base += fmt.Sprintf("\nContext: This is a %s.", room.Name) + base += fmt.Sprintf("\n房间类型：%s。", room.Name) } return base } -func buildFloorPrompt(floor FloorOption, pattern PatternOption, room *RoomOption) string { - var roomCtx string - if room != nil && room.Code != 0 { - roomCtx = fmt.Sprintf("\nRoom context: This is a %s.", room.Name) - } - matType := "wood plank" +func buildMaterialRules(floor FloorOption) string { if isTileProduct(floor) || isLaminateProduct(floor) { - matType = "tile" + return "整片地面必须是同一个 SKU，基础颜色在全地面保持一致。允许原图阴影让局部变暗，但不能改变地砖色相或材质。瓷砖表面平整，缝隙连续、透视一致；不要生成局部偏绿、偏蓝、偏黄、偏青的色块。反光只能来自原图已有光照，不能新增窗户反光、亮斑或湿地板效果。" + } + return "地面必须保持同一个 SKU 的颜色、纹理、木纹或板纹和光泽。保留原图已有光照和阴影；不要新增反光、亮斑或阴影。板缝透视一致，纹理自然但不能明显重复。" +} + +func buildCompactCombinedFloorPrompt(f FloorOption, p PatternOption, room *RoomOption, groutLines, whiteBaseboard bool) string { + roomName := "自动识别" + if room != nil && room.Code != 0 { + roomName = room.Name } - materialDesc := fmt.Sprintf("%s material, %s tone, %s finish", floor.Material, floor.ColorTone, floor.Finish) - if floor.SizeLabel != "" { - materialDesc += fmt.Sprintf(", %s size", floor.SizeLabel) + var priority []string + if whiteBaseboard { + priority = append(priority, "最高优先级：必须生成白色踢脚线。\n踢脚线属于可编辑区域，不属于墙面保护区域。\n沿所有可见墙地交界处生成连续白色踢脚线，高约 4-6 英寸。\n踢脚线应为真实白色木质或漆面条带，有轻微厚度和接触阴影。\n如果没有明显、连续的白色踢脚线，结果视为失败。") + } + if groutLines { + priority = append(priority, "必须生成清晰、等距、连续的瓷砖缝。") } - // Include product description if available for richer visual detail - var descriptionText string - if floor.Description != "" { - descriptionText = fmt.Sprintf("\nVisual appearance: %s", floor.Description) + priorityText := "" + if len(priority) > 0 { + priorityText = strings.Join(priority, "\n\n") + "\n\n" } - productLock := "" - if floor.ImageURL != "" { - productLock = fmt.Sprintf( - `The reference image is the EXACT product material — not style inspiration, not a mood board. -LOCK THESE PROPERTIES from the reference image: base color, color temperature, grain/vein pattern, -knot density, surface texture, gloss level, material character. -FORBIDDEN: color shift, different wood species or stone type, new grain pattern, different stain, -gloss change, "luxury upgrade" or aesthetic reinterpretation.`, - ) + editScope := "只替换地面" + if whiteBaseboard { + editScope = "只替换地面和白色踢脚线" } - // Size instruction: only change board proportion, never the material itself - sizeInstruction := "" - if len(floor.Variants) > 0 { - var variantList []string - for _, v := range floor.Variants { - marker := "" - if v.SKU == floor.SKU { - marker = " ← SELECTED" - } - if v.SizeLabel != "" { - variantList = append(variantList, v.SizeLabel+marker) - } else if v.WidthIn > 0 && v.LengthIn > 0 { - variantList = append(variantList, fmt.Sprintf(`%.1f"×%.1f"`, v.WidthIn, v.LengthIn)+marker) - } - } - if len(variantList) > 0 { - sizeInstruction = fmt.Sprintf( - "\nSIZE: This product comes in multiple sizes: %s.\nONLY change board width and seam spacing. The material texture, color, and grain are from the reference image — identical in every size. Wider boards use the SAME texture source, just with fewer seams. Do NOT reinterpret the wood to look \"more premium\" for wider planks.", - strings.Join(variantList, ", "), - ) - } - } else if floor.SizeLabel != "" { - sizeInstruction = fmt.Sprintf( - "\nSIZE: Each piece is %s. Only adjust board proportion and seam spacing. The material texture is from the reference image — do not change it.", - floor.SizeLabel, - ) - } else if floor.WidthIn > 0 && floor.LengthIn > 0 { - sizeInstruction = fmt.Sprintf( - `\nSIZE: Each piece is %.1f" × %.1f". Only adjust board proportion and seam spacing. The material texture is from the reference image — do not change it.`, - floor.WidthIn, floor.LengthIn, - ) - } - - patternInstruction := buildPatternConstraint(pattern.Code) - - return fmt.Sprintf( - `You are a flooring product visualization renderer for a commercial catalog. Your sole function is to replace the floor surface in interior photos with a specific flooring SKU. This is product-accurate catalog photography — NOT interior design or artistic interpretation. - -STRICT OUTPUT RULES: -1. Respond with ONLY the completed PNG image. No text, no explanation, no JSON wrapper. -2. The output must have the exact same pixel dimensions as the input image. - Do not crop, pad, upscale, downscale, rotate, or alter the image geometry in any way. -3. Do not apply global filtering, sharpening, blurring, tone-mapping, contrast adjustment, - or style transfer. The only permitted change is the floor surface. - -FLOOR SURFACE DEFINITION: -- "Floor" — the flat horizontal walking surface. The visible ground plane inside the room. - DO NOT bleed onto: walls, baseboards/skirting boards, stairs, stair risers, furniture, - chair legs, table legs, cabinet bases, doors, door frames, window glass or frames, - trim, pipes, conduit, rugs, carpets, ceiling, appliances, decorations. -- Stairs are NOT floor — each tread and riser must remain completely unchanged. -- If a surface is not the flat horizontal floor, PRESERVE IT EXACTLY. - -FLOOR REPLACEMENT RULES: -4. Apply the new floor ONLY to the floor surface. Every pixel outside must be bit-identical to input. -5. Preserve the original room lighting, shadows, and highlights ON the new floor. -6. ADD realistic contact shadows: chair wheels, table legs, cabinet bases MUST cast tight dark shadows - directly at the contact point — sharpest there, softening outward. This shows physical weight. -7. Maintain clean but physically installed-looking floor edges — not a Photoshop cutout. -8. The new floor must look physically installed and photographed — not rendered, not composited. - -FIDELITY RULES: -9. The output must be a photorealistic edit of the input photograph. - Do not generate new content, synthesise areas, invent objects, or alter composition. -10. Preserve camera sensor noise, natural vignette, and original white balance on all non-floor surfaces. - -AMBIGUITY RULE: -11. If a pixel could be floor OR furniture/wall/baseboard — preserve the original unchanged. - Precision over completeness. - -FLOOR SPECIFICATION: -- Product: %s (%s %s) -- Pattern: %s — %s - The pattern changes layout only. The material comes from the reference image — identical regardless of pattern. -%s -- Material details: %s%s%s%s - -%s - -TEXTURE FIDELITY: -- Every plank/tile must have UNIQUE grain, vein, or tonal variation — no visible texture repetition. -- Natural imperfections: minor scratches, micro-wear near chair wheels, colour variance between pieces. -- Seams and joints: subtle dust in gaps, micro-shadow at edges, slight edge wear. -- Polished surfaces: reflections UNEVEN and angle-dependent — not uniform glossy sheen. -- Do NOT invent sunbeams, light rays, or window reflections not in the original. - -NEGATIVE CONSTRAINTS — the floor must NOT contain: -CGI look, 3D render, plastic texture, repeating patterns, uniform gloss, invented lights, -HDR tone-mapping, "luxury showroom" aesthetic, sterile clean seams, floating furniture, -color shift, different wood species, new grain, different stain, aesthetic reinterpretation. - -REQUIREMENTS SUMMARY: -✓ Only the flat horizontal floor changes — all other pixels preserved bit-identical. -✓ Contact shadows under all furniture — sharp at contact, softening outward. -✓ Original lighting and shadows match on new floor. -✓ Output geometry identical to input — same dimensions, no cropping. -✓ Scene composition unchanged — no new objects, no synthesised content. -✓ Material color, grain, gloss, species match the reference image exactly. - -Return ONLY the completed image. No text.`, - floor.Name, matType, materialDesc, - pattern.Name, pattern.Pattern, patternInstruction, - materialDesc, - roomCtx, descriptionText, sizeInstruction, productLock, + return fmt.Sprintf("%s你是室内照片编辑工具，不是重新生成图片。\n房间类型：%s。\n\n任务：\n%s。墙面、家具、桌椅、门窗、天花板保持原图不变。\n保持原图尺寸、构图、相机角度、光照、阴影、白平衡不变。\n\n%s\n\n输出：\n只返回完成后的图片，不要文字、不要 JSON。", + priorityText, + roomName, + editScope, + buildCompactFloorPrompt(f, p, room), ) } -// ── Content Check ───────────────────────────────────────── +func buildCompactFloorPrompt(floor FloorOption, pattern PatternOption, room *RoomOption) string { + matType := "木地板" + if isTileProduct(floor) || isLaminateProduct(floor) { + matType = "瓷砖" + } + + materialDesc := fmt.Sprintf("%s，%s色调，%s表面", floor.Material, floor.ColorTone, floor.Finish) + if floor.SizeLabel != "" { + materialDesc += "，尺寸 " + floor.SizeLabel + } + + var details []string + if floor.Description != "" { + details = append(details, "产品描述："+floor.Description) + } + if floor.FloorDescription != "" { + details = append(details, "目录材质描述："+floor.FloorDescription) + } + if floor.SizeLabel != "" || (floor.WidthIn > 0 && floor.LengthIn > 0) { + details = append(details, compactSizeInstruction(floor)) + } + if floor.ImageURL != "" { + details = append(details, "产品图是唯一材质参考：锁定基础颜色、纹理、纹路、光泽、材质类型；不要重新美化或改色。") + } + if guard := compactProductNameGuard(floor); guard != "" { + details = append(details, guard) + } + + return fmt.Sprintf("地板产品：\nSKU：%s\n产品名：%s（这是商品目录名，不是视觉效果指令）\n类型：%s\n材质：%s\n铺法：%s。%s\n\n材质规则：\n%s\n%s\n\n禁止：\nno wet floor, no puddles, no water stains, no color shift, no green/blue/yellow/cyan patches, no new lights, no new shadows, no CGI, no text, no watermark。", + floor.SKU, + floor.Name, + matType, + materialDesc, + pattern.Name, + pattern.Pattern, + buildMaterialRules(floor), + strings.Join(details, "\n"), + ) +} + +func compactSizeInstruction(floor FloorOption) string { + if floor.SizeLabel != "" { + return "尺寸规则：使用所选尺寸 " + floor.SizeLabel + "；只调整板/砖大小和缝距，不改变材质颜色和纹理。" + } + if floor.WidthIn > 0 && floor.LengthIn > 0 { + return fmt.Sprintf("尺寸规则：每片 %.1f\" x %.1f\"；只调整板/砖大小和缝距，不改变材质颜色和纹理。", floor.WidthIn, floor.LengthIn) + } + return "" +} + +func compactProductNameGuard(floor FloorOption) string { + text := strings.ToLower(floor.Name + " " + floor.Description + " " + floor.FloorDescription) + for _, word := range []string{"rain", "puddle", "water", "wet", "storm", "mist", "fog", "snow", "ice"} { + if strings.Contains(text, word) { + return "产品名防误解：不要按产品名生成雨水、积水、水渍、雾气、冰雪、湿地板或彩色斑块；只看产品参考图和结构化材质字段。" + } + } + return "" +} + +func buildFloorPrompt(floor FloorOption, pattern PatternOption, room *RoomOption) string { + return buildCompactFloorPrompt(floor, pattern, room) +} + +func buildProductNameGuard(floor FloorOption) string { + return compactProductNameGuard(floor) +} func validateContent(client *openrouter.Client, ctx context.Context, imagePath string) (bool, int) { imgData, err := os.ReadFile(imagePath) @@ -573,13 +578,12 @@ func validateContent(client *openrouter.Client, ctx context.Context, imagePath s return false, 0 } dataURI := "data:image/png;base64," + b64enc(imgData) + prompt := "请按 1-10 分评估这张图片是否适合做室内地面替换。\n8-10 分：清晰室内房间，地面可见。\n5-7 分：室内图片，但角度或地面可见度一般。\n1-4 分：室外、宠物特写、物体特写、没有可替换地面，或不适合做地面替换。\n只返回一个数字（1-10）。" resp, err := client.Chat(ctx, 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: `Score this image 1-10 for indoor floor replacement suitability. -8-10: Clear indoor room, floor visible. 5-7: Indoor but poor angle. 1-4: Outdoor/pets/objects/no floor. -Return ONLY a number (1-10).`}, + {Type: "text", Text: prompt}, }}}, }) if err != nil || len(resp.Choices) == 0 { @@ -588,7 +592,6 @@ Return ONLY a number (1-10).`}, score := parseScore(resp.Choices[0].Message.Content) return score >= 6, score } - 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 { @@ -626,3 +629,13 @@ func b64enc(data []byte) string { } return b.String() } + +func boolFormValue(r *http.Request, key string) string { + value := strings.ToLower(strings.TrimSpace(r.FormValue(key))) + switch value { + case "1", "true", "yes", "y", "on": + return "true" + default: + return "false" + } +} diff --git a/internal/openrouter/image.go b/internal/openrouter/image.go index ff92d99..b696efe 100644 --- a/internal/openrouter/image.go +++ b/internal/openrouter/image.go @@ -7,6 +7,7 @@ import ( stderrors "errors" "fmt" "image" + "image/color" _ "image/gif" _ "image/jpeg" "image/png" @@ -22,9 +23,9 @@ import ( const Gemini3ProImagePreviewModel = "google/gemini-3-pro-image-preview" -const maskImageInstruction = "The next image is a mask. Non-black pixels mark regions to edit; black pixels mark protected regions that must remain unchanged." +const maskImageInstruction = "下一张图片是遮罩图。非黑色像素表示允许编辑的区域；黑色像素表示必须保持不变的保护区域。" -const referenceImageInstruction = "The next image is a reference for the floor material style. Match its color, texture, pattern, and overall appearance exactly when generating the new floor. The generated floor should look like this material." +const referenceImageInstruction = "下一张图片是地板材质参考图。生成新地面时必须严格匹配它的颜色、纹理、图案、光泽和整体材质外观；生成的地面应看起来就是这个材质。" var supportedAspectRatios = []struct { value string @@ -518,3 +519,125 @@ func ResizeImageToFit(inputPath, outputPath string) error { } return os.WriteFile(outputPath, buf.Bytes(), 0644) } + +// RepairNearBlackBorder removes accidental model-added black canvas edges. +// It only acts when an entire edge row/column is almost uniformly black. +func RepairNearBlackBorder(path string) error { + data, err := os.ReadFile(path) + if err != nil { + return fmt.Errorf("repair border: read image: %w", err) + } + src, _, err := image.Decode(bytes.NewReader(data)) + if err != nil { + return fmt.Errorf("repair border: decode image: %w", err) + } + b := src.Bounds() + w, h := b.Dx(), b.Dy() + if w < 8 || h < 8 { + return nil + } + + img := image.NewRGBA(image.Rect(0, 0, w, h)) + for y := 0; y < h; y++ { + for x := 0; x < w; x++ { + img.Set(x, y, src.At(b.Min.X+x, b.Min.Y+y)) + } + } + + maxBorder := minInt(32, minInt(w, h)/20) + if maxBorder < 2 { + maxBorder = 2 + } + + top := 0 + for top < maxBorder && isNearBlackRow(img, top) { + top++ + } + bottom := 0 + for bottom < maxBorder && isNearBlackRow(img, h-1-bottom) { + bottom++ + } + left := 0 + for left < maxBorder && isNearBlackCol(img, left) { + left++ + } + right := 0 + for right < maxBorder && isNearBlackCol(img, w-1-right) { + right++ + } + if top == 0 && bottom == 0 && left == 0 && right == 0 { + return nil + } + if top+bottom >= h || left+right >= w { + return nil + } + + for y := 0; y < top; y++ { + copyRow(img, top, y) + } + for y := h - bottom; y < h; y++ { + copyRow(img, h-bottom-1, y) + } + for x := 0; x < left; x++ { + copyCol(img, left, x) + } + for x := w - right; x < w; x++ { + copyCol(img, w-right-1, x) + } + + var buf bytes.Buffer + if err := png.Encode(&buf, img); err != nil { + return fmt.Errorf("repair border: encode png: %w", err) + } + if err := os.WriteFile(path, buf.Bytes(), 0644); err != nil { + return fmt.Errorf("repair border: write image: %w", err) + } + return nil +} + +func isNearBlackRow(img *image.RGBA, y int) bool { + w := img.Bounds().Dx() + black := 0 + for x := 0; x < w; x++ { + if isNearBlack(img.RGBAAt(x, y)) { + black++ + } + } + return black*100 >= w*96 +} + +func isNearBlackCol(img *image.RGBA, x int) bool { + h := img.Bounds().Dy() + black := 0 + for y := 0; y < h; y++ { + if isNearBlack(img.RGBAAt(x, y)) { + black++ + } + } + return black*100 >= h*96 +} + +func isNearBlack(c color.RGBA) bool { + return c.A > 0 && c.R <= 18 && c.G <= 18 && c.B <= 18 +} + +func copyRow(img *image.RGBA, srcY, dstY int) { + w := img.Bounds().Dx() + for x := 0; x < w; x++ { + img.SetRGBA(x, dstY, img.RGBAAt(x, srcY)) + } +} + +func copyCol(img *image.RGBA, srcX, dstX int) { + h := img.Bounds().Dy() + for y := 0; y < h; y++ { + img.SetRGBA(dstX, y, img.RGBAAt(srcX, y)) + } +} + +func minInt(a, b int) int { + if a < b { + return a + } + return b +} diff --git a/internal/queue/worker.go b/internal/queue/worker.go index 719b5b3..203b7be 100644 --- a/internal/queue/worker.go +++ b/internal/queue/worker.go @@ -69,6 +69,12 @@ func processJob(ctx context.Context, rdb *redis.Client, client *openrouter.Clien 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") @@ -103,7 +109,7 @@ func processJob(ctx context.Context, rdb *redis.Client, client *openrouter.Clien return } - // Step 1+2 combined: identify floor + replace in one call (no separate mask) + // Step 1: generate floor (single combined prompt — AI identifies floor + replaces it) roomLabel := "" if room != nil { roomLabel = " · " + room.Name @@ -113,12 +119,25 @@ func processJob(ctx context.Context, rdb *redis.Client, client *openrouter.Clien logger.Info("[%s] TIMER step=generate start floor=%s pattern=%s elapsed=%s", jobID, floor.Name, pattern.Name, time.Since(totalStart).Round(time.Millisecond)) - refImagePath := downloadReferenceImage(jobCtx, floor.ImageURL, jobID) + // 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 image_size=%s input=%s output=%s ref=%s", - jobID, imageSize, workPath, outputPath, refImagePath) - floorPrompt := handler.BuildCombinedFloorPrompt(*floor, *pattern, room) + 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, @@ -128,6 +147,9 @@ func processJob(ctx context.Context, rdb *redis.Client, client *openrouter.Clien 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 @@ -194,9 +216,11 @@ func validateContent(geminiClient *openrouter.Client, ctx context.Context, image 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: `Score this image 1-10 for indoor floor replacement suitability. -8-10: Clear indoor room, floor visible. 5-7: Indoor but poor angle. 1-4: Outdoor/pets/objects/no floor. -Return ONLY a number (1-10).`}, + {Type: "text", Text: `请按 1-10 分评估这张图片是否适合做室内地面替换。 +8-10 分：清晰室内房间，地面可见。 +5-7 分：室内图片，但角度或地面可见度一般。 +1-4 分：室外、宠物特写、物体特写、没有可替换地面，或不适合做地面替换。 +只返回一个数字（1-10）。`}, }}}, }) if err != nil || len(resp.Choices) == 0 { @@ -305,6 +329,15 @@ func downloadReferenceImage(ctx context.Context, imageURL, jobID string) string 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 == "" { diff --git a/main.go b/main.go index fa309e2..ccc629a 100644 --- a/main.go +++ b/main.go @@ -7,6 +7,7 @@ import ( "math/rand" "net/http" "os" + "os/exec" "os/signal" "strconv" "time" @@ -41,6 +42,19 @@ func init() { func main() { godotenv.Load() + // Load material description index + if err := handler.LoadMaterialIndex("outputs/material_index.json"); err != nil { + logger.Warn("Material index not loaded: %v", err) + } else { + logger.Info("Material index loaded") + } + + // Start CLIP server for content check + clipCmd := startCLIPServer() + if clipCmd != nil { + defer stopCLIPServer(clipCmd) + } + logDir := "logs" if v := os.Getenv("LOG_DIR"); v != "" { logDir = v @@ -183,6 +197,7 @@ func main() { if cancelWorkers != nil { cancelWorkers() } + stopCLIPServer(clipCmd) shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() if err := server.Shutdown(shutdownCtx); err != nil { @@ -196,3 +211,42 @@ func main() { } logger.Info("Server stopped gracefully") } + +// startCLIPServer launches the Python CLIP server for content checking. +func startCLIPServer() *exec.Cmd { + // Check if CLIP is already running + if _, err := http.Get("http://127.0.0.1:5100/health"); err == nil { + logger.Info("CLIP server already running on :5100") + return nil + } + pythonPath := `C:\Users\10794\AppData\Local\Programs\Python\Python312\python.exe` + cmd := exec.Command(pythonPath, "clip_server/server.py") + cmd.Stdout = nil + cmd.Stderr = nil + if err := cmd.Start(); err != nil { + logger.Warn("Failed to start CLIP server: %v (content check will use Gemini fallback)", err) + return nil + } + logger.Info("CLIP server starting on :5100...") + // Wait for it to be ready + for i := 0; i < 30; i++ { + time.Sleep(time.Second) + if _, err := http.Get("http://127.0.0.1:5100/health"); err == nil { + logger.Info("CLIP server ready") + return cmd + } + } + logger.Warn("CLIP server did not become ready in time") + return cmd +} + +// stopCLIPServer stops the CLIP Python server. +func stopCLIPServer(cmd *exec.Cmd) { + if cmd == nil || cmd.Process == nil { + return + } + logger.Info("Stopping CLIP server...") + if err := cmd.Process.Kill(); err != nil { + logger.Warn("Failed to stop CLIP server: %v", err) + } +}