119 lines
3.8 KiB
Go
119 lines
3.8 KiB
Go
package formula
|
|
|
|
import (
|
|
"net/http"
|
|
"path/filepath"
|
|
|
|
"gitea.com/bitwsd/document_ai/internal/model/formula"
|
|
"gitea.com/bitwsd/document_ai/internal/service"
|
|
"gitea.com/bitwsd/document_ai/internal/storage/dao"
|
|
"gitea.com/bitwsd/document_ai/pkg/common"
|
|
"gitea.com/bitwsd/document_ai/pkg/utils"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
type FormulaEndpoint struct {
|
|
recognitionService *service.RecognitionService
|
|
}
|
|
|
|
func NewFormulaEndpoint() *FormulaEndpoint {
|
|
return &FormulaEndpoint{
|
|
recognitionService: service.NewRecognitionService(),
|
|
}
|
|
}
|
|
|
|
// CreateTask godoc
|
|
// @Summary Create a formula recognition task
|
|
// @Description Create a new formula recognition task from image
|
|
// @Tags Formula
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Param request body CreateFormulaRecognitionRequest true "Create task request"
|
|
// @Success 200 {object} common.Response{data=CreateTaskResponse}
|
|
// @Failure 400 {object} common.Response
|
|
// @Failure 500 {object} common.Response
|
|
// @Router /v1/formula/recognition [post]
|
|
func (endpoint *FormulaEndpoint) CreateTask(ctx *gin.Context) {
|
|
var req formula.CreateFormulaRecognitionRequest
|
|
if err := ctx.BindJSON(&req); err != nil {
|
|
ctx.JSON(http.StatusOK, common.ErrorResponse(ctx, common.CodeParamError, "Invalid parameters"))
|
|
return
|
|
}
|
|
|
|
if !utils.InArray(req.TaskType, []string{string(dao.TaskTypeFormula), string(dao.TaskTypeFormula)}) {
|
|
ctx.JSON(http.StatusOK, common.ErrorResponse(ctx, common.CodeParamError, "Invalid task type"))
|
|
return
|
|
}
|
|
|
|
fileExt := filepath.Ext(req.FileName)
|
|
if !utils.InArray(fileExt, []string{".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"}) {
|
|
ctx.JSON(http.StatusOK, common.ErrorResponse(ctx, common.CodeParamError, "Invalid file type"))
|
|
return
|
|
}
|
|
|
|
task, err := endpoint.recognitionService.CreateRecognitionTask(ctx, &req)
|
|
if err != nil {
|
|
ctx.JSON(http.StatusOK, common.ErrorResponse(ctx, common.CodeSystemError, "Failed to create task"))
|
|
return
|
|
}
|
|
|
|
ctx.JSON(http.StatusOK, common.SuccessResponse(ctx, &formula.CreateTaskResponse{
|
|
TaskNo: task.TaskUUID,
|
|
Status: int(task.Status),
|
|
}))
|
|
}
|
|
|
|
// GetTaskStatus godoc
|
|
// @Summary Get formula recognition task status
|
|
// @Description Get the status and results of a formula recognition task
|
|
// @Tags Formula
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Param task_no path string true "Task No"
|
|
// @Success 200 {object} common.Response{data=GetTaskStatusResponse}
|
|
// @Failure 400 {object} common.Response
|
|
// @Failure 404 {object} common.Response
|
|
// @Failure 500 {object} common.Response
|
|
// @Router /v1/formula/recognition/{task_no} [get]
|
|
func (endpoint *FormulaEndpoint) GetTaskStatus(c *gin.Context) {
|
|
var req formula.GetRecognitionStatusRequest
|
|
if err := c.ShouldBindUri(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, common.ErrorResponse(c, common.CodeParamError, "invalid task no"))
|
|
return
|
|
}
|
|
|
|
task, err := endpoint.recognitionService.GetFormualTask(c, req.TaskNo)
|
|
if err != nil {
|
|
c.JSON(http.StatusOK, common.ErrorResponse(c, common.CodeSystemError, "failed to get task status"))
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, common.SuccessResponse(c, task))
|
|
}
|
|
|
|
// AI增强识别
|
|
// @Summary AI增强识别
|
|
// @Description AI增强识别
|
|
// @Tags Formula
|
|
// @Accept json
|
|
// @Produce json
|
|
// @Param request body formula.AIEnhanceRecognitionRequest true "AI增强识别请求"
|
|
// @Success 200 {object} common.Response{data=formula.AIEnhanceRecognitionResponse}
|
|
// @Router /v1/formula/ai_enhance [post]
|
|
func (endpoint *FormulaEndpoint) AIEnhanceRecognition(c *gin.Context) {
|
|
var req formula.AIEnhanceRecognitionRequest
|
|
if err := c.BindJSON(&req); err != nil {
|
|
c.JSON(http.StatusOK, common.ErrorResponse(c, common.CodeParamError, "Invalid parameters"))
|
|
return
|
|
}
|
|
|
|
_, err := endpoint.recognitionService.AIEnhanceRecognition(c, &req)
|
|
if err != nil {
|
|
c.JSON(http.StatusOK, common.ErrorResponse(c, common.CodeSystemError, err.Error()))
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, common.SuccessResponse(c, nil))
|
|
}
|