85 lines
2.1 KiB
Go
85 lines
2.1 KiB
Go
package handler
|
||
|
||
import (
|
||
"net/http"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"tencent_ocr/pkg/service"
|
||
"tencent_ocr/pkg/errors"
|
||
)
|
||
|
||
type OCRHandler struct {
|
||
ocrService *service.OCRService
|
||
geminiService *service.GeminiService
|
||
}
|
||
|
||
func NewOCRHandler(ocrService *service.OCRService, geminiService *service.GeminiService) *OCRHandler {
|
||
return &OCRHandler{
|
||
ocrService: ocrService,
|
||
geminiService: geminiService,
|
||
}
|
||
}
|
||
|
||
type OCRRequest struct {
|
||
ImageBase64 string `json:"image_base64"`
|
||
ImageURL string `json:"image_url"`
|
||
Scene string `json:"scene"`
|
||
}
|
||
|
||
type OCRResponse struct {
|
||
OriginalText string `json:"original_text"`
|
||
Result string `json:"result"`
|
||
Success bool `json:"success"`
|
||
Error string `json:"error,omitempty"`
|
||
}
|
||
|
||
func (h *OCRHandler) HandleOCR(c *gin.Context) {
|
||
var req OCRRequest
|
||
if err := c.ShouldBindJSON(&req); err != nil {
|
||
c.JSON(http.StatusBadRequest, OCRResponse{
|
||
Success: false,
|
||
Error: "Invalid request format",
|
||
})
|
||
return
|
||
}
|
||
|
||
// Validate that at least one of ImageURL or ImageBase64 is provided
|
||
if req.ImageURL == "" && req.ImageBase64 == "" {
|
||
c.JSON(http.StatusBadRequest, OCRResponse{
|
||
Success: false,
|
||
Error: "Either image_url or image_base64 must be provided",
|
||
})
|
||
return
|
||
}
|
||
|
||
// Process image
|
||
ocrText, err := h.ocrService.ProcessImage(c.Request.Context(), req.ImageBase64)
|
||
if err != nil {
|
||
status := http.StatusInternalServerError
|
||
if errors.IsClientError(err) {
|
||
status = http.StatusBadRequest
|
||
}
|
||
c.JSON(status, OCRResponse{
|
||
Success: false,
|
||
Error: err.Error(),
|
||
})
|
||
return
|
||
}
|
||
|
||
// Process with Gemini
|
||
prompt := "你是一个专业的助手，负责纠正OCR识别结果中的文本。只需要输出识别结果，不需要输出任何解释。\n\n" + ocrText
|
||
processedText, err := h.geminiService.ProcessText(c.Request.Context(), ocrText)
|
||
if err != nil {
|
||
c.JSON(http.StatusInternalServerError, OCRResponse{
|
||
Success: false,
|
||
Error: "Text processing failed: " + err.Error(),
|
||
})
|
||
return
|
||
}
|
||
|
||
c.JSON(http.StatusOK, OCRResponse{
|
||
Success: true,
|
||
OriginalText: ocrText,
|
||
Result: processedText,
|
||
})
|
||
} |