Update app.py

app.py

@@ -211,24 +211,23 @@ def predict_wsi(image):
     num_classes=5
     pred_labels = torch.zeros(1, num_classes+1, image.shape[2], image.shape[3]).cuda()
     counter_tensor = torch.zeros(1, 1, image.shape[2], image.shape[3]).cuda()
-    for i in range(0, image.shape[2] - patch_size + 1, stride):
-        for j in range(0, image.shape[3] - patch_size + 1, stride):
+    for i in range(0, image.shape[2] - patch_size, stride):
+        for j in range(0, image.shape[3] - patch_size, stride):
             i_lowered = min(i, image.shape[2] - patch_size)
             j_lowered = min(j, image.shape[3] - patch_size)
             patch = image[:, :, i_lowered:i_lowered + patch_size, j_lowered:j_lowered + patch_size]
             pred_labels_patch = model(patch.float())
-            update_region_i = i_lowered + (patch_size - generator_output_size) // 2
-            update_region_j = j_lowered + (patch_size - generator_output_size) // 2
+            update_region_i = i_lowered # + (patch_size - generator_output_size) // 2
+            update_region_j = j_lowered # + (patch_size - generator_output_size) // 2
             pred_labels[:, :, update_region_i:update_region_i + generator_output_size,
             update_region_j:update_region_j + generator_output_size] += pred_labels_patch
             counter_tensor[:, :, update_region_i:update_region_i + generator_output_size,
             update_region_j:update_region_j + generator_output_size] += 1
-    pred_labels /= counter_tensor
+    # pred_labels /= counter_tensor
     return pred_labels
 
 @spaces.GPU(duration=120)
 def segment_image(image):
-    # img = Image.open(image_path)
     img = image
     img = np.asarray(img)
     if (np.max(img) > 100):