Initial orientator

internal/pipeline/pipeline.go

@@ -0,0 +1,198 @@
+// Package pipeline orchestrates orientation detection and rotation across
+// images, PDFs, and archives.
+package pipeline
+
+import (
+	"archive/zip"
+	"bytes"
+	"fmt"
+	"image"
+	_ "image/jpeg"
+	_ "image/png"
+	"io"
+	"path/filepath"
+	"strings"
+	"sync"
+
+	_ "golang.org/x/image/tiff"
+	_ "golang.org/x/image/webp"
+
+	"orientator/internal/detector"
+	"orientator/internal/extract"
+	"orientator/internal/rasterize"
+)
+
+// PageReport describes what happened to a single page or image.
+type PageReport struct {
+	File       string  `json:"file"`
+	Page       int     `json:"page,omitempty"` // 0 for plain images
+	Angle      int     `json:"angle"`
+	Confidence float64 `json:"confidence"`
+}
+
+// Result holds the rewritten output and a per-page report.
+type Result struct {
+	OutputName  string
+	OutputBytes []byte
+	Reports     []PageReport
+}
+
+// Process processes a single uploaded file and returns the corrected output.
+// For archives the output is a zip of all corrected entries.
+func Process(name string, data []byte) (*Result, error) {
+	det, err := detector.Get()
+	if err != nil {
+		return nil, fmt.Errorf("detector init: %w", err)
+	}
+
+	if extract.IsArchive(name) {
+		return processArchive(det, name, data)
+	}
+	return processSingle(det, name, data)
+}
+
+func processSingle(det *detector.Detector, name string, data []byte) (*Result, error) {
+	ext := strings.ToLower(filepath.Ext(name))
+	switch ext {
+	case ".pdf":
+		out, reports, err := processPDF(det, name, data)
+		if err != nil {
+			return nil, err
+		}
+		return &Result{OutputName: name, OutputBytes: out, Reports: reports}, nil
+	case ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp", ".bmp":
+		out, outName, report, err := processImage(det, name, data)
+		if err != nil {
+			return nil, err
+		}
+		return &Result{OutputName: outName, OutputBytes: out, Reports: []PageReport{report}}, nil
+	}
+	return nil, fmt.Errorf("unsupported file type: %s", ext)
+}
+
+func processImage(det *detector.Detector, name string, data []byte) ([]byte, string, PageReport, error) {
+	img, format, err := image.Decode(bytes.NewReader(data))
+	if err != nil {
+		return nil, "", PageReport{}, fmt.Errorf("decode %s: %w", name, err)
+	}
+	angle, conf, err := det.Predict(img)
+	if err != nil {
+		return nil, "", PageReport{}, err
+	}
+	rotated := RotateImage(img, angle)
+	out, ext, err := EncodeImage(rotated, format)
+	if err != nil {
+		return nil, "", PageReport{}, err
+	}
+	outName := replaceExt(name, "."+ext)
+	return out, outName, PageReport{File: name, Angle: angle, Confidence: conf}, nil
+}
+
+func processPDF(det *detector.Detector, name string, data []byte) ([]byte, []PageReport, error) {
+	pdf, err := rasterize.Open(data)
+	if err != nil {
+		return nil, nil, err
+	}
+	defer pdf.Close()
+
+	n := pdf.NumPages()
+	angles := make([]int, n)
+	reports := make([]PageReport, 0, n)
+	for i := 0; i < n; i++ {
+		img, err := pdf.RenderPage(i)
+		if err != nil {
+			return nil, nil, err
+		}
+		angle, conf, err := det.Predict(img)
+		if err != nil {
+			return nil, nil, err
+		}
+		angles[i] = angle
+		reports = append(reports, PageReport{File: name, Page: i + 1, Angle: angle, Confidence: conf})
+	}
+	out, err := rasterize.RotatePDFPages(data, angles)
+	if err != nil {
+		return nil, nil, err
+	}
+	return out, reports, nil
+}
+
+func processArchive(det *detector.Detector, name string, data []byte) (*Result, error) {
+	var (
+		mu      sync.Mutex
+		buf     bytes.Buffer
+		zw      = zip.NewWriter(&buf)
+		reports []PageReport
+	)
+
+	err := extract.Walk(name, data, func(e extract.Entry) error {
+		// Recurse into nested archives.
+		if extract.IsArchive(e.Name) {
+			sub, err := Process(e.Name, e.Data)
+			if err != nil {
+				return writeZipFailure(zw, e.Name, err)
+			}
+			mu.Lock()
+			reports = append(reports, sub.Reports...)
+			mu.Unlock()
+			return writeZipEntry(zw, sub.OutputName, sub.OutputBytes)
+		}
+
+		ext := strings.ToLower(filepath.Ext(e.Name))
+		switch ext {
+		case ".pdf":
+			out, rs, err := processPDF(det, e.Name, e.Data)
+			if err != nil {
+				return writeZipFailure(zw, e.Name, err)
+			}
+			mu.Lock()
+			reports = append(reports, rs...)
+			mu.Unlock()
+			return writeZipEntry(zw, e.Name, out)
+		case ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp", ".bmp":
+			out, outName, r, err := processImage(det, e.Name, e.Data)
+			if err != nil {
+				return writeZipFailure(zw, e.Name, err)
+			}
+			mu.Lock()
+			reports = append(reports, r)
+			mu.Unlock()
+			return writeZipEntry(zw, outName, out)
+		default:
+			// Pass-through any non-document files.
+			return writeZipEntry(zw, e.Name, e.Data)
+		}
+	})
+	if err != nil {
+		return nil, err
+	}
+	if err := zw.Close(); err != nil {
+		return nil, err
+	}
+
+	outName := replaceExt(name, ".zip")
+	return &Result{OutputName: outName, OutputBytes: buf.Bytes(), Reports: reports}, nil
+}
+
+func writeZipEntry(zw *zip.Writer, name string, data []byte) error {
+	w, err := zw.Create(name)
+	if err != nil {
+		return err
+	}
+	_, err = io.Copy(w, bytes.NewReader(data))
+	return err
+}
+
+func writeZipFailure(zw *zip.Writer, name string, cause error) error {
+	w, err := zw.Create(name + ".error.txt")
+	if err != nil {
+		return err
+	}
+	_, _ = w.Write([]byte(cause.Error()))
+	return nil
+}
+
+func replaceExt(name, newExt string) string {
+	old := filepath.Ext(name)
+	return strings.TrimSuffix(name, old) + newExt
+}