Fix test fail: Due to recompute_scale_factor missing in UpsamplingBil…

…inear2d (#68) * Fix test fail: Due to recompute_scale_factor missing in UpsamplingBilinear2d * Add comments to explain the changes
DevoLearn · Aug 3, 2022 · e3c4a1c · e3c4a1c
1 parent 904f454
commit e3c4a1c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/devolearn/cell_membrane_segmentor/cell_membrane_segmentor.py b/devolearn/cell_membrane_segmentor/cell_membrane_segmentor.py
@@ -138,6 +138,11 @@ def predict(self, image_path, pred_size = (350,250), centroid_mode = False):
         im = cv2.imread(image_path,0)
 
         tensor = self.preprocess(im)
+
+        # The model has issues with the latest PyTorch versions, which causes
+        # AttributeError: 'Upsample' object has no attribute 'recompute_scale_factor'
+        # To avoid this, we manually set the recompute_scale_factor attribute to None
+        self.model.segmentation_head[1].recompute_scale_factor = None
         res = self.model(tensor).detach().cpu().numpy()[0][0]
 
         res = cv2.resize(res,pred_size)

diff --git a/devolearn/cell_nucleus_segmentor/cell_nucleus_segmentor.py b/devolearn/cell_nucleus_segmentor/cell_nucleus_segmentor.py
@@ -98,6 +98,12 @@ def predict(self, image_path, pred_size = (350,250)):
 
         im = cv2.imread(image_path,0)
         tensor = self.preprocess(im)
+
+        # The model has issues with the latest PyTorch versions, which causes
+        # AttributeError: 'Upsample' object has no attribute 'recompute_scale_factor'
+        # To avoid this, we manually set the recompute_scale_factor attribute to None
+        self.model.segmentation_head[1].recompute_scale_factor = None
+
         res = self.model(tensor).detach().cpu().numpy()[0][0]
         res = cv2.resize(res,pred_size)
         return res