IPSExtractor/services/services-count-facturas-medicas-serving/internal/pipeline/processor.go

package pipeline

import (
	"fmt"
	"log"
	"time"

	"github.com/pys/pdf-count/internal/domain"
	"github.com/pys/pdf-count/internal/model"
	"github.com/pys/pdf-count/internal/ocr"
)

type Options struct {
	DPI               int
	Lang              string
	DBSCANEps         float64
	DBSCANMinSamples  int
	IncludeRejected   bool
	LogPipelineTiming bool
}

type Processor struct {
	predictor *model.Predictor
	ocrEngine *ocr.Engine
	opts      Options
}

func NewProcessor(predictor *model.Predictor, ocrEngine *ocr.Engine, opts Options) *Processor {
	return &Processor{predictor: predictor, ocrEngine: ocrEngine, opts: opts}
}

func (p *Processor) ProcessPDF(filename string, pdfData []byte) domain.PDFResult {
	start := time.Now()
	result := domain.PDFResult{Filename: filename, Status: domain.StatusOK}
	var renderDur, ocrDur, clusterDur, inferDur time.Duration

	tRender := time.Now()
	img, pageW, pageH, err := ocr.RenderPage(pdfData, p.opts.DPI)
	renderDur = time.Since(tRender)
	if err != nil {
		result.Status = domain.StatusError
		result.Error = err.Error()
		p.finishTiming(&result, renderDur, 0, 0, 0, time.Since(start))
		return result
	}

	tOCR := time.Now()
	boxes, err := p.ocrEngine.OCRImage(img)
	ocrDur = time.Since(tOCR)
	if err != nil {
		result.Status = domain.StatusError
		result.Error = err.Error()
		p.finishTiming(&result, renderDur, ocrDur, 0, 0, time.Since(start))
		return result
	}
	if len(boxes) == 0 {
		result.ServiceCount = 0
		result.Services = []domain.ServiceRow{}
		p.finishTiming(&result, renderDur, ocrDur, 0, 0, time.Since(start))
		return result
	}

	tCluster := time.Now()
	indexed := make([]indexedBox, len(boxes))
	for i, b := range boxes {
		indexed[i] = indexedBox{Text: b.Text, X: b.X, Y: b.Y, W: b.W, H: b.H}
	}
	ClusterRows(indexed, p.opts.DBSCANEps, p.opts.DBSCANMinSamples)

	rows := buildFeatureRows(indexed, pageW, pageH)
	clusterDur = time.Since(tCluster)
	if len(rows) == 0 {
		result.ServiceCount = 0
		result.Services = []domain.ServiceRow{}
		p.finishTiming(&result, renderDur, ocrDur, clusterDur, 0, time.Since(start))
		return result
	}

	// Sin modelo ML: contar todas las filas agrupadas como servicios (heurística).
	if p.predictor == nil {
		services := make([]domain.ServiceRow, 0, len(rows))
		for _, row := range rows {
			services = append(services, domain.ServiceRow{
				Text:        row.TextoFila,
				Probability: 1.0,
			})
		}
		result.Services = services
		result.ServiceCount = len(rows)
		p.finishTiming(&result, renderDur, ocrDur, clusterDur, 0, time.Since(start))
		return result
	}

	features := make([][]float32, len(rows))
	for i, r := range rows {
		features[i] = r.Features[:]
	}

	tInfer := time.Now()
	probs, err := p.predictor.PredictProba(features)
	inferDur = time.Since(tInfer)
	if err != nil {
		result.Status = domain.StatusError
		result.Error = fmt.Sprintf("inference: %v", err)
		p.finishTiming(&result, renderDur, ocrDur, clusterDur, inferDur, time.Since(start))
		return result
	}

	threshold := p.predictor.Threshold()
	services := make([]domain.ServiceRow, 0)
	serviceCount := 0
	for i, row := range rows {
		prob := model.RoundProb(probs[i])
		isService := prob >= threshold
		if isService {
			serviceCount++
		}
		if p.opts.IncludeRejected {
			v := isService
			services = append(services, domain.ServiceRow{
				Text:        row.TextoFila,
				Probability: prob,
				IsService:   &v,
			})
		} else if isService {
			services = append(services, domain.ServiceRow{
				Text:        row.TextoFila,
				Probability: prob,
			})
		}
	}

	result.Services = services
	result.ServiceCount = serviceCount
	p.finishTiming(&result, renderDur, ocrDur, clusterDur, inferDur, time.Since(start))
	return result
}

func (p *Processor) finishTiming(result *domain.PDFResult, render, ocrD, cluster, infer, total time.Duration) {
	result.RenderMs = render.Milliseconds()
	result.OCRMs = ocrD.Milliseconds()
	result.ClusterMs = cluster.Milliseconds()
	result.InferMs = infer.Milliseconds()
	result.TotalMs = total.Milliseconds()
	if !p.opts.LogPipelineTiming {
		return
	}
	log.Printf(
		"pipeline filename=%s ocr_render=%s ocr_tesseract=%s cluster=%s infer=%s total=%s",
		result.Filename, render, ocrD, cluster, infer, total,
	)
}