From 8a2090097ad1b042996fa1cfb2358c1384a4de5c Mon Sep 17 00:00:00 2001
From: Dominique Makowski <dom.mak19@gmail.com>
Date: Fri, 25 Feb 2022 09:26:58 +0800
Subject: [PATCH] add 'features' option to analyze_image

---
 pyllusion/utilities/analyze_color.py |  2 +-
 pyllusion/utilities/analyze_image.py | 41 ++++++++++++++++++----------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/pyllusion/utilities/analyze_color.py b/pyllusion/utilities/analyze_color.py
index d11fc0f..bdae303 100644
--- a/pyllusion/utilities/analyze_color.py
+++ b/pyllusion/utilities/analyze_color.py
@@ -76,8 +76,8 @@ def analyze_color(image, average=True):
     out["RedGreen"] = lab[:, :, 1]
     out["BlueYellow"] = lab[:, :, 2]
 
+    # Average all elements
     if average is True:
-        # Average all elements
         out = {key: np.mean(value) for (key, value) in out.items()}
 
     return out
diff --git a/pyllusion/utilities/analyze_image.py b/pyllusion/utilities/analyze_image.py
index e95215c..a62c42d 100644
--- a/pyllusion/utilities/analyze_image.py
+++ b/pyllusion/utilities/analyze_image.py
@@ -4,7 +4,7 @@
 from .analyze_luminance import analyze_luminance
 
 
-def analyze_image(image):
+def analyze_image(image, features="all"):
     """Compute Objective Characteristics of an Image
 
     Compute the physical characteristics of an image.
@@ -13,6 +13,9 @@ def analyze_image(image):
     ----------
     image : ndarray
         Array for R, G and B channels.
+    features : str
+        Which features to extract. Can be 'all' or a list that can contain 'luminance', 'color',
+        'entropy', 'colorfulness', 'contrast', 'structure'.
 
     Returns
     ----------
@@ -30,7 +33,7 @@ def analyze_image(image):
     >>> # Visualize: plt.imshow(image)
     >>>
     >>> # Compute color
-    >>> out = pyllusion.analyze_image(image)
+    >>> out = pyllusion.analyze_image(image, features="all")
     >>> out["Entropy"]
 
     """
@@ -45,16 +48,26 @@ def analyze_image(image):
             "Please install it first (`pip install scikit-image`).",
         )
 
-    out = analyze_luminance(image, average=True)
-    out.update(analyze_color(image, average=True))
-    out["Entropy"] = skimage.measure.shannon_entropy(image, base=2)
-    # SD of the HUE axis (the colors)
-    out["Colorfulness"] = np.std(skimage.color.rgb2hsv(image)[:, :, 0])
-    # SD of the Luminance axis
-    out["Contrast"] = np.std(skimage.color.rgb2lab(image)[:, :, 0])
-
-    # Edge detection
-    bw = skimage.color.rgb2gray(image)
-    edges = skimage.filters.sobel(bw)  # skimage.filters.roberts or skimage.feature.canny
-    out["Structure"] = skimage.measure.shannon_entropy(edges, base=2)
+    if features == "all":
+        features = ["luminance", "color", "entropy", "colorfulness", "contrast", "structure"]
+
+    out = {}
+    if "luminance" in features:
+        out.update(analyze_luminance(image, average=True))
+    if "color" in features:
+        out.update(analyze_color(image, average=True))
+    if "entropy" in features:
+        out["Entropy"] = skimage.measure.shannon_entropy(image, base=2)
+    if "colorfulness" in features:
+        # SD of the HUE axis (the colors)
+        out["Colorfulness"] = np.std(skimage.color.rgb2hsv(image)[:, :, 0])
+    if "contrast" in features:
+        # SD of the Luminance axis
+        out["Contrast"] = np.std(skimage.color.rgb2lab(image)[:, :, 0])
+
+    if "structure" in features:
+        bw = skimage.color.rgb2gray(image)
+        # Edge detection
+        edges = skimage.filters.sobel(bw)  # skimage.filters.roberts or skimage.feature.canny
+        out["Structure"] = skimage.measure.shannon_entropy(edges, base=2)
     return out