update

inference.py

@@ -28,7 +28,7 @@ class Inference(object):
     def inference(self, image, modality):
         assert modality in self.modality_mapping, "Modality '{}' not supported".format(modality)
         
-        image = self.load_image(image)
+        image, raw_h, raw_w = self.load_image(image)
         modality_idx = self.modality_mapping[modality]
         modality_idx = torch.tensor([modality_idx])
         with torch.no_grad():
@@ -36,19 +36,20 @@ class Inference(object):
         output = output.data.cpu().numpy()[0][0]
         output = sigmoid(output) * 255
         output = output.astype(np.uint8)
+        output = cv2.resize(output, (raw_w, raw_h))
         return output
         
     def load_image(self, image):
         # Load the image and preprocess it
         if isinstance(image, str):
             image = cv2.imread(image)[:, :, [2, 1, 0]]
-        #image = image
+        raw_h, raw_w = image.shape[:2]
         image = cv2.resize(image, (self.model_params['size_w'], self.model_params['size_h']))
         image = image.astype(np.float32) / 255.0
         image = np.transpose(image, (2, 0, 1))
         image = np.expand_dims(image, axis=0)
         image = torch.tensor(image)
-        return image
+        return image, raw_h, raw_w
     
     def load_model(self):
         print('Loading model from {}'.format(self.model_path))