From 294151cd409bb9e73c2b7a4cf06313ee35020026 Mon Sep 17 00:00:00 2001 From: lafay <2021211506@stu.hit.edu.cn> Date: Mon, 16 Mar 2026 13:15:56 +0800 Subject: [PATCH] =?UTF-8?q?refactor:=20=E5=B0=86recognizer=E6=A8=A1?= =?UTF-8?q?=E5=9D=97=E6=8B=86=E5=88=86=E5=88=B0=E7=8E=B0=E6=9C=89=E6=A8=A1?= =?UTF-8?q?=E5=9D=97=E4=B8=AD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 将课表识别功能合并到captcha模块 - 将RecognizerClient添加到client模块 - 将输出函数移至models/output.go - 更新API结构体添加扩展字段 - 删除recognizer目录 --- captcha/recognize.go | 196 +++++++++++++++++++++++++ client/http.go | 9 +- main.go | 60 +++++--- models/api.go | 8 +- models/output.go | 39 +++++ recognizer/recognizer.go | 298 --------------------------------------- 6 files changed, 287 insertions(+), 323 deletions(-) create mode 100644 models/output.go delete mode 100644 recognizer/recognizer.go diff --git a/captcha/recognize.go b/captcha/recognize.go index 0546397..859a8b9 100644 --- a/captcha/recognize.go +++ b/captcha/recognize.go @@ -1,14 +1,21 @@ package captcha import ( + "bytes" + "encoding/base64" + "encoding/json" "fmt" "io" + "log" "net/http" "os" + "regexp" "strings" + "time" "jwts/client" "jwts/config" + "jwts/models" "github.com/PuerkitoBio/goquery" ) @@ -70,3 +77,192 @@ func GetAndRecognize(client *http.Client) (string, error) { func ParseHTML(html string) (*goquery.Document, error) { return goquery.NewDocumentFromReader(strings.NewReader(html)) } + +// RecognizeScheduleFromFile 从图片文件识别课表 +func RecognizeScheduleFromFile(imagePath string) ([]models.Course, error) { + fmt.Printf("[识别] 正在读取图片: %s\n", imagePath) + + imageData, err := os.ReadFile(imagePath) + if err != nil { + return nil, fmt.Errorf("读取图片失败: %v", err) + } + + imageBase64 := base64.StdEncoding.EncodeToString(imageData) + fmt.Printf("[识别] 图片大小: %d bytes\n", len(imageData)) + + return callScheduleAPI(imageBase64) +} + +// RecognizeScheduleFromBytes 从图片数据识别课表 +func RecognizeScheduleFromBytes(imageData []byte) ([]models.Course, error) { + fmt.Printf("[识别] 正在处理图片数据, 大小: %d bytes\n", len(imageData)) + + imageBase64 := base64.StdEncoding.EncodeToString(imageData) + + return callScheduleAPI(imageBase64) +} + +// RecognizeScheduleFromBase64 从base64数据识别课表 +func RecognizeScheduleFromBase64(imageBase64 string) ([]models.Course, error) { + fmt.Printf("[识别] 正在处理base64图片数据, 长度: %d\n", len(imageBase64)) + + return callScheduleAPI(imageBase64) +} + +// callScheduleAPI 调用API识别课表（带重试机制） +func callScheduleAPI(imageBase64 string) ([]models.Course, error) { + maxRetries := 3 + for retry := 0; retry < maxRetries; retry++ { + if retry > 0 { + fmt.Printf("[识别] 第 %d 次重试...\n", retry+1) + time.Sleep(5 * time.Second) + } + + courses, err := callScheduleAPISingle(imageBase64) + if err == nil { + return courses, nil + } + + fmt.Printf("[识别] 尝试 %d/%d 失败: %v\n", retry+1, maxRetries, err) + + if !strings.Contains(err.Error(), "504") { + return nil, err + } + } + + return nil, fmt.Errorf("已达到最大重试次数") +} + +// callScheduleAPISingle 单次API调用 +func callScheduleAPISingle(imageBase64 string) ([]models.Course, error) { + if config.RecognizerAPIKey == "" { + return nil, fmt.Errorf("RECOGNIZER_API_KEY 未配置") + } + + falseVal := false + zero := 0 + apiReq := models.APIRequest{ + Model: config.RecognizerModel, + MaxTokens: 4000, + EnableThinking: &falseVal, + ThinkingBudget: &zero, + Messages: []models.Message{ + { + Role: "user", + Content: []models.ContentItem{ + { + Type: "text", + Text: `请解析以下课表图片，提取所有课程信息并以JSON数组格式返回。 + +JSON格式要求： +[{"name":"课程名","teacher":"老师名","room":"教室","schedule":[{"day":"星期X","section":[节次数组],"weeks":[周数数组]}]}] + +注意事项： +1. day使用"星期一"到"星期日" +2. section节次：上午1-2节为[1,2]，上午3-4节为[3,4]，下午5-6节为[5,6]，下午7-8节为[7,8]，晚上9-10节为[9,10]，晚上11-12节为[11,12] +3. weeks周数：格式如[1,2,3,4,5,6,7,8]或[9,10,11,12,13,14,15,16] +4. 只返回JSON数组，不要任何其他文字。`, + }, + { + Type: "image_url", + ImageURL: &models.ImageURL{ + URL: "data:image/png;base64," + imageBase64, + }, + }, + }, + }, + }, + } + + jsonData, err := json.Marshal(apiReq) + if err != nil { + return nil, err + } + + httpReq, err := http.NewRequest("POST", config.RecognizerAPIURL, bytes.NewBuffer(jsonData)) + if err != nil { + return nil, err + } + + httpReq.Header = http.Header{ + "Content-Type": []string{"application/json"}, + "Authorization": []string{"Bearer " + config.RecognizerAPIKey}, + } + + fmt.Println("[识别] 正在调用API进行识别...") + resp, err := client.RecognizerClient.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("API调用失败: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != 200 { + respBody, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("API调用失败: %d, %s", resp.StatusCode, string(respBody)) + } + + var apiResp models.APIResponse + err = json.NewDecoder(resp.Body).Decode(&apiResp) + if err != nil { + return nil, fmt.Errorf("API响应解析失败: %v", err) + } + + if len(apiResp.Choices) == 0 { + return nil, fmt.Errorf("API返回为空") + } + + var resultText string + switch c := apiResp.Choices[0].Message.Content.(type) { + case string: + resultText = c + case []interface{}: + for _, item := range c { + if itemMap, ok := item.(map[string]interface{}); ok { + if text, ok := itemMap["text"].(string); ok { + resultText = text + break + } + } + } + } + + fmt.Println("[识别] API响应已接收，正在解析...") + fmt.Printf("[识别] 原始响应: %s\n", resultText) + + var jsonMatch string + + jsonBlockMatch := regexp.MustCompile("(?s)```json\\s*(\\[[\\s\\S]*?\\])\\s*```").FindStringSubmatch(resultText) + if len(jsonBlockMatch) > 1 { + jsonMatch = jsonBlockMatch[1] + fmt.Println("[识别] 从json代码块中提取JSON") + } else { + jsonBlockMatch = regexp.MustCompile("(?s)```\\s*(\\[[\\s\\S]*?\\])\\s*```").FindStringSubmatch(resultText) + if len(jsonBlockMatch) > 1 { + jsonMatch = jsonBlockMatch[1] + fmt.Println("[识别] 从代码块中提取JSON") + } else { + jsonMatch = regexp.MustCompile(`(\[[\s\S]*\])`).FindString(resultText) + if jsonMatch == "" { + return nil, fmt.Errorf("未找到JSON数据，原始响应: %s", resultText) + } + fmt.Println("[识别] 从文本中提取JSON") + } + } + + var courses []models.Course + err = json.Unmarshal([]byte(jsonMatch), &courses) + if err != nil { + return nil, fmt.Errorf("JSON解析失败: %v, 数据: %s", err, jsonMatch) + } + + fmt.Printf("[识别] 成功解析 %d 门课程\n", len(courses)) + return courses, nil +} + +func init() { + if config.RecognizerAPIKey == "" { + log.Println("警告: RECOGNIZER_API_KEY 未设置，图片识别功能将无法使用") + } + + log.Println("captcha模块已初始化") +} diff --git a/client/http.go b/client/http.go index 5a15ebc..5fbafad 100644 --- a/client/http.go +++ b/client/http.go @@ -7,14 +7,15 @@ import ( "net/http/cookiejar" "net/url" "strings" + "time" "jwts/config" ) -// 全局HTTP客户端 var Client *http.Client -// 初始化HTTP客户端 +var RecognizerClient *http.Client + func InitClient() { jar, err := cookiejar.New(nil) if err != nil { @@ -23,6 +24,10 @@ func InitClient() { Client = &http.Client{ Jar: jar, } + + RecognizerClient = &http.Client{ + Timeout: 120 * time.Second, + } } // 设置请求头 diff --git a/main.go b/main.go index 0a4ff33..1ec9770 100644 --- a/main.go +++ b/main.go @@ -16,8 +16,8 @@ import ( "jwts/client" "jwts/config" grpcClient "jwts/grpc" + "jwts/models" "jwts/parser" - "jwts/recognizer" "github.com/gin-gonic/gin" ) @@ -49,35 +49,55 @@ func runGRPCClient(serverAddr string) { log.Printf("启动 gRPC 客户端, Runner ID: %s, 版本: %s", runnerID, version) - client := grpcClient.NewClient(serverAddr, runnerID, version) - handler := grpcClient.NewTaskHandler() - client.SetHandler(handler) - ctx, cancel := context.WithCancel(context.Background()) defer cancel() - // 处理信号 sigChan := make(chan os.Signal, 1) signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) - if err := client.Connect(ctx); err != nil { - log.Printf("连接失败: %v", err) - os.Exit(1) - } - go func() { <-sigChan log.Println("正在关闭...") - client.Close() cancel() }() - // 启动心跳 - go client.HeartbeatLoop(ctx, 30*time.Second) + reconnectInterval := 5 * time.Second - // 运行主循环 - if err := client.Run(ctx); err != nil { - log.Printf("客户端错误: %v", err) + for { + select { + case <-ctx.Done(): + log.Println("客户端已停止") + return + default: + } + + client := grpcClient.NewClient(serverAddr, runnerID, version) + handler := grpcClient.NewTaskHandler() + client.SetHandler(handler) + + if err := client.Connect(ctx); err != nil { + log.Printf("连接失败: %v, %v 后重试...", err, reconnectInterval) + time.Sleep(reconnectInterval) + continue + } + + heartbeatCtx, heartbeatCancel := context.WithCancel(ctx) + go client.HeartbeatLoop(heartbeatCtx, 30*time.Second) + + err := client.Run(ctx) + heartbeatCancel() + client.Close() + + if ctx.Err() != nil { + return + } + + if err != nil { + log.Printf("连接断开: %v, %v 后重连...", err, reconnectInterval) + } else { + log.Printf("连接已关闭, %v 后重连...", reconnectInterval) + } + time.Sleep(reconnectInterval) } } @@ -117,15 +137,15 @@ func runHTTPServer() { return } - courses, err := recognizer.RecognizeFromBytes(imageData) + courses, err := captcha.RecognizeScheduleFromBytes(imageData) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": "识别失败", "details": err.Error()}) return } - recognizer.PrintCourses(courses) + models.PrintCourses(courses) - err = recognizer.SaveToJSON(courses, "schedule_result.json") + err = models.SaveToJSON(courses, "schedule_result.json") if err != nil { log.Printf("保存JSON失败: %v", err) } diff --git a/models/api.go b/models/api.go index d2c2b19..bc9af61 100644 --- a/models/api.go +++ b/models/api.go @@ -2,9 +2,11 @@ package models // APIRequest API请求结构 type APIRequest struct { - Model string `json:"model"` - Messages []Message `json:"messages"` - MaxTokens int `json:"max_tokens"` + Model string `json:"model"` + Messages []Message `json:"messages"` + MaxTokens int `json:"max_tokens"` + EnableThinking *bool `json:"enable_thinking,omitempty"` + ThinkingBudget *int `json:"thinking_budget,omitempty"` } // Message 消息结构 diff --git a/models/output.go b/models/output.go new file mode 100644 index 0000000..b7d719d --- /dev/null +++ b/models/output.go @@ -0,0 +1,39 @@ +package models + +import ( + "encoding/json" + "fmt" + "os" +) + +func SaveToJSON(courses []Course, filename string) error { + jsonData, err := json.MarshalIndent(courses, "", " ") + if err != nil { + return err + } + + err = os.WriteFile(filename, jsonData, 0644) + if err != nil { + return err + } + + fmt.Printf("[输出] JSON结果已保存到: %s\n", filename) + return nil +} + +func PrintCourses(courses []Course) { + fmt.Println("\n==================================================") + fmt.Println("识别结果：") + fmt.Println("==================================================") + for i, course := range courses { + fmt.Printf("\n课程 %d: %s\n", i+1, course.Name) + fmt.Printf(" 教师: %s\n", course.Teacher) + fmt.Printf(" 教室: %s\n", course.Room) + fmt.Printf(" 时间:") + for _, sch := range course.Schedule { + fmt.Printf(" %s %v 周%v", sch.Day, sch.Section, sch.Weeks) + } + fmt.Println() + } + fmt.Println("\n==================================================") +} diff --git a/recognizer/recognizer.go b/recognizer/recognizer.go deleted file mode 100644 index c7b1bde..0000000 --- a/recognizer/recognizer.go +++ /dev/null @@ -1,298 +0,0 @@ -package recognizer - -import ( - "bytes" - "encoding/base64" - "encoding/json" - "fmt" - "io" - "log" - "net/http" - "os" - "regexp" - "strings" - "time" - - "jwts/config" - "jwts/models" -) - -// APIRequest API请求结构 -type APIRequest struct { - Model string `json:"model"` - Messages []Message `json:"messages"` - MaxTokens int `json:"max_tokens"` - EnableThinking *bool `json:"enable_thinking,omitempty"` - ThinkingBudget *int `json:"thinking_budget,omitempty"` -} - -// Message 消息结构 -type Message struct { - Role string `json:"role"` - Content interface{} `json:"content"` -} - -// ContentItem 内容项 -type ContentItem struct { - Type string `json:"type,omitempty"` - Text string `json:"text,omitempty"` - ImageURL *ImageURL `json:"image_url,omitempty"` -} - -// ImageURL 图片URL -type ImageURL struct { - URL string `json:"url"` -} - -// APIResponse API响应结构 -type APIResponse struct { - Choices []Choice `json:"choices"` -} - -// Choice 选择 -type Choice struct { - Message Message `json:"message"` -} - -// 全局HTTP客户端 -var client *http.Client - -// 初始化HTTP客户端 -func init() { - client = &http.Client{ - Timeout: 120 * time.Second, - } - - if config.RecognizerAPIKey == "" { - log.Println("警告: RECOGNIZER_API_KEY 未设置，图片识别功能将无法使用") - } - - log.Println("recognizer模块已初始化") -} - -// RecognizeFromFile 从图片文件识别课表 -func RecognizeFromFile(imagePath string) ([]models.Course, error) { - fmt.Printf("[识别] 正在读取图片: %s\n", imagePath) - - // 读取图片文件 - imageData, err := os.ReadFile(imagePath) - if err != nil { - return nil, fmt.Errorf("读取图片失败: %v", err) - } - - // 将图片转为base64 - imageBase64 := base64.StdEncoding.EncodeToString(imageData) - fmt.Printf("[识别] 图片大小: %d bytes\n", len(imageData)) - - return callAPI(imageBase64) -} - -// RecognizeFromBytes 从图片数据识别课表 -func RecognizeFromBytes(imageData []byte) ([]models.Course, error) { - fmt.Printf("[识别] 正在处理图片数据, 大小: %d bytes\n", len(imageData)) - - // 将图片转为base64 - imageBase64 := base64.StdEncoding.EncodeToString(imageData) - - return callAPI(imageBase64) -} - -// RecognizeFromBase64 从base64数据识别课表 -func RecognizeFromBase64(imageBase64 string) ([]models.Course, error) { - fmt.Printf("[识别] 正在处理base64图片数据, 长度: %d\n", len(imageBase64)) - - return callAPI(imageBase64) -} - -// callAPI 调用API识别课表（带重试机制） -func callAPI(imageBase64 string) ([]models.Course, error) { - // 重试机制 - maxRetries := 3 - for retry := 0; retry < maxRetries; retry++ { - if retry > 0 { - fmt.Printf("[识别] 第 %d 次重试...\n", retry+1) - time.Sleep(5 * time.Second) // 等待5秒后重试 - } - - courses, err := callAPISingle(imageBase64) - if err == nil { - return courses, nil - } - - fmt.Printf("[识别] 尝试 %d/%d 失败: %v\n", retry+1, maxRetries, err) - - // 如果是504超时，继续重试；否则直接返回错误 - if !strings.Contains(err.Error(), "504") { - return nil, err - } - } - - return nil, fmt.Errorf("已达到最大重试次数") -} - -// callAPISingle 单次API调用 -func callAPISingle(imageBase64 string) ([]models.Course, error) { - if config.RecognizerAPIKey == "" { - return nil, fmt.Errorf("RECOGNIZER_API_KEY 未配置") - } - - falseVal := false - zero := 0 - apiReq := APIRequest{ - Model: config.RecognizerModel, - MaxTokens: 4000, - EnableThinking: &falseVal, - ThinkingBudget: &zero, - Messages: []Message{ - { - Role: "user", - Content: []ContentItem{ - { - Type: "text", - Text: `请解析以下课表图片，提取所有课程信息并以JSON数组格式返回。 - -JSON格式要求： -[{"name":"课程名","teacher":"老师名","room":"教室","schedule":[{"day":"星期X","section":[节次数组],"weeks":[周数数组]}]}] - -注意事项： -1. day使用"星期一"到"星期日" -2. section节次：上午1-2节为[1,2]，上午3-4节为[3,4]，下午5-6节为[5,6]，下午7-8节为[7,8]，晚上9-10节为[9,10]，晚上11-12节为[11,12] -3. weeks周数：格式如[1,2,3,4,5,6,7,8]或[9,10,11,12,13,14,15,16] -4. 只返回JSON数组，不要任何其他文字。`, - }, - { - Type: "image_url", - ImageURL: &ImageURL{ - URL: "data:image/png;base64," + imageBase64, - }, - }, - }, - }, - }, - } - - jsonData, err := json.Marshal(apiReq) - if err != nil { - return nil, err - } - - httpReq, err := http.NewRequest("POST", config.RecognizerAPIURL, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, err - } - - httpReq.Header = http.Header{ - "Content-Type": []string{"application/json"}, - "Authorization": []string{"Bearer " + config.RecognizerAPIKey}, - } - - fmt.Println("[识别] 正在调用API进行识别...") - resp, err := client.Do(httpReq) - if err != nil { - return nil, fmt.Errorf("API调用失败: %v", err) - } - defer resp.Body.Close() - - if resp.StatusCode != 200 { - respBody, _ := io.ReadAll(resp.Body) - return nil, fmt.Errorf("API调用失败: %d, %s", resp.StatusCode, string(respBody)) - } - - var apiResp APIResponse - err = json.NewDecoder(resp.Body).Decode(&apiResp) - if err != nil { - return nil, fmt.Errorf("API响应解析失败: %v", err) - } - - if len(apiResp.Choices) == 0 { - return nil, fmt.Errorf("API返回为空") - } - - // 处理Content，可能是string或[]ContentItem - var resultText string - switch c := apiResp.Choices[0].Message.Content.(type) { - case string: - resultText = c - case []interface{}: - for _, item := range c { - if itemMap, ok := item.(map[string]interface{}); ok { - if text, ok := itemMap["text"].(string); ok { - resultText = text - break - } - } - } - } - - fmt.Println("[识别] API响应已接收，正在解析...") - fmt.Printf("[识别] 原始响应: %s\n", resultText) - - // 提取JSON部分 - var jsonMatch string - - // 1. 先尝试匹配 ```json ... ``` 代码块 - jsonBlockMatch := regexp.MustCompile("(?s)```json\\s*(\\[[\\s\\S]*?\\])\\s*```").FindStringSubmatch(resultText) - if len(jsonBlockMatch) > 1 { - jsonMatch = jsonBlockMatch[1] - fmt.Println("[识别] 从json代码块中提取JSON") - } else { - // 2. 尝试匹配 ``` ... ``` 代码块中的数组 - jsonBlockMatch = regexp.MustCompile("(?s)```\\s*(\\[[\\s\\S]*?\\])\\s*```").FindStringSubmatch(resultText) - if len(jsonBlockMatch) > 1 { - jsonMatch = jsonBlockMatch[1] - fmt.Println("[识别] 从代码块中提取JSON") - } else { - // 3. 查找最外层的JSON数组（匹配第一个 [ 到最后一个 ]） - // 使用非贪婪匹配找到完整的JSON数组 - jsonMatch = regexp.MustCompile(`(\[[\s\S]*\])`).FindString(resultText) - if jsonMatch == "" { - return nil, fmt.Errorf("未找到JSON数据，原始响应: %s", resultText) - } - fmt.Println("[识别] 从文本中提取JSON") - } - } - - // 解析JSON - var courses []models.Course - err = json.Unmarshal([]byte(jsonMatch), &courses) - if err != nil { - return nil, fmt.Errorf("JSON解析失败: %v, 数据: %s", err, jsonMatch) - } - - fmt.Printf("[识别] 成功解析 %d 门课程\n", len(courses)) - return courses, nil -} - -// SaveToJSON 保存课程数据到JSON文件 -func SaveToJSON(courses []models.Course, filename string) error { - jsonData, err := json.MarshalIndent(courses, "", " ") - if err != nil { - return err - } - - err = os.WriteFile(filename, jsonData, 0644) - if err != nil { - return err - } - - fmt.Printf("[输出] JSON结果已保存到: %s\n", filename) - return nil -} - -// PrintCourses 打印课程信息 -func PrintCourses(courses []models.Course) { - fmt.Println("\n==================================================") - fmt.Println("识别结果：") - fmt.Println("==================================================") - for i, course := range courses { - fmt.Printf("\n课程 %d: %s\n", i+1, course.Name) - fmt.Printf(" 教师: %s\n", course.Teacher) - fmt.Printf(" 教室: %s\n", course.Room) - fmt.Printf(" 时间:") - for _, sch := range course.Schedule { - fmt.Printf(" %s %v 周%v", sch.Day, sch.Section, sch.Weeks) - } - fmt.Println() - } - fmt.Println("\n==================================================") -}