diff options
| -rw-r--r-- | cmd/rescribe/main.go | 100 |
1 files changed, 100 insertions, 0 deletions
diff --git a/cmd/rescribe/main.go b/cmd/rescribe/main.go index 16c284c..96f2853 100644 --- a/cmd/rescribe/main.go +++ b/cmd/rescribe/main.go @@ -15,6 +15,8 @@ import ( "context" "flag" "fmt" + "image" + "image/draw" "image/jpeg" "image/png" "io" @@ -397,6 +399,12 @@ func extractPdfImgs(ctx context.Context, path string) (string, error) { if p.Page(pgnum).V.IsNull() { continue } + var rotate int64 + for v := p.Page(pgnum).V; !v.IsNull(); v = v.Key("Parent") { + if r := v.Key("Rotate"); !r.IsNull() { + rotate = r.Int64() + } + } res := p.Page(pgnum).Resources() if res.Kind() != pdf.Dict { continue @@ -432,6 +440,13 @@ func extractPdfImgs(ctx context.Context, path string) (string, error) { if err != nil { return tempdir, fmt.Errorf("Error removing extracted image %s from PDF: %v\n", fn, err) } + + if rotate != 0 { + err = rotateImage(path, rotate) + if err != nil { + return tempdir, fmt.Errorf("Error rotating extracted image %s from PDF: %v\n", fn, err) + } + } } } // TODO: check for places where there are multiple images per page, and only keep largest ones where that's the case @@ -512,6 +527,91 @@ func rmIfNotImage(f string) error { return nil } +// rotateImage rotates an image at the given path by the given angle +func rotateImage(path string, angle int64) error { + switch angle { + case 90: + // proceed with the rest of the function + case 180, 270: + // rotate the image again first, as many times as necessary. + // this is inefficient but easy. + err := rotateImage(path, angle-90) + if err != nil { + return fmt.Errorf("error with a rotation run: %w", err) + } + default: + return fmt.Errorf("Rotation angle of %d is not supported", angle) + } + + r, err := os.Open(path) + defer r.Close() + if err != nil { + return fmt.Errorf("Failed to open image: %w", err) + } + img, err := png.Decode(r) + if err != nil { + r.Close() + r, err = os.Open(path) + defer r.Close() + if err != nil { + return fmt.Errorf("Failed to open image: %w", err) + } + img, err = jpeg.Decode(r) + } + if err != nil { + r.Close() + r, err = os.Open(path) + defer r.Close() + if err != nil { + return fmt.Errorf("Failed to open image: %w", err) + } + img, err = tiff.Decode(r) + } + if err != nil { + return fmt.Errorf("Failed to decode image as png, jpeg or tiff: %w", err) + } + + b := img.Bounds() + + orig := image.NewRGBA(b) + draw.Draw(orig, b, img, b.Min, draw.Src) + + newb := image.Rectangle{ + Min: image.Point{X: 0, Y: 0}, + Max: image.Point{X: b.Dy(), Y: b.Dx()}, + } + new := image.NewRGBA(newb) + + for x := b.Min.X; x < b.Max.X; x++ { + desty := newb.Min.Y + x + for y := b.Max.Y; y > b.Min.Y; y-- { + destx := b.Dy() - y + newb.Min.X + new.SetRGBA(destx, desty, orig.RGBAAt(x, y)) + } + } + + err = r.Close() + if err != nil { + return fmt.Errorf("Failed to close image: %w", err) + } + w, err := os.Create(path) + if err != nil { + return fmt.Errorf("Failed to create rotated image: %w", err) + } + defer w.Close() + + if !strings.HasSuffix(path, ".jpg") { + err = jpeg.Encode(w, new, nil) + } else { + err = png.Encode(w, new) + } + if err != nil { + return fmt.Errorf("Failed to encode rotated image: %w", err) + } + + return nil +} + func startProcess(ctx context.Context, logger *log.Logger, tessCommand string, bookdir string, bookname string, trainingName string, savedir string, tessdir string, nowipe bool, fullpdf bool) error { cmd := exec.Command(tessCommand, "--help") pipeline.HideCmd(cmd) |