[data.utils;engine.predictor] Re-use package image overlaying technique

bob/ip/binseg/data/utils.py

@@ -54,8 +54,8 @@ def overlayed_image(
         An RGB PIL image that represents the original image for analysis
 
     label : PIL.Image.Image
-        A PIL image in mode "1" that represents the labelled elements in the
-        image.  White pixels represent the labelled object.  Black pixels
+        A PIL image in mode "1" or "L" that represents the labelled elements in
+        the image.  White pixels represent the labelled object.  Black pixels
         represent background.
 
     mask : py:class:`PIL.Image.Image`, Optional
@@ -95,7 +95,13 @@ def overlayed_image(
     # slight "label_color" tone on top, then composite with original image, not
     # to loose brightness.
     retval = PIL.Image.blend(img, label_colored, alpha)
-    retval = PIL.Image.composite(img, retval, invert_mode1_image(label))
+    if label.mode == "1":
+        composite_mask = invert_mode1_image(label)
+    elif label.mode == "L":
+        composite_mask = PIL.ImageOps.invert(label)
+    else:
+        raise TypeError(f"Label image mode {label.mode} != ('1', 'L')")
+    retval = PIL.Image.composite(img, retval, composite_mask)
 
     # creates a representation of the mask negative with the right color
     if mask is not None:

bob/ip/binseg/engine/predictor.py

@@ -15,6 +15,7 @@ import torchvision.transforms.functional as VF
 import h5py
 
 from ..utils.summary import summary
+from ..data.utils import overlayed_image
 
 import logging
 logger = logging.getLogger(__name__)
@@ -34,7 +35,7 @@ def _save_hdf5(stem, prob, output_folder):
         Monochrome Image with prediction maps
 
     output_folder : str
-        path where to store overlayed results
+        path where to store predictions
 
     """
 
@@ -101,11 +102,7 @@ def _save_overlayed_png(stem, image, prob, output_folder):
 
     image = VF.to_pil_image(image)
     prob = VF.to_pil_image(prob.cpu())
-
-    # color and overlay
-    prob_green = PIL.ImageOps.colorize(prob, (0, 0, 0), (0, 255, 0))
-    overlayed = PIL.Image.blend(image, prob_green, 0.4)
-    _save_image(stem, '.png', overlayed, output_folder)
+    _save_image(stem, '.png', overlayed_image(image, prob), output_folder)
 
 
 def run(model, data_loader, device, output_folder, overlayed_folder):