添加了通过识别图片导入课表的模块

This commit is contained in:
WuYuuuub
2026-03-15 21:49:10 +08:00
parent eebf239528
commit 9addc89a66
2 changed files with 341 additions and 0 deletions

44
main.go
View File

@@ -3,6 +3,7 @@ package main
import ( import (
"context" "context"
"flag" "flag"
"io"
"log" "log"
"net/http" "net/http"
"os" "os"
@@ -16,6 +17,7 @@ import (
"jwts/config" "jwts/config"
grpcClient "jwts/grpc" grpcClient "jwts/grpc"
"jwts/parser" "jwts/parser"
"jwts/recognizer"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
) )
@@ -93,6 +95,48 @@ func runHTTPServer() {
}) })
}) })
// 课表图片识别接口
r.POST("/recognize", func(c *gin.Context) {
// 读取上传的图片文件
file, header, err := c.Request.FormFile("image")
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "请上传图片文件", "details": err.Error()})
return
}
defer file.Close()
log.Printf("收到图片上传: %s, 大小: %d bytes", header.Filename, header.Size)
// 读取图片数据
imageData, err := io.ReadAll(file)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "读取图片失败", "details": err.Error()})
return
}
// 调用识别模块
courses, err := recognizer.RecognizeFromBytes(imageData)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "识别失败", "details": err.Error()})
return
}
// 打印识别结果
recognizer.PrintCourses(courses)
// 保存到JSON文件
err = recognizer.SaveToJSON(courses, "schedule_result.json")
if err != nil {
log.Printf("保存JSON失败: %v", err)
}
c.JSON(http.StatusOK, gin.H{
"message": "识别成功",
"courses": courses,
"count": len(courses),
})
})
// 任务处理接口 // 任务处理接口
r.POST("/task", func(c *gin.Context) { r.POST("/task", func(c *gin.Context) {
// 这里可以根据请求体参数决定执行什么任务 // 这里可以根据请求体参数决定执行什么任务

297
recognizer/recognizer.go Normal file
View File

@@ -0,0 +1,297 @@
package recognizer
import (
"bytes"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"os"
"regexp"
"strings"
"time"
"jwts/models"
)
// 配置 - 使用您提供的API配置
const (
APIURL = "https://api.littlelan.cn/v1/chat/completions"
APIKey = "sk-Ll7P7vuiLum3N0pj7XaYfwHhD5UW1rDYlFoQ7F5JmeBfTewZ"
Model = "qwen3.5-plus" // 使用qwen3.5-plus模型
)
// 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,
}
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) {
// 构建请求 - 发送图片给AI识别
// 禁用qwen3.5的思考模式避免产生大量不必要的token消耗
falseVal := false
zero := 0
apiReq := APIRequest{
Model: Model,
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", APIURL, bytes.NewBuffer(jsonData))
if err != nil {
return nil, err
}
httpReq.Header = http.Header{
"Content-Type": []string{"application/json"},
"Authorization": []string{"Bearer " + APIKey},
}
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==================================================")
}