From b1c7094adb17c8534964f2056d07e4ef9755bac6 Mon Sep 17 00:00:00 2001
From: Kok Wei <52521979+kwcckw@users.noreply.github.com>
Date: Mon, 16 Oct 2023 15:05:54 +0800
Subject: [PATCH] Fixed bug in PaperFactory.

---
 augraphy/base/paperfactory.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/augraphy/base/paperfactory.py b/augraphy/base/paperfactory.py
index 78370e4..8072cf3 100644
--- a/augraphy/base/paperfactory.py
+++ b/augraphy/base/paperfactory.py
@@ -188,7 +188,18 @@ def __call__(self, image, layer=None, mask=None, keypoints=None, bounding_boxes=
                 brighten_max = 1 + brighten_ratio
                 brightness = Brightness(brightness_range=(brighten_min, brighten_max), min_brightness=1)
                 texture = brightness(texture)
-            return texture
+
+            # check for additional output of mask, keypoints and bounding boxes
+            outputs_extra = []
+            if mask is not None or keypoints is not None or bounding_boxes is not None:
+                outputs_extra = [mask, keypoints, bounding_boxes]
+
+            # returns additional mask, keypoints and bounding boxes if there is additional input
+            if outputs_extra:
+                # returns in the format of [image, mask, keypoints, bounding_boxes]
+                return [texture] + outputs_extra
+            else:
+                return texture
 
     def check_paper_edges(self, texture):
         """Crop image section with better texture.