From e3c4a1c9ed5d56b0ece3be2d38c709f9c429e0de Mon Sep 17 00:00:00 2001
From: Geo Jolly <69108486+kingjuno@users.noreply.github.com>
Date: Wed, 3 Aug 2022 20:22:03 +0530
Subject: [PATCH] Fix test fail: Due to recompute_scale_factor missing in
 UpsamplingBilinear2d (#68)

* Fix test fail: Due to recompute_scale_factor missing in UpsamplingBilinear2d

* Add comments to explain the changes
---
 .../cell_membrane_segmentor/cell_membrane_segmentor.py      | 5 +++++
 devolearn/cell_nucleus_segmentor/cell_nucleus_segmentor.py  | 6 ++++++
 2 files changed, 11 insertions(+)

diff --git a/devolearn/cell_membrane_segmentor/cell_membrane_segmentor.py b/devolearn/cell_membrane_segmentor/cell_membrane_segmentor.py
index f772eaa..093f3f9 100755
--- 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
index 9641141..d1e57c2 100755
--- 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