294151cd40
- 将课表识别功能合并到captcha模块 - 将RecognizerClient添加到client模块 - 将输出函数移至models/output.go - 更新API结构体添加扩展字段 - 删除recognizer目录
269 lines
7.0 KiB
Go
269 lines
7.0 KiB
Go
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"
|
||
)
|
||
|
||
// 保存验证码图片
|
||
func SaveCaptchaImage(httpClient *http.Client) ([]byte, error) {
|
||
fmt.Println("[2] 获取验证码...")
|
||
|
||
req, err := http.NewRequest("GET", config.BaseURL+"/captchaImage", nil)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
req.Header = client.GetHeaders()
|
||
|
||
resp, err := httpClient.Do(req)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
body, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
if err := os.WriteFile("captcha.png", body, 0644); err != nil {
|
||
fmt.Printf(" 警告: 无法保存验证码图片: %v\n", err)
|
||
} else {
|
||
fmt.Println(" 验证码已保存到: captcha.png")
|
||
}
|
||
|
||
return body, nil
|
||
}
|
||
|
||
// 识别验证码 - 使用固定验证码用于测试
|
||
func Recognize(imagePath string) (string, error) {
|
||
fmt.Println("[3] 使用固定验证码...")
|
||
// 固定验证码用于测试
|
||
return "....", nil
|
||
}
|
||
|
||
// 获取并识别验证码
|
||
func GetAndRecognize(client *http.Client) (string, error) {
|
||
_, err := SaveCaptchaImage(client)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
captchaCode, err := Recognize("captcha.png")
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
fmt.Printf(" 识别结果: %s\n", captchaCode)
|
||
return captchaCode, nil
|
||
}
|
||
|
||
// 用于解析HTML（登录页面）
|
||
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模块已初始化")
|
||
}
|