Fix shift in segmentations
Browse files- Model_Seg.py +1 -6
Model_Seg.py
CHANGED
|
@@ -72,7 +72,6 @@ post_transforms = Compose([
|
|
| 72 |
|
| 73 |
|
| 74 |
def load_and_segment_image(input_image_path, device):
|
| 75 |
-
|
| 76 |
image_tensor = pre_transforms(input_image_path)
|
| 77 |
image_tensor = image_tensor.unsqueeze(0).to(device)
|
| 78 |
|
|
@@ -84,11 +83,7 @@ def load_and_segment_image(input_image_path, device):
|
|
| 84 |
|
| 85 |
outputs = outputs.squeeze(0)
|
| 86 |
|
| 87 |
-
processed_outputs = post_transforms(outputs)
|
| 88 |
-
|
| 89 |
-
# rotate
|
| 90 |
-
rotate = Rotate90(spatial_axes=(0, 1), k=3)
|
| 91 |
-
processed_outputs = rotate(processed_outputs).to('cpu')
|
| 92 |
|
| 93 |
output_array = processed_outputs.squeeze().detach().numpy().astype(np.uint8)
|
| 94 |
|
|
|
|
| 72 |
|
| 73 |
|
| 74 |
def load_and_segment_image(input_image_path, device):
|
|
|
|
| 75 |
image_tensor = pre_transforms(input_image_path)
|
| 76 |
image_tensor = image_tensor.unsqueeze(0).to(device)
|
| 77 |
|
|
|
|
| 83 |
|
| 84 |
outputs = outputs.squeeze(0)
|
| 85 |
|
| 86 |
+
processed_outputs = post_transforms(outputs).to('cpu')
|
|
|
|
|
|
|
|
|
|
|
|
|
| 87 |
|
| 88 |
output_array = processed_outputs.squeeze().detach().numpy().astype(np.uint8)
|
| 89 |
|