ocr worker
This commit is contained in:
@@ -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()
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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},
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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"`
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
|
||||||
}
|
}
|
||||||
return s.client.Recognize(data)
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user