Merge pull request #16 from normster/master

add corruption category option to get_corruption_names
bethgelab · Dec 3, 2020 · 89b6c41 · 89b6c41
2 parents 4bd4cda + 501c517
commit 89b6c41
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/imagecorruptions/__init__.py b/imagecorruptions/__init__.py
@@ -75,5 +75,13 @@ def get_corruption_names(subset='common'):
         return [f.__name__ for f in corruption_tuple[15:]]
     elif subset == 'all':
         return [f.__name__ for f in corruption_tuple]
+    elif subset == 'noise':
+        return [f.__name__ for f in corruption_tuple[0:3]]
+    elif subset == 'blur':
+        return [f.__name__ for f in corruption_tuple[3:7]]
+    elif subset == 'weather':
+        return [f.__name__ for f in corruption_tuple[7:11]]
+    elif subset == 'digital':
+        return [f.__name__ for f in corruption_tuple[11:15]]
     else:
         raise ValueError("subset must be one of ['common', 'validation', 'all']")