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==================================================") }