feat(personal): render bounded PDF OCR pages

2026-07-12 21:17:00 +08:00
parent 88bfe1d815
commit 11bcef208b
3 changed files with 234 additions and 0 deletions
@@ -54,6 +54,12 @@
            <version>3.2.2</version>
        </dependency>

+        <dependency>
+            <groupId>org.apache.pdfbox</groupId>
+            <artifactId>pdfbox</artifactId>
+            <version>3.0.5</version>
+        </dependency>
+
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-test</artifactId>
@@ -0,0 +1,171 @@
+package org.dromara.aihr.personal.service;
+
+import org.apache.pdfbox.Loader;
+import org.apache.pdfbox.pdmodel.PDDocument;
+import org.apache.pdfbox.rendering.ImageType;
+import org.apache.pdfbox.rendering.PDFRenderer;
+import org.springframework.stereotype.Component;
+
+import javax.imageio.IIOImage;
+import javax.imageio.ImageIO;
+import javax.imageio.ImageWriteParam;
+import javax.imageio.ImageWriter;
+import javax.imageio.stream.ImageOutputStream;
+import java.awt.Graphics2D;
+import java.awt.RenderingHints;
+import java.awt.image.BufferedImage;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+@Component
+public class PersonalPdfPageRenderer {
+
+    public static final int MAX_PAGES = 200;
+    public static final int BATCH_SIZE = 20;
+    public static final int MAX_IMAGE_BYTES = 4 * 1024 * 1024;
+    private static final int MAX_IMAGE_EDGE = 2200;
+    private static final float RENDER_DPI = 144F;
+    private static final float JPEG_QUALITY = 0.82F;
+
+    public int pageCount(byte[] pdfBytes) {
+        try (PDDocument document = load(pdfBytes)) {
+            return requireSupportedPageCount(document.getNumberOfPages());
+        } catch (PdfPageLimitException exception) {
+            throw exception;
+        } catch (IOException exception) {
+            throw new PdfRenderException("PERSONAL_PDF_INVALID", exception);
+        }
+    }
+
+    public int requireSupportedPageCount(int pageCount) {
+        if (pageCount < 1) {
+            throw new PdfRenderException("PERSONAL_PDF_EMPTY");
+        }
+        if (pageCount > MAX_PAGES) {
+            throw new PdfPageLimitException(pageCount);
+        }
+        return pageCount;
+    }
+
+    public List<RenderedPage> render(byte[] pdfBytes, int startPageIndex, int batchSize) {
+        if (startPageIndex < 0 || batchSize < 1 || batchSize > BATCH_SIZE) {
+            throw new IllegalArgumentException("invalid PDF OCR batch bounds");
+        }
+        try (PDDocument document = load(pdfBytes)) {
+            int pageCount = requireSupportedPageCount(document.getNumberOfPages());
+            if (startPageIndex >= pageCount) {
+                return List.of();
+            }
+            PDFRenderer renderer = new PDFRenderer(document);
+            int end = Math.min(pageCount, startPageIndex + batchSize);
+            List<RenderedPage> pages = new ArrayList<>(end - startPageIndex);
+            for (int pageIndex = startPageIndex; pageIndex < end; pageIndex++) {
+                BufferedImage image = renderer.renderImageWithDPI(pageIndex, RENDER_DPI, ImageType.RGB);
+                pages.add(new RenderedPage(pageIndex, encodeBoundedJpeg(scaleToLimit(image)), "image/jpeg"));
+            }
+            return List.copyOf(pages);
+        } catch (PdfPageLimitException exception) {
+            throw exception;
+        } catch (IOException exception) {
+            throw new PdfRenderException("PERSONAL_PDF_RENDER_FAILED", exception);
+        }
+    }
+
+    private PDDocument load(byte[] pdfBytes) throws IOException {
+        if (pdfBytes == null || pdfBytes.length == 0) {
+            throw new PdfRenderException("PERSONAL_PDF_EMPTY");
+        }
+        return Loader.loadPDF(pdfBytes);
+    }
+
+    private BufferedImage scaleToLimit(BufferedImage source) {
+        int longest = Math.max(source.getWidth(), source.getHeight());
+        if (longest <= MAX_IMAGE_EDGE) {
+            return source;
+        }
+        double ratio = (double) MAX_IMAGE_EDGE / longest;
+        return scale(source, Math.max(1, (int) Math.round(source.getWidth() * ratio)),
+            Math.max(1, (int) Math.round(source.getHeight() * ratio)));
+    }
+
+    private byte[] encodeBoundedJpeg(BufferedImage source) throws IOException {
+        BufferedImage current = source;
+        for (int attempt = 0; attempt < 4; attempt++) {
+            byte[] encoded = encodeJpeg(current, Math.max(0.55F, JPEG_QUALITY - attempt * 0.08F));
+            if (encoded.length <= MAX_IMAGE_BYTES) {
+                return encoded;
+            }
+            current = scale(current, Math.max(1, current.getWidth() * 3 / 4),
+                Math.max(1, current.getHeight() * 3 / 4));
+        }
+        throw new PdfRenderException("PERSONAL_PDF_PAGE_TOO_LARGE");
+    }
+
+    private byte[] encodeJpeg(BufferedImage image, float quality) throws IOException {
+        Iterator<ImageWriter> writers = ImageIO.getImageWritersByFormatName("jpeg");
+        if (!writers.hasNext()) {
+            throw new PdfRenderException("PERSONAL_PDF_JPEG_UNAVAILABLE");
+        }
+        ImageWriter writer = writers.next();
+        try (ByteArrayOutputStream output = new ByteArrayOutputStream();
+             ImageOutputStream imageOutput = ImageIO.createImageOutputStream(output)) {
+            writer.setOutput(imageOutput);
+            ImageWriteParam params = writer.getDefaultWriteParam();
+            params.setCompressionMode(ImageWriteParam.MODE_EXPLICIT);
+            params.setCompressionQuality(quality);
+            writer.write(null, new IIOImage(image, null, null), params);
+            return output.toByteArray();
+        } finally {
+            writer.dispose();
+        }
+    }
+
+    private BufferedImage scale(BufferedImage source, int width, int height) {
+        BufferedImage target = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
+        Graphics2D graphics = target.createGraphics();
+        try {
+            graphics.setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BILINEAR);
+            graphics.drawImage(source, 0, 0, width, height, null);
+        } finally {
+            graphics.dispose();
+        }
+        return target;
+    }
+
+    public record RenderedPage(int pageIndex, byte[] bytes, String mimeType) {
+    }
+
+    public static class PdfRenderException extends RuntimeException {
+        private final String code;
+
+        public PdfRenderException(String code) {
+            super(code);
+            this.code = code;
+        }
+
+        public PdfRenderException(String code, Throwable cause) {
+            super(code, cause);
+            this.code = code;
+        }
+
+        public String code() {
+            return code;
+        }
+    }
+
+    public static final class PdfPageLimitException extends PdfRenderException {
+        private final int actualPages;
+
+        public PdfPageLimitException(int actualPages) {
+            super("PERSONAL_PDF_PAGE_LIMIT");
+            this.actualPages = actualPages;
+        }
+
+        public int actualPages() {
+            return actualPages;
+        }
+    }
+}
@@ -0,0 +1,57 @@
+package org.dromara.aihr.personal;
+
+import org.apache.pdfbox.pdmodel.PDDocument;
+import org.apache.pdfbox.pdmodel.PDPage;
+import org.dromara.aihr.personal.service.PersonalPdfPageRenderer;
+import org.dromara.aihr.personal.service.PersonalPdfPageRenderer.PdfPageLimitException;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayOutputStream;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@Tag("dev")
+class PersonalPdfPageRendererTest {
+
+    private final PersonalPdfPageRenderer renderer = new PersonalPdfPageRenderer();
+
+    @Test
+    void countsPagesAndRendersOnlyRequestedBatch() throws Exception {
+        byte[] pdf = pdfWithPages(8);
+
+        assertEquals(8, renderer.pageCount(pdf));
+        List<PersonalPdfPageRenderer.RenderedPage> pages = renderer.render(pdf, 0, 2);
+
+        assertEquals(List.of(0, 1), pages.stream().map(PersonalPdfPageRenderer.RenderedPage::pageIndex).toList());
+        assertTrue(pages.stream().allMatch(page -> "image/jpeg".equals(page.mimeType())));
+        assertTrue(pages.stream().allMatch(page -> page.bytes().length > 0));
+        assertTrue(pages.stream().allMatch(page -> page.bytes().length <= PersonalPdfPageRenderer.MAX_IMAGE_BYTES));
+    }
+
+    @Test
+    void rejectsPdfAboveHardPageLimit() {
+        assertThrows(PdfPageLimitException.class, () -> renderer.requireSupportedPageCount(201));
+        assertEquals(200, renderer.requireSupportedPageCount(200));
+    }
+
+    @Test
+    void rejectsInvalidBatchBounds() throws Exception {
+        byte[] pdf = pdfWithPages(2);
+        assertThrows(IllegalArgumentException.class, () -> renderer.render(pdf, -1, 1));
+        assertThrows(IllegalArgumentException.class, () -> renderer.render(pdf, 0, 21));
+    }
+
+    private byte[] pdfWithPages(int count) throws Exception {
+        try (PDDocument document = new PDDocument(); ByteArrayOutputStream output = new ByteArrayOutputStream()) {
+            for (int index = 0; index < count; index++) {
+                document.addPage(new PDPage());
+            }
+            document.save(output);
+            return output.toByteArray();
+        }
+    }
+}