fix bugs with python2

andrewekhalel · Nov 21, 2018 · 772604c · 772604c
1 parent 1cf27cc
commit 772604c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/edafa/BasePredictor.py b/edafa/BasePredictor.py
@@ -99,8 +99,8 @@ def apply_aug(self,img):
 		
 		:returns: a set of augmentations of original image
 		"""
-		shp = img.shape
-		aug_patch = np.zeros((len(self.augs),*shp),dtype=img.dtype)
+		h,w,c = img.shape
+		aug_patch = np.zeros((len(self.augs),h,w,c),dtype=img.dtype)
 		for i,aug in enumerate(self.augs):
 			aug_patch[i] = apply(aug,img,self.bits)
 		return aug_patch

diff --git a/edafa/SegPredictor.py b/edafa/SegPredictor.py
@@ -53,8 +53,8 @@ def _predict_single(self,img,overlap=0):
 
 		:return: prediction on the image
 		"""
-		shp = img.shape[:2]
-		output = np.zeros((*shp,self.out_channels))
+		h,w = img.shape[:2]
+		output = np.zeros((h,w,self.out_channels))
 		times = np.zeros(img.shape[:2])
 
 		img = add_reflections(img,self.in_patch_size,self.out_patch_size)