-
Notifications
You must be signed in to change notification settings - Fork 16
Expand file tree
/
Copy pathsimpledemo2d.py
More file actions
109 lines (82 loc) · 2.63 KB
/
simpledemo2d.py
File metadata and controls
109 lines (82 loc) · 2.63 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import FastGeodis
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
# FastGeodis Method
device = "cuda" if torch.cuda.is_available() else "cpu"
image = np.asarray(Image.open("data/img2d.png"), np.float32)
image_pt = torch.from_numpy(image).unsqueeze_(0).unsqueeze_(0)
image_pt = image_pt.to(device)
mask_pt = torch.ones_like(image_pt)
mask_pt[..., 100, 100] = 0
v = 1e10
iterations = 2
lamb = 1.0 # <-- Geodesic distance transform
geodesic_dist = FastGeodis.generalised_geodesic2d(
image_pt, mask_pt, v, lamb, iterations
)
geodesic_dist = np.squeeze(geodesic_dist.cpu().numpy())
lamb = 0.0 # <-- Euclidean distance transform
euclidean_dist = FastGeodis.generalised_geodesic2d(
image_pt, mask_pt, v, lamb, iterations
)
euclidean_dist = np.squeeze(euclidean_dist.cpu().numpy())
plt.figure(figsize=(12, 4))
plt.subplot(1, 3, 1)
plt.imshow(image)
plt.subplot(1, 3, 2)
plt.imshow(geodesic_dist)
plt.plot(100, 100, "mo")
plt.subplot(1, 3, 3)
plt.imshow(euclidean_dist)
plt.plot(100, 100, "mo")
plt.show()
# Toivanen's Raster Method
# Toivanen's method only support CPU
image_pt = image_pt.to("cpu")
mask_pt = mask_pt.to("cpu")
lamb = 1.0 # <-- Geodesic distance transform
geodesic_dist_toivanen = FastGeodis.generalised_geodesic2d_toivanen(
image_pt, mask_pt, v, lamb, iterations
)
geodesic_dist_toivanen = np.squeeze(geodesic_dist_toivanen.cpu().numpy())
lamb = 0.0 # <-- Euclidean distance transform
euclidean_dist_toivanen = FastGeodis.generalised_geodesic2d_toivanen(
image_pt, mask_pt, v, lamb, iterations
)
euclidean_dist_toivanen = np.squeeze(euclidean_dist_toivanen.cpu().numpy())
plt.figure(figsize=(12, 4))
plt.subplot(1, 3, 1)
plt.imshow(image)
plt.subplot(1, 3, 2)
plt.imshow(geodesic_dist_toivanen)
plt.plot(100, 100, "mo")
plt.subplot(1, 3, 3)
plt.imshow(euclidean_dist_toivanen)
plt.plot(100, 100, "mo")
plt.show()
# Fast Marching Method
# Fast Marching method only support CPU
image_pt = image_pt.to("cpu")
mask_pt = mask_pt.to("cpu")
lamb = 1.0 # <-- Geodesic distance transform
geodesic_dist_fastmarch = FastGeodis.geodesic2d_fastmarch(
image_pt, mask_pt, lamb
)
geodesic_dist_fastmarch = np.squeeze(geodesic_dist_fastmarch.cpu().numpy())
lamb = 0.0 # <-- Euclidean distance transform
euclidean_dist_fastmarch = FastGeodis.geodesic2d_fastmarch(
image_pt, mask_pt, lamb
)
euclidean_dist_fastmarch = np.squeeze(euclidean_dist_fastmarch.cpu().numpy())
plt.figure(figsize=(12, 4))
plt.subplot(1, 3, 1)
plt.imshow(image)
plt.subplot(1, 3, 2)
plt.imshow(geodesic_dist_fastmarch)
plt.plot(100, 100, "mo")
plt.subplot(1, 3, 3)
plt.imshow(euclidean_dist_fastmarch)
plt.plot(100, 100, "mo")
plt.show()