Merge pull request #27 from EKami/fix/segfault

Add OD mask check

staintools/stain_extractors/vahadane_stain_extractor.py

@@ -6,6 +6,11 @@
 from staintools.utils.misc_utils import convert_RGB_to_OD, normalize_rows, get_luminosity_mask
 
 
+class VahadaneStainExtractorException(Exception):
+    def __init__(self, *args, **kwargs):
+        Exception.__init__(self, *args, **kwargs)
+
+
 class VahadaneStainExtractor(StainExtractor):
 
     @staticmethod
@@ -25,6 +30,10 @@ def get_stain_matrix(I, luminosity_threshold=0.8, dictionary_regularizer=0.1):
         OD = convert_RGB_to_OD(I).reshape((-1, 3))
         OD = OD[tissue_mask]
 
+        # Check if OD is fully masked
+        if OD.size == 0:
+            raise VahadaneStainExtractorException("Could not compute stain norm for the given input array")
+
         # do the dictionary learning
         dictionary = spams.trainDL(X=OD.T, K=2, lambda1=dictionary_regularizer, mode=2,
                                    modeD=0, posAlpha=True, posD=True, verbose=False).T

tests/test_misc.py

@@ -1,3 +1,7 @@
+import os
+from pathlib import Path
+import staintools
+from staintools.stain_extractors.vahadane_stain_extractor import VahadaneStainExtractor, VahadaneStainExtractorException
 from staintools.utils.misc_utils import *
 import numpy as np
 
@@ -20,3 +24,19 @@ def test_is_uint8_image():
     assert is_uint8_image(x)
     x = x / 255
     assert not is_uint8_image(x)
+
+
+def test_vahadane_stain_extractor():
+    script_dir = Path(os.path.dirname(os.path.abspath(__file__)))
+    i1 = staintools.read_image(str(script_dir / ".." / "data" / "i1.png"))
+    x = np.zeros(shape=i1.shape, dtype=np.uint8)
+    x[:] = 255
+
+    # Ensure the proper exception is raised
+    try:
+        VahadaneStainExtractor.get_stain_matrix(x)
+    except VahadaneStainExtractorException as e:
+        print(e)
+        assert True
+        return
+    assert False