From 501c51719d811d119ac81798d9888c486dd8b0ef Mon Sep 17 00:00:00 2001 From: Norman Mu Date: Sun, 29 Nov 2020 15:19:43 -0800 Subject: [PATCH] add corruption category option to get_corruption_names --- imagecorruptions/__init__.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/imagecorruptions/__init__.py b/imagecorruptions/__init__.py index e7fa311..351d531 100644 --- 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']")