replaced strings with enums.val and formatting

cleanlab · Oct 5, 2023 · 78715f7 · 78715f7
1 parent b57b636
commit 78715f7
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 19 deletions.
diff --git a/src/cleanvision/utils/base_issue_manager.py b/src/cleanvision/utils/base_issue_manager.py
@@ -39,7 +39,8 @@ def check_params(**kwargs: Any) -> None:
                 raise ValueError(f"{name} is not a valid keyword argument.")
             if value is not None and not isinstance(value, allowed_kwargs[name]):
                 raise ValueError(
-                    f"Valid type for keyword argument {name} can only be {allowed_kwargs[name]}. {name} cannot be type {type(name)}. "
+                    f"Valid type for keyword argument {name} can only be "
+                    f"{allowed_kwargs[name]}. {name} cannot be type {type(name)}."
                 )
 
     @abstractmethod
@@ -55,7 +56,9 @@ def get_default_params(self) -> Dict[str, Any]:
 
     @abstractmethod
     def update_params(self, params: Dict[str, Any]) -> None:
-        """Sets params for an issue manager. Default params will be overridden by user provided params"""
+        """Sets params for an issue manager.
+        Default params will be overridden by user provided params
+        """
         raise NotImplementedError
 
     @staticmethod

diff --git a/src/cleanvision/utils/constants.py b/src/cleanvision/utils/constants.py
@@ -1,18 +1,34 @@
 from typing import List
 
+from cleanvision.issue_managers import IssueType
+
 IMAGE_PROPERTY: str = "image_property"
 DUPLICATE: str = "duplicate"
 
+# class IssueType(Enum):
+#     DARK = "dark"
+#     LIGHT = "light"
+#     ODD_ASPECT_RATIO = "odd_aspect_ratio"
+#     LOW_INFORMATION = "low_information"
+#     EXACT_DUPLICATES = "exact_duplicates"
+#     NEAR_DUPLICATES = "near_duplicates"
+#     BLURRY = "blurry"
+#     GRAYSCALE = "grayscale"
+#     ODD_SIZE = "odd_size"
+
 IMAGE_PROPERTY_ISSUE_TYPES_LIST: List[str] = [
-    "dark",
-    "light",
-    "odd_aspect_ratio",
-    "low_information",
-    "blurry",
-    "grayscale",
-    "odd_size",
+    IssueType.DARK.value,  # "dark",
+    IssueType.LIGHT.value,  # "light",
+    IssueType.ODD_ASPECT_RATIO.value,  # "odd_aspect_ratio",
+    IssueType.LOW_INFORMATION.value,  # "low_information",
+    IssueType.BLURRY.value,  # "blurry",
+    IssueType.GRAYSCALE.value,  # "grayscale",
+    IssueType.ODD_SIZE.value,  # "odd_size",
+]
+DUPLICATE_ISSUE_TYPES_LIST: List[str] = [
+    IssueType.EXACT_DUPLICATES.value,  # "exact_duplicates"
+    IssueType.NEAR_DUPLICATES.value,  #  "near_duplicates"
 ]
-DUPLICATE_ISSUE_TYPES_LIST: List[str] = ["exact_duplicates", "near_duplicates"]
 SETS: str = "sets"
 
 # max number of processes that can be forked/spawned for multiprocessing
@@ -38,13 +54,22 @@
 ]  # filetypes supported by PIL
 
 DEFAULT_ISSUE_TYPES = [
-    "dark",
-    "light",
-    "odd_aspect_ratio",
-    "low_information",
-    "exact_duplicates",
-    "near_duplicates",
-    "blurry",
-    "grayscale",
-    "odd_size",
+    IssueType.DARK.value,  # "dark",
+    IssueType.LIGHT.value,  # "light",
+    IssueType.ODD_ASPECT_RATIO.value,  # "odd_aspect_ratio",
+    IssueType.LOW_INFORMATION.value,  # "low_information",
+    IssueType.EXACT_DUPLICATES.value,  # "exact_duplicates"
+    IssueType.NEAR_DUPLICATES.value,  # "near_duplicates"
+    IssueType.BLURRY.value,  # "blurry",
+    IssueType.GRAYSCALE.value,  # "grayscale",
+    IssueType.ODD_SIZE.value,  # "odd_size",
+    # "dark",
+    # "light",
+    # "odd_aspect_ratio",
+    # "low_information",
+    # "exact_duplicates",
+    # "near_duplicates",
+    # "blurry",
+    # "grayscale",
+    # "odd_size",
 ]