fix(preprocessor): add numpy transform

gnes-ai · Oct 15, 2019 · 24e41be · 24e41be
1 parent cfee124
commit 24e41be
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/gnes/preprocessor/helper.py b/gnes/preprocessor/helper.py
@@ -388,9 +388,12 @@ def detect_peak_boundary(distances: List[float],
 
 
 def torch_transform(img):
-    import torchvision.transforms as transforms
-    return transforms.Compose([transforms.ToTensor(),
-                               transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])(img)
+    try:
+        import torchvision.transforms as transforms
+        return transforms.Compose([transforms.ToTensor(),
+                                   transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])(img)
+    except (ImportError, ModuleNotFoundError):
+        return np.asarray(img) * 2 / 255. - 1.
 
 
 def get_all_subarea(img):