ocr worker

This commit is contained in:
m
2026-07-12 14:35:30 +02:00
parent 911592758a
commit 67096c9482
6 changed files with 73 additions and 13 deletions
+4 -1
View File
@@ -27,9 +27,12 @@ func main() {
queries := db.New(database) queries := db.New(database)
fileSvc := service.NewFileService(queries, cfg) fileSvc := service.NewFileService(queries, cfg)
ocrSvc := service.NewOCRService(cfg) ocrSvc := service.NewOCRService(cfg, fileSvc)
urlSvc := service.NewURLService(cfg.HMACSecret, cfg.ServerHost) urlSvc := service.NewURLService(cfg.HMACSecret, cfg.ServerHost)
ocrSvc.Start()
defer ocrSvc.Stop()
h := handler.New(fileSvc, ocrSvc, urlSvc) h := handler.New(fileSvc, ocrSvc, urlSvc)
r := gin.Default() r := gin.Default()
+4
View File
@@ -13,6 +13,7 @@ import (
type FileHandler struct { type FileHandler struct {
files *service.FileService files *service.FileService
urls *service.URLService urls *service.URLService
ocr *service.OCRService
} }
func (h *FileHandler) Upload(c *gin.Context) { func (h *FileHandler) Upload(c *gin.Context) {
@@ -35,6 +36,9 @@ func (h *FileHandler) Upload(c *gin.Context) {
api.Error(c, http.StatusInternalServerError, "UPLOAD_ERROR", err.Error()) api.Error(c, http.StatusInternalServerError, "UPLOAD_ERROR", err.Error())
return return
} }
h.ocr.Enqueue(result.ID, result.Path)
results = append(results, gin.H{ results = append(results, gin.H{
"id": result.ID, "id": result.ID,
"name": result.Name, "name": result.Name,
+1 -1
View File
@@ -12,7 +12,7 @@ type Handler struct {
func New(fileSvc *service.FileService, ocrSvc *service.OCRService, urlSvc *service.URLService) *Handler { func New(fileSvc *service.FileService, ocrSvc *service.OCRService, urlSvc *service.URLService) *Handler {
return &Handler{ return &Handler{
File: &FileHandler{files: fileSvc, urls: urlSvc}, File: &FileHandler{files: fileSvc, urls: urlSvc, ocr: ocrSvc},
OCR: &OCRHandler{ocr: ocrSvc, files: fileSvc}, OCR: &OCRHandler{ocr: ocrSvc, files: fileSvc},
Health: &HealthHandler{ocr: ocrSvc}, Health: &HealthHandler{ocr: ocrSvc},
} }
+1
View File
@@ -15,4 +15,5 @@ type File struct {
type UploadResult struct { type UploadResult struct {
ID string `json:"id"` ID string `json:"id"`
Name string `json:"name"` Name string `json:"name"`
Path string `json:"path"`
} }
+1
View File
@@ -62,6 +62,7 @@ func (s *FileService) Upload(file *multipart.FileHeader) (*model.UploadResult, e
return &model.UploadResult{ return &model.UploadResult{
ID: dbFile.ID, ID: dbFile.ID,
Name: dbFile.Name, Name: dbFile.Name,
Path: dst,
}, nil }, nil
} }
+58 -7
View File
@@ -1,7 +1,7 @@
package service package service
import ( import (
"fmt" "log"
"os" "os"
"strings" "strings"
@@ -9,22 +9,69 @@ import (
"github.com/vaultdrop/backend/internal/ocr" "github.com/vaultdrop/backend/internal/ocr"
) )
type OCRJob struct {
FileID string
FilePath string
}
type OCRService struct { type OCRService struct {
client *ocr.Client client *ocr.Client
fileSvc *FileService
jobs chan OCRJob
} }
func NewOCRService(cfg *config.Config) *OCRService { func NewOCRService(cfg *config.Config, fileSvc *FileService) *OCRService {
return &OCRService{ return &OCRService{
client: ocr.NewClient(cfg.OCREndpoint), client: ocr.NewClient(cfg.OCREndpoint),
fileSvc: fileSvc,
jobs: make(chan OCRJob, 100),
} }
} }
func (s *OCRService) RecognizeFromFile(filePath string) ([]ocr.TextBlock, error) { func (s *OCRService) Start() {
data, err := os.ReadFile(filePath) go s.worker()
if err != nil { log.Println("[OCR] Worker started")
return nil, fmt.Errorf("read file: %w", err) }
func (s *OCRService) Stop() {
close(s.jobs)
log.Println("[OCR] Worker stopped")
}
func (s *OCRService) Enqueue(fileID, filePath string) {
s.jobs <- OCRJob{FileID: fileID, FilePath: filePath}
log.Printf("[OCR] Enqueued file %s", fileID)
}
func (s *OCRService) worker() {
for job := range s.jobs {
s.process(job)
} }
return s.client.Recognize(data) }
func (s *OCRService) process(job OCRJob) {
log.Printf("[OCR] Processing file %s", job.FileID)
data, err := os.ReadFile(job.FilePath)
if err != nil {
log.Printf("[OCR] Failed to read file %s: %v", job.FileID, err)
return
}
blocks, err := s.client.Recognize(data)
if err != nil {
log.Printf("[OCR] Failed to recognize file %s: %v", job.FileID, err)
return
}
text := s.FlattenResults(blocks)
if err := s.fileSvc.UpdateOCRText(job.FileID, text); err != nil {
log.Printf("[OCR] Failed to update ocr_text for file %s: %v", job.FileID, err)
return
}
log.Printf("[OCR] Completed file %s (%d chars)", job.FileID, len(text))
} }
func (s *OCRService) RecognizeFromBytes(data []byte) ([]ocr.TextBlock, error) { func (s *OCRService) RecognizeFromBytes(data []byte) ([]ocr.TextBlock, error) {
@@ -42,3 +89,7 @@ func (s *OCRService) FlattenResults(blocks []ocr.TextBlock) string {
func (s *OCRService) HealthCheck() error { func (s *OCRService) HealthCheck() error {
return s.client.HealthCheck() return s.client.HealthCheck()
} }
func (s *OCRService) QueueLength() int {
return len(s.jobs)
}