Revise the watermark removal algorithm to be both generalizable and highly specific

Author: valandi

src/main/java/com/applitools/imagetester/ImageTester.java

@@ -705,11 +705,12 @@ private static Options getOptions() {
 
         options.addOption(Option.builder("rwauto")
             .longOpt("removeWatermarkAuto")
-            .desc("Auto-detect a vector watermark template by intersecting path " +
-                  "shapes across all PDFs in the input directory, then strip the " +
-                  "template from each. Requires at least 2 input PDFs. With -rwo, " +
-                  "writes cleaned PDFs to that directory and exits. Without -rwo, " +
-                  "cleans to a temp directory and uploads cleaned PDFs to Applitools.")
+            .desc("Auto-detect a vector watermark shared across all PDFs in the input " +
+                  "directory by its fill color, then strip only paths in that color " +
+                  "from each PDF, leaving all other content intact. Requires at least " +
+                  "2 input PDFs from the same source. With -rwo, writes cleaned PDFs to " +
+                  "that directory and exits. Without -rwo, cleans to a temp directory " +
+                  "and uploads cleaned PDFs to Applitools.")
             .hasArg(false)
             .build());
 

src/main/java/com/applitools/imagetester/lib/ColorPathStripper.java

@@ -0,0 +1,123 @@
+package com.applitools.imagetester.lib;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Deque;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import org.apache.pdfbox.contentstream.operator.Operator;
+import org.apache.pdfbox.cos.COSName;
+import org.apache.pdfbox.cos.COSStream;
+import org.apache.pdfbox.pdfparser.PDFStreamParser;
+import org.apache.pdfbox.pdfwriter.ContentStreamWriter;
+import org.apache.pdfbox.pdmodel.PDDocument;
+import org.apache.pdfbox.pdmodel.PDPage;
+
+/**
+ * Removes filled vector paths drawn in a target non-stroking (fill) color,
+ * leaving everything else — body text, images, strokes, clip paths, and fills
+ * of any other color — untouched.
+ *
+ * Watermarks of the kind this targets are stamped as filled outlines in a
+ * single muted color distinct from the document's real content, so keying
+ * removal on fill color strips the watermark without touching shared branding.
+ */
+public final class ColorPathStripper {
+
+    private static final Set<String> PATH_CONSTRUCTION_OPS = new HashSet<>(Arrays.asList(
+            "m", "l", "c", "v", "y", "h", "re", "W", "W*"));
+    private static final Set<String> FILL_PAINT_OPS = new HashSet<>(Arrays.asList(
+            "f", "F", "f*", "b", "b*", "B", "B*"));
+    private static final Set<String> NON_FILL_PAINT_OPS = new HashSet<>(Arrays.asList(
+            "S", "s", "n"));
+
+    private ColorPathStripper() {
+    }
+
+    public static void removeFromAllPages(PDDocument doc, float[] targetRgb, float tolerance) throws IOException {
+        if (targetRgb == null) return;
+        for (int i = 0; i < doc.getNumberOfPages(); i++) {
+            PDPage page = doc.getPage(i);
+            PDFStreamParser parser = new PDFStreamParser(page);
+            parser.parse();
+            List<Object> cleaned = strip(parser.getTokens(), targetRgb, tolerance);
+
+            COSStream newStream = new COSStream();
+            try (OutputStream out = newStream.createOutputStream()) {
+                new ContentStreamWriter(out).writeTokens(cleaned);
+            }
+            page.getCOSObject().setItem(COSName.CONTENTS, newStream);
+        }
+    }
+
+    public static List<Object> strip(List<Object> tokens, float[] targetRgb, float tolerance) {
+        List<Object> result = new ArrayList<>();
+        List<Object> argBuffer = new ArrayList<>();
+        List<Object> currentPath = new ArrayList<>();
+        boolean inPath = false;
+
+        Deque<float[]> stateStack = new ArrayDeque<>();
+        float[] fill = {0f, 0f, 0f};
+
+        for (Object t : tokens) {
+            if (!(t instanceof Operator)) {
+                argBuffer.add(t);
+                continue;
+            }
+            String op = ((Operator) t).getName();
+
+            if (PATH_CONSTRUCTION_OPS.contains(op)) {
+                currentPath.addAll(argBuffer);
+                currentPath.add(t);
+                argBuffer.clear();
+                inPath = true;
+            } else if (inPath && (FILL_PAINT_OPS.contains(op) || NON_FILL_PAINT_OPS.contains(op))) {
+                currentPath.addAll(argBuffer);
+                currentPath.add(t);
+                argBuffer.clear();
+                boolean drop = FILL_PAINT_OPS.contains(op) && colorMatches(fill, targetRgb, tolerance);
+                if (!drop) result.addAll(currentPath);
+                currentPath.clear();
+                inPath = false;
+            } else {
+                if (inPath) {
+                    result.addAll(currentPath);
+                    currentPath.clear();
+                    inPath = false;
+                }
+                fill = applyColorState(op, argBuffer, fill, stateStack);
+                result.addAll(argBuffer);
+                result.add(t);
+                argBuffer.clear();
+            }
+        }
+        result.addAll(currentPath);
+        result.addAll(argBuffer);
+        return result;
+    }
+
+    private static float[] applyColorState(String op, List<Object> args, float[] fill, Deque<float[]> stack) {
+        switch (op) {
+            case "q":
+                stack.push(fill.clone());
+                return fill;
+            case "Q":
+                return stack.isEmpty() ? fill : stack.pop();
+            default:
+                return DeviceColor.fromOperator(op, args, fill);
+        }
+    }
+
+    private static boolean colorMatches(float[] fill, float[] target, float tol) {
+        if (target == null) return false;
+        for (int i = 0; i < 3; i++) {
+            if (Math.abs(fill[i] - target[i]) > tol) return false;
+        }
+        return true;
+    }
+}

src/main/java/com/applitools/imagetester/lib/DeviceColor.java

@@ -0,0 +1,76 @@
+package com.applitools.imagetester.lib;
+
+import java.util.List;
+
+import org.apache.pdfbox.cos.COSNumber;
+
+/**
+ * Interprets PDF non-stroking color operators (rg/g/k and sc/scn over device
+ * color spaces) into normalized RGB. Pattern and named-colorspace operands that
+ * can't be read as plain device components leave the color unchanged.
+ */
+final class DeviceColor {
+
+    private DeviceColor() {
+    }
+
+    /** Returns the RGB a color operator sets, or {@code current} if the operator isn't a color op. */
+    static float[] fromOperator(String op, List<Object> args, float[] current) {
+        switch (op) {
+            case "rg":
+                return rgbOrCurrent(numbers(args, 3), current);
+            case "g":
+                float[] g = numbers(args, 1);
+                return g == null ? current : new float[] {g[0], g[0], g[0]};
+            case "k":
+                return cmykToRgb(numbers(args, 4), current);
+            case "sc":
+            case "scn":
+                return fromComponents(args, current);
+            default:
+                return current;
+        }
+    }
+
+    private static float[] fromComponents(List<Object> args, float[] current) {
+        int n = 0;
+        for (int i = args.size() - 1; i >= 0 && args.get(i) instanceof COSNumber; i--) n++;
+        switch (n) {
+            case 1:
+                float v = ((COSNumber) args.get(args.size() - 1)).floatValue();
+                return new float[] {v, v, v};
+            case 3:
+                return rgbOrCurrent(numbers(args, 3), current);
+            case 4:
+                return cmykToRgb(numbers(args, 4), current);
+            default:
+                return current;
+        }
+    }
+
+    private static float[] rgbOrCurrent(float[] vals, float[] current) {
+        return vals == null ? current : vals;
+    }
+
+    private static float[] cmykToRgb(float[] cmyk, float[] current) {
+        if (cmyk == null) return current;
+        return new float[] {
+                (1f - cmyk[0]) * (1f - cmyk[3]),
+                (1f - cmyk[1]) * (1f - cmyk[3]),
+                (1f - cmyk[2]) * (1f - cmyk[3])
+        };
+    }
+
+    /** Returns the last {@code count} numeric operands, or null if too few are present. */
+    private static float[] numbers(List<Object> args, int count) {
+        if (args.size() < count) return null;
+        float[] out = new float[count];
+        int start = args.size() - count;
+        for (int i = 0; i < count; i++) {
+            Object o = args.get(start + i);
+            if (!(o instanceof COSNumber)) return null;
+            out[i] = ((COSNumber) o).floatValue();
+        }
+        return out;
+    }
+}

src/main/java/com/applitools/imagetester/lib/PdfVectorWatermarkAutoMode.java

@@ -9,7 +9,6 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Set;
 
 import org.apache.pdfbox.pdmodel.PDDocument;
 import org.apache.pdfbox.pdmodel.PDPage;
@@ -28,6 +27,9 @@
  */
 public final class PdfVectorWatermarkAutoMode {
 
+    /** Per-channel RGB tolerance when matching a path's fill color to the detected watermark color. */
+    private static final float WATERMARK_COLOR_TOLERANCE = 0.05f;
+
     static final String SINGLE_PDF_NOTICE =
         "\nHeads up: you're testing one PDF on its own." +
         "\n" +
@@ -80,22 +82,24 @@ public static int run(File inputRoot, File outDir, String optionalTextHint, Logg
             List<File> pdfs = entry.getValue();
             String groupLabel = groupLabel(inputRoot, groupDir);
 
-            Set<String> pathFingerprint;
-            Set<String> varyingOpSeqs;
-            Set<String> varyingForms;
+            float[] watermarkColor;
             try {
-                pathFingerprint = PathFingerprinter.intersection(pdfs);
-                varyingOpSeqs = OpSequenceVarianceFinder.findVarying(pdfs);
-                varyingForms = VaryingFormFinder.findVarying(pdfs);
+                watermarkColor = WatermarkColorDetector.detect(pdfs);
             } catch (IOException e) {
                 logger.printMessage(String.format(
-                        "Failed to compute watermark fingerprint for %s: %s",
-                        groupLabel, e.getMessage()));
+                        "Failed to detect watermark for %s: %s", groupLabel, e.getMessage()));
                 continue;
             }
-            logger.printMessage(String.format(
-                    "[%s] Watermark fingerprint: %d shared shape(s), %d varying op-seq(s), %d varying form(s) across %d PDF(s)",
-                    groupLabel, pathFingerprint.size(), varyingOpSeqs.size(), varyingForms.size(), pdfs.size()));
+            if (watermarkColor == null) {
+                logger.printMessage(String.format(
+                        "[%s] No vector watermark detected across %d PDF(s); copying unchanged",
+                        groupLabel, pdfs.size()));
+            } else {
+                logger.printMessage(String.format(
+                        "[%s] Watermark color rgb(%d, %d, %d) detected across %d PDF(s)",
+                        groupLabel, to255(watermarkColor[0]), to255(watermarkColor[1]),
+                        to255(watermarkColor[2]), pdfs.size()));
+            }
 
             for (File pdf : pdfs) {
                 File output = resolveOutput(inputRoot, outDir, pdf);
@@ -104,7 +108,7 @@ public static int run(File inputRoot, File outDir, String optionalTextHint, Logg
                     continue;
                 }
                 try {
-                    cleanOnePdf(pdf, pathFingerprint, varyingOpSeqs, varyingForms, optionalTextHint, output);
+                    cleanOnePdf(pdf, watermarkColor, optionalTextHint, output);
                     processed++;
                 } catch (IOException e) {
                     logger.printMessage("Failed to clean " + pdf.getAbsolutePath() + ": " + e.getMessage());
@@ -116,8 +120,7 @@ public static int run(File inputRoot, File outDir, String optionalTextHint, Logg
         return 0;
     }
 
-    private static void cleanOnePdf(File input, Set<String> pathFingerprint, Set<String> varyingOpSeqs,
-                                     Set<String> varyingForms, String optionalTextHint, File output)
+    private static void cleanOnePdf(File input, float[] watermarkColor, String optionalTextHint, File output)
             throws IOException {
         boolean hasTextHint = optionalTextHint != null && !optionalTextHint.trim().isEmpty();
         try (PDDocument source = PDDocument.load(input);
@@ -129,12 +132,15 @@ private static void cleanOnePdf(File input, Set<String> pathFingerprint, Set<Str
                 }
                 cleaned.addPage(page);
             }
-            VectorWatermarkRemover.removeFromAllPages(cleaned, pathFingerprint, varyingOpSeqs);
-            FormXObjectStripper.emptyForms(cleaned, varyingForms);
+            ColorPathStripper.removeFromAllPages(cleaned, watermarkColor, WATERMARK_COLOR_TOLERANCE);
             cleaned.save(output);
         }
     }
 
+    private static int to255(float component) {
+        return Math.round(component * 255f);
+    }
+
     private static Map<File, List<File>> groupPdfsByDirectory(File root) {
         Map<File, List<File>> result = new LinkedHashMap<>();
         if (root.isFile()) {