Initial orientator
This commit is contained in:
80
internal/pipeline/imgrot.go
Normal file
80
internal/pipeline/imgrot.go
Normal file
@@ -0,0 +1,80 @@
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/jpeg"
|
||||
"image/png"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/image/tiff"
|
||||
_ "golang.org/x/image/webp" // decode-only; we re-encode webp as png
|
||||
)
|
||||
|
||||
// RotateImage rotates src clockwise by angle (must be 0/90/180/270).
|
||||
func RotateImage(src image.Image, angle int) image.Image {
|
||||
angle = ((angle % 360) + 360) % 360
|
||||
if angle == 0 {
|
||||
return src
|
||||
}
|
||||
b := src.Bounds()
|
||||
w, h := b.Dx(), b.Dy()
|
||||
switch angle {
|
||||
case 180:
|
||||
dst := image.NewRGBA(image.Rect(0, 0, w, h))
|
||||
for y := 0; y < h; y++ {
|
||||
for x := 0; x < w; x++ {
|
||||
dst.Set(w-1-x, h-1-y, src.At(b.Min.X+x, b.Min.Y+y))
|
||||
}
|
||||
}
|
||||
return dst
|
||||
case 90:
|
||||
dst := image.NewRGBA(image.Rect(0, 0, h, w))
|
||||
for y := 0; y < h; y++ {
|
||||
for x := 0; x < w; x++ {
|
||||
dst.Set(h-1-y, x, src.At(b.Min.X+x, b.Min.Y+y))
|
||||
}
|
||||
}
|
||||
return dst
|
||||
case 270:
|
||||
dst := image.NewRGBA(image.Rect(0, 0, h, w))
|
||||
for y := 0; y < h; y++ {
|
||||
for x := 0; x < w; x++ {
|
||||
dst.Set(y, w-1-x, src.At(b.Min.X+x, b.Min.Y+y))
|
||||
}
|
||||
}
|
||||
return dst
|
||||
}
|
||||
return src
|
||||
}
|
||||
|
||||
// EncodeImage re-encodes img using format hint from the original filename / detected format.
|
||||
// Falls back to PNG when format is not directly writable.
|
||||
func EncodeImage(img image.Image, format string) ([]byte, string, error) {
|
||||
var buf bytes.Buffer
|
||||
switch strings.ToLower(format) {
|
||||
case "jpeg", "jpg":
|
||||
if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: 92}); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return buf.Bytes(), "jpg", nil
|
||||
case "tiff", "tif":
|
||||
if err := tiff.Encode(&buf, img, &tiff.Options{Compression: tiff.Deflate}); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return buf.Bytes(), "tiff", nil
|
||||
case "png", "":
|
||||
if err := png.Encode(&buf, img); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return buf.Bytes(), "png", nil
|
||||
case "webp":
|
||||
// Re-encode to PNG (no webp encoder in stdlib/x/image).
|
||||
if err := png.Encode(&buf, img); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return buf.Bytes(), "png", nil
|
||||
}
|
||||
return nil, "", fmt.Errorf("unsupported format: %s", format)
|
||||
}
|
||||
198
internal/pipeline/pipeline.go
Normal file
198
internal/pipeline/pipeline.go
Normal file
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user