wechat_ai/agent/regions.go

205 lines
5.9 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package main
import (
"encoding/json"
"fmt"
"math"
"os"
"strings"
)
func loadAnnotation(path string) (AnnotationFile, error) {
data, err := os.ReadFile(path)
if err != nil {
return AnnotationFile{}, err
}
var annotation AnnotationFile
if err := json.Unmarshal(data, &annotation); err != nil {
return AnnotationFile{}, err
}
return annotation, nil
}
func summarizeRegions(regions []Region) []RegionSummary {
summaries := make([]RegionSummary, 0, len(regions))
for _, region := range regions {
bboxSource := region.BBoxSource
if isZeroBox(bboxSource) {
scaleFactor := region.ScaleFactor
if scaleFactor == 0 {
scaleFactor = 1
}
bboxSource = [4]float64{
region.BBoxImage[0] / scaleFactor,
region.BBoxImage[1] / scaleFactor,
region.BBoxImage[2] / scaleFactor,
region.BBoxImage[3] / scaleFactor,
}
}
summaries = append(summaries, RegionSummary{
Name: region.Name,
Description: region.Description,
Type: region.Type,
BBoxScreen: region.BBoxScreen,
BBoxSource: bboxSource,
CenterScreen: centerOf(region.BBoxScreen),
CenterSource: centerOf(bboxSource),
})
}
return summaries
}
func centerOf(box [4]float64) [2]float64 {
return [2]float64{(box[0] + box[2]) / 2, (box[1] + box[3]) / 2}
}
func regionByType(summaries []RegionSummary) map[string]RegionSummary {
index := make(map[string]RegionSummary, len(summaries))
for _, summary := range summaries {
if _, exists := index[summary.Type]; !exists {
index[summary.Type] = summary
}
}
return index
}
func validateRequiredRegions(summaries []RegionSummary) error {
index := regionByType(summaries)
missing := make([]string, 0)
for _, regionType := range []string{"chat_content", "input_box"} {
if _, exists := index[regionType]; !exists {
missing = append(missing, regionType)
}
}
if len(missing) > 0 {
return fmt.Errorf("missing required regions: %v", missing)
}
return nil
}
func inferMissingRegionTypes(summaries []RegionSummary) ([]RegionSummary, []string) {
index := regionByType(summaries)
warnings := make([]string, 0)
customs := make([]RegionSummary, 0)
for _, summary := range summaries {
if inferredType := inferTypeFromDescription(summary.Description); inferredType != "" {
if _, exists := index[inferredType]; !exists {
inferred := summary
inferred.Type = inferredType
inferred.Name = inferred.Name + "(按描述推断)"
summaries = append(summaries, inferred)
index[inferredType] = inferred
warnings = append(warnings, fmt.Sprintf("区域 %q 已按描述推断为 %s", summary.Name, inferredType))
}
}
if summary.Type == "custom" {
customs = append(customs, summary)
}
}
if len(customs) == 0 {
return summaries, warnings
}
if _, exists := index["contact_list"]; !exists {
best := customs[0]
for _, candidate := range customs[1:] {
if candidate.CenterScreen[0] < best.CenterScreen[0] && boxHeight(candidate.BBoxScreen) > boxHeight(best.BBoxScreen)*0.5 {
best = candidate
}
}
best.Type = "contact_list"
best.Name = best.Name + "(推断联系人区域)"
summaries = append(summaries, best)
index[best.Type] = best
warnings = append(warnings, "未标注 contact_list已按最左侧大区域推断")
}
if _, exists := index["input_box"]; !exists {
best := customs[0]
for _, candidate := range customs[1:] {
if candidate.CenterScreen[1] > best.CenterScreen[1] || nearlyEqual(candidate.CenterScreen[1], best.CenterScreen[1]) && boxArea(candidate.BBoxScreen) > boxArea(best.BBoxScreen) {
best = candidate
}
}
best.Type = "input_box"
best.Name = best.Name + "(推断输入框区域)"
summaries = append(summaries, best)
index[best.Type] = best
warnings = append(warnings, "未标注 input_box已按底部区域推断")
}
if _, exists := index["chat_content"]; !exists {
input := index["input_box"]
contact := index["contact_list"]
best := customs[0]
bestScore := math.Inf(-1)
for _, candidate := range customs {
if sameBox(candidate.BBoxScreen, input.BBoxScreen) || sameBox(candidate.BBoxScreen, contact.BBoxScreen) {
continue
}
score := boxArea(candidate.BBoxScreen) - math.Abs(candidate.CenterScreen[1]-input.CenterScreen[1])*0.2
if score > bestScore {
best = candidate
bestScore = score
}
}
best.Type = "chat_content"
best.Name = best.Name + "(推断聊天内容区域)"
summaries = append(summaries, best)
index[best.Type] = best
warnings = append(warnings, "未标注 chat_content已按大面积内容区推断")
}
return summaries, warnings
}
func sameBox(a, b [4]float64) bool {
return nearlyEqual(a[0], b[0]) && nearlyEqual(a[1], b[1]) && nearlyEqual(a[2], b[2]) && nearlyEqual(a[3], b[3])
}
func isZeroBox(box [4]float64) bool {
return box[0] == 0 && box[1] == 0 && box[2] == 0 && box[3] == 0
}
func boxWidth(box [4]float64) float64 {
return math.Abs(box[2] - box[0])
}
func boxHeight(box [4]float64) float64 {
return math.Abs(box[3] - box[1])
}
func boxArea(box [4]float64) float64 {
return boxWidth(box) * boxHeight(box)
}
func nearlyEqual(a, b float64) bool {
return math.Abs(a-b) < 8
}
func inferTypeFromDescription(description string) string {
text := strings.ToLower(description)
if text == "" {
return ""
}
if strings.Contains(text, "chat_content") || strings.Contains(text, "聊天") || strings.Contains(text, "消息") || strings.Contains(text, "记录") {
return "chat_content"
}
if strings.Contains(text, "input_box") || strings.Contains(text, "输入") || strings.Contains(text, "文本框") || strings.Contains(text, "编辑框") {
return "input_box"
}
if strings.Contains(text, "send_button") || strings.Contains(text, "发送") || strings.Contains(text, "send") {
return "send_button"
}
if strings.Contains(text, "contact_list") || strings.Contains(text, "联系人") || strings.Contains(text, "会话列表") || strings.Contains(text, "好友") {
return "contact_list"
}
return ""
}