205 lines
5.9 KiB
Go
205 lines
5.9 KiB
Go
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 ""
|
||
}
|