Reduce cyclomatic complexity readbayer

plantcv/plantcv/readbayer.py

@@ -8,25 +8,25 @@
 
 
 def readbayer(filename, bayerpattern='BG', alg='default'):
-    """
-    Read image from file that has a Bayer mosaic.
-
-    Inputs:
-    filename = name of image file
-    bayerpattern = arrangement of the pixels. Often found by trial and error. ("BG","GB","RG","GR")
-    alg = algorithm with which to demosaic the image. ("default","EdgeAware","VariableNumberGradients")
+    """Read image from file that has a Bayer mosaic.
 
-    Returns:
-    img      = image object as numpy array
-    path     = path to image file
-    img_name = name of image file
+    Parameters
+    ----------
+    filename : str
+        Name of image file
+    bayerpattern : str, optional
+        Arrangement of the pixels. ("BG","GB","RG","GR"), by default 'BG'
+    alg : str, optional
+        Algorithm with which to demosaic the image ("default","EdgeAware","VariableNumberGradients"), by default 'default'
 
-    :param filename: str
-    :param alg: str
-    :param bayerpattern: str
-    :return img: numpy.ndarray
-    :return path: str
-    :return img_name: str
+    Returns
+    -------
+    numpy.ndarray
+        Image with Bayer mosaic demosaiced
+    str
+        Path to the image file
+    str
+        Name of the image file
     """
     # bayerpattern is defined as the colors of the pixels in the 2nd and 3rd column of the 2nd row.
     # see https://docs.opencv.org/3.2.0/de/d25/imgproc_color_conversions.html
@@ -37,35 +37,38 @@ def readbayer(filename, bayerpattern='BG', alg='default'):
     image_raw = cv2.imread(filename, flags=-1)
 
     if image_raw is None:
-        fatal_error("Failed to open " + filename)
+        fatal_error(f"Failed to open {filename}")
+
+    # Itemize the Bayer pattern and algorithm
+    converters = {
+        "DEFAULT": {
+            "BG": cv2.COLOR_BayerBG2BGR,
+            "GB": cv2.COLOR_BayerGB2BGR,
+            "RG": cv2.COLOR_BayerRG2BGR,
+            "GR": cv2.COLOR_BayerGR2BGR
+        },
+        "EDGEAWARE": {
+            "BG": cv2.COLOR_BayerBG2BGR_EA,
+            "GB": cv2.COLOR_BayerGB2BGR_EA,
+            "RG": cv2.COLOR_BayerRG2BGR_EA,
+            "GR": cv2.COLOR_BayerGR2BGR_EA
+        },
+        "VARIABLENUMBERGRADIENTS": {
+            "BG": cv2.COLOR_BayerBG2BGR_VNG,
+            "GB": cv2.COLOR_BayerGB2BGR_VNG,
+            "RG": cv2.COLOR_BayerRG2BGR_VNG,
+            "GR": cv2.COLOR_BayerGR2BGR_VNG
+        }
+    }
 
-    if alg.upper() == 'DEFAULT':
-        if bayerpattern.upper() == 'BG':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerBG2BGR)
-        elif bayerpattern.upper() == 'GB':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerGB2BGR)
-        elif bayerpattern.upper() == 'RG':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerRG2BGR)
-        elif bayerpattern.upper() == 'GR':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerGR2BGR)
-    elif alg.upper() == 'EDGEAWARE':
-        if bayerpattern.upper() == 'BG':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerBG2BGR_EA)
-        elif bayerpattern.upper() == 'GB':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerGB2BGR_EA)
-        elif bayerpattern.upper() == 'RG':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerRG2BGR_EA)
-        elif bayerpattern.upper() == 'GR':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerGR2BGR_EA)
-    elif alg.upper() == 'VARIABLENUMBERGRADIENTS':
-        if bayerpattern.upper() == 'BG':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerBG2BGR_VNG)
-        elif bayerpattern.upper() == 'GB':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerGB2BGR_VNG)
-        elif bayerpattern.upper() == 'RG':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerRG2BGR_VNG)
-        elif bayerpattern.upper() == 'GR':
-            img = cv2.cvtColor(image_raw, cv2.COLOR_BayerGR2BGR_VNG)
+    # Check if the algorithm is valid
+    if alg.upper() not in converters:
+        fatal_error(f"Invalid algorithm '{alg}'. Choose from {list(converters.keys())}.")
+    # Check if the bayer pattern is valid
+    if bayerpattern.upper() not in converters[alg.upper()]:
+        fatal_error(f"Invalid Bayer pattern '{bayerpattern}'. Choose from {list(converters[alg.upper()].keys())}.")
+    # Convert the image using the specified algorithm and Bayer pattern
+    img = cv2.cvtColor(image_raw, converters[alg.upper()][bayerpattern.upper()])
 
     # Split path from filename
     path, img_name = os.path.split(filename)

tests/plantcv/test_readbayer.py

@@ -24,3 +24,15 @@ def test_readbayer_default_bad_input():
     """Test for PlantCV."""
     with pytest.raises(RuntimeError):
         _, _, _ = readbayer(filename="no-image.png", bayerpattern="GR", alg="default")
+
+
+def test_readbayer_bad_algorithm(test_data):
+    """Test for PlantCV."""
+    with pytest.raises(RuntimeError):
+        _, _, _ = readbayer(filename=test_data.bayer_img, bayerpattern="GR", alg="bad_algorithm")
+
+
+def test_readbayer_bad_bayerpattern(test_data):
+    """Test for PlantCV."""
+    with pytest.raises(RuntimeError):
+        _, _, _ = readbayer(filename=test_data.bayer_img, bayerpattern="bad_bayerpattern", alg="default")