Xiao Furen
264e17b13a
- Implemented xfr_pso.py to handle PSO optimization tasks for lumbar levels, including GPU management and logging. - Introduced xfr_record_planes.py to supplement existing volumes with symmetry and endplate plane data in transform.json. - Added functionality to skip already recorded levels unless forced, ensuring efficient processing. - Enhanced logging for better traceability of operations and errors during execution.
450 lines
No EOL
14 KiB
Python
450 lines
No EOL
14 KiB
Python
|
||
from scipy.spatial import KDTree
|
||
|
||
import torch
|
||
import numpy as np
|
||
|
||
from config.constant import ALLOWED_DIAMETERS, ALLOWED_LENGTHS
|
||
|
||
def create_coordinate_grid(
|
||
shape: tuple[int, int, int],
|
||
device: torch.device,
|
||
dtype: torch.dtype = torch.float32
|
||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||
"""
|
||
建立固定的 3D voxel coordinate grid
|
||
returns:
|
||
z_t, y_t, x_t with shape = (Z, Y, X)
|
||
"""
|
||
z_t = torch.arange(shape[0], device=device, dtype=dtype)
|
||
y_t = torch.arange(shape[1], device=device, dtype=dtype)
|
||
x_t = torch.arange(shape[2], device=device, dtype=dtype)
|
||
|
||
z_t, y_t, x_t = torch.meshgrid(z_t, y_t, x_t, indexing='ij')
|
||
return z_t, y_t, x_t
|
||
|
||
def snap_to_discrete_values(diameter_raw, length_raw):
|
||
"""
|
||
將連續值映射到最接近的允許離散值
|
||
|
||
Parameters:
|
||
diameter_raw: PSO 給的連續直徑值
|
||
length_raw: PSO 給的連續長度值
|
||
|
||
Returns:
|
||
diameter_discrete, length_discrete: 離散化後的值
|
||
"""
|
||
# 找最接近的 diameter
|
||
diameter_discrete = min(ALLOWED_DIAMETERS, key=lambda x: abs(x - diameter_raw))
|
||
|
||
# 找最接近的 length
|
||
length_discrete = min(ALLOWED_LENGTHS, key=lambda x: abs(x - length_raw))
|
||
|
||
return diameter_discrete, length_discrete
|
||
|
||
points = (
|
||
(4.5, 25),
|
||
(4.5, 30),
|
||
(4.5, 35),
|
||
|
||
(5.0, 30),
|
||
(5.0, 35),
|
||
(5.0, 40),
|
||
|
||
(5.5, 30),
|
||
(5.5, 35),
|
||
(5.5, 40),
|
||
(5.5, 45),
|
||
(5.5, 50),
|
||
|
||
(6.0, 35),
|
||
(6.0, 40),
|
||
(6.0, 45),
|
||
(6.0, 50),
|
||
(6.0, 55),
|
||
|
||
(6.5, 35),
|
||
(6.5, 40),
|
||
(6.5, 45),
|
||
(6.5, 50),
|
||
(6.5, 55),
|
||
|
||
(7.0, 35),
|
||
(7.0, 40),
|
||
(7.0, 45),
|
||
(7.0, 50),
|
||
|
||
(7.5, 60),
|
||
(7.5, 70),
|
||
(7.5, 80),
|
||
)
|
||
|
||
tree = KDTree(points)
|
||
|
||
def snap_to_discrete_values_xfr(diameter_raw, length_raw):
|
||
return snap_to_discrete_values(diameter_raw, length_raw)
|
||
|
||
distance, index = tree.query((diameter_raw, length_raw))
|
||
return points[index][0], points[index][1]
|
||
|
||
def round_down_to_discrete_values_xfr(diameter_raw, length_raw):
|
||
"""
|
||
將連續值映射到小於等於該值的最接近允許離散值 (向下取整)
|
||
|
||
Parameters:
|
||
diameter_raw: PSO 給的連續直徑值
|
||
length_raw: PSO 給的連續長度值
|
||
|
||
Returns:
|
||
diameter_discrete, length_discrete: 離散化後的值
|
||
"""
|
||
# 找小於等於 raw 的最接近 diameter，如果都大於 raw 則取最小值
|
||
valid_diameters = [d for d in ALLOWED_DIAMETERS if d <= diameter_raw]
|
||
diameter_discrete = max(valid_diameters) if valid_diameters else min(ALLOWED_DIAMETERS)
|
||
|
||
# 找小於等於 raw 的最接近 length，如果都大於 raw 則取最小值
|
||
valid_lengths = [l for l in ALLOWED_LENGTHS if l <= length_raw]
|
||
length_discrete = max(valid_lengths) if valid_lengths else min(ALLOWED_LENGTHS)
|
||
|
||
return diameter_discrete, length_discrete
|
||
|
||
|
||
def generate_cylinder_n_torch(
|
||
diameter: float,
|
||
length: float,
|
||
position_z: float,
|
||
position_y: float,
|
||
position_x: float,
|
||
azimuth: float,
|
||
altitude: float,
|
||
shape: tuple[int, int, int],
|
||
spacing: list[float],
|
||
device: torch.device,
|
||
grid: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
|
||
) -> torch.Tensor:
|
||
|
||
"""
|
||
Generate a "forward" (positive z) cylinder mask in 3D space using PyTorch tensors.
|
||
Returns a binary mask (torch.uint8) on the specified device.
|
||
"""
|
||
|
||
if grid is None:
|
||
z_t, y_t, x_t = create_coordinate_grid(shape, device)
|
||
else:
|
||
z_t, y_t, x_t = grid
|
||
|
||
azimuth_rad_t = torch.deg2rad(torch.tensor(azimuth, device=device, dtype=torch.float32))
|
||
altitude_rad_t = torch.deg2rad(torch.tensor(altitude, device=device, dtype=torch.float32))
|
||
|
||
# Shift
|
||
z_t = z_t - position_z
|
||
y_t = y_t - position_y
|
||
x_t = x_t - position_x
|
||
|
||
# Apply rotation (same formula as your NumPy version, but in torch)
|
||
x_rot = (
|
||
x_t * torch.cos(azimuth_rad_t) * torch.cos(altitude_rad_t)
|
||
+ y_t * torch.sin(azimuth_rad_t) * torch.cos(altitude_rad_t)
|
||
- z_t * torch.sin(altitude_rad_t)
|
||
)
|
||
y_rot = -x_t * torch.sin(azimuth_rad_t) + y_t * torch.cos(azimuth_rad_t)
|
||
z_rot = (
|
||
x_t * torch.cos(azimuth_rad_t) * torch.sin(altitude_rad_t)
|
||
+ y_t * torch.sin(azimuth_rad_t) * torch.sin(altitude_rad_t)
|
||
+ z_t * torch.cos(altitude_rad_t)
|
||
)
|
||
|
||
# Handle spacing
|
||
# You can expand or generalize for more spacing options
|
||
if spacing == [1, 1, 1]:
|
||
radius = diameter / 2.0
|
||
mask = (
|
||
(x_rot**2 + y_rot**2 <= radius**2)
|
||
& (z_rot >= 0)
|
||
& (z_rot <= length)
|
||
)
|
||
elif spacing == [0.5, 0.5, 0.5]:
|
||
radius = (diameter / 2.0) * 2
|
||
mask = (
|
||
(x_rot**2 + y_rot**2 <= radius**2)
|
||
& (z_rot >= 0)
|
||
& (z_rot <= length * 2)
|
||
)
|
||
else:
|
||
raise ValueError(f"Unsupported spacing: {spacing}")
|
||
|
||
# Convert boolean mask to uint8
|
||
cylinder_mask = mask.to(torch.uint8)
|
||
|
||
return cylinder_mask
|
||
|
||
def generate_cylinder_o_torch(
|
||
diameter: float,
|
||
length: float,
|
||
position_z: float,
|
||
position_y: float,
|
||
position_x: float,
|
||
azimuth: float,
|
||
altitude: float,
|
||
shape: tuple[int, int, int],
|
||
spacing: list[float],
|
||
device: torch.device,
|
||
grid: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
|
||
butt_mm: float = 5.0
|
||
) -> torch.Tensor:
|
||
"""
|
||
Generate an "opposite" (negative z) cylinder mask in 3D space using PyTorch tensors.
|
||
Returns a binary mask (torch.uint8) on the specified device.
|
||
|
||
butt_mm: 入口端（z_rot=0，螺絲 butt 面）附近前 butt_mm 不計入此 mask，
|
||
讓 bone surface 附近允許最多 butt_mm 的 miss（入口端略陷骨面內
|
||
該帶 voxel 不再被 scoring 的 cyl_o∩bone 懲罰）。
|
||
"""
|
||
|
||
if grid is None:
|
||
z_t, y_t, x_t = create_coordinate_grid(shape, device)
|
||
else:
|
||
z_t, y_t, x_t = grid
|
||
|
||
# Convert angles to torch
|
||
azimuth_rad_t = torch.deg2rad(torch.tensor(azimuth, device=device, dtype=torch.float32))
|
||
altitude_rad_t = torch.deg2rad(torch.tensor(altitude, device=device, dtype=torch.float32))
|
||
|
||
# Shift
|
||
z_t = z_t - position_z
|
||
y_t = y_t - position_y
|
||
x_t = x_t - position_x
|
||
|
||
# Apply rotation
|
||
x_rot = (
|
||
x_t * torch.cos(azimuth_rad_t) * torch.cos(altitude_rad_t)
|
||
+ y_t * torch.sin(azimuth_rad_t) * torch.cos(altitude_rad_t)
|
||
- z_t * torch.sin(altitude_rad_t)
|
||
)
|
||
y_rot = -x_t * torch.sin(azimuth_rad_t) + y_t * torch.cos(azimuth_rad_t)
|
||
z_rot = (
|
||
x_t * torch.cos(azimuth_rad_t) * torch.sin(altitude_rad_t)
|
||
+ y_t * torch.sin(azimuth_rad_t) * torch.sin(altitude_rad_t)
|
||
+ z_t * torch.cos(altitude_rad_t)
|
||
)
|
||
|
||
# Handle spacing
|
||
if spacing == [1, 1, 1]:
|
||
radius = diameter / 2.0
|
||
mask = (
|
||
(x_rot**2 + y_rot**2 <= radius**2)
|
||
& (z_rot <= -butt_mm)
|
||
& (z_rot >= -length)
|
||
)
|
||
elif spacing == [0.5, 0.5, 0.5]:
|
||
radius = (diameter / 2.0) * 2
|
||
mask = (
|
||
(x_rot**2 + y_rot**2 <= radius**2)
|
||
& (z_rot <= -butt_mm * 2)
|
||
& (z_rot >= -length * 2)
|
||
)
|
||
else:
|
||
raise ValueError(f"Unsupported spacing: {spacing}")
|
||
|
||
cylinder_mask_o = mask.to(torch.uint8)
|
||
|
||
return cylinder_mask_o
|
||
|
||
def generate_cylinder_numpy(diameter, length, position_z, position_y, position_x, azimuth, altitude, shape, spacing):
|
||
|
||
azimuth = np.radians(azimuth)
|
||
altitude = np.radians(altitude)
|
||
|
||
cylinder_mask = np.zeros(shape, dtype=np.uint8)
|
||
|
||
z, y, x = np.mgrid[0:shape[0], 0:shape[1], 0:shape[2]].astype(np.float64)
|
||
|
||
z -= float(position_z)
|
||
y -= float(position_y)
|
||
x -= float(position_x)
|
||
|
||
x_rot = x * np.cos(azimuth) * np.cos(altitude) + y * np.sin(azimuth) * np.cos(altitude) - z * np.sin(altitude)
|
||
y_rot = -x * np.sin(azimuth) + y * np.cos(azimuth)
|
||
z_rot = x * np.cos(azimuth) * np.sin(altitude) + y * np.sin(azimuth) * np.sin(altitude) + z * np.cos(altitude)
|
||
|
||
if spacing == [1, 1, 1]:
|
||
|
||
radius = diameter / 2.0
|
||
cylinder = (x_rot**2 + y_rot**2 <= radius**2) & (z_rot >= 0) & (z_rot <= length)
|
||
|
||
elif spacing == [0.5, 0.5, 0.5]:
|
||
radius = diameter / 2.0 * 2 # *2 is for resampling
|
||
cylinder = (x_rot**2 + y_rot**2 <= radius**2) & (z_rot >= 0) & (z_rot <= length * 2)
|
||
|
||
cylinder_mask[cylinder] = 1
|
||
|
||
return cylinder_mask
|
||
|
||
def generate_cylinder_butt_torch(
|
||
diameter,
|
||
position_z, position_y, position_x,
|
||
azimuth, altitude,
|
||
shape, spacing, device, grid=None,
|
||
butt_mm=5.0 # 入口端最後 butt_mm（mm）：but 豁免帶寬度
|
||
) -> torch.Tensor:
|
||
"""生成圓柱「最後端」mask：入口端（z_rot=0，遠離 VBODY 的一端）
|
||
最靠近的 butt_mm 圓柱短柱（與 generate_cylinder_n_torch 同慣例、同直徑）。
|
||
尖端在 z_rot=length（靠 VBODY 端），不在此 mask。
|
||
此 mask 是完整圓柱的子集，scoring 中 mask 內 not_in_bone 的 voxel
|
||
豁免 10000/voxel 扣分與 HARD_BREACH_PER_VOX 硬約束（螺絲入口端允許
|
||
伸出骨外 butt_mm）。"""
|
||
|
||
if grid is None:
|
||
z_t, y_t, x_t = create_coordinate_grid(shape, device)
|
||
else:
|
||
z_t, y_t, x_t = grid
|
||
|
||
azimuth_rad_t = torch.deg2rad(torch.tensor(azimuth, device=device, dtype=torch.float32))
|
||
altitude_rad_t = torch.deg2rad(torch.tensor(altitude, device=device, dtype=torch.float32))
|
||
|
||
z_t = z_t - position_z
|
||
y_t = y_t - position_y
|
||
x_t = x_t - position_x
|
||
|
||
x_rot = (
|
||
x_t * torch.cos(azimuth_rad_t) * torch.cos(altitude_rad_t)
|
||
+ y_t * torch.sin(azimuth_rad_t) * torch.cos(altitude_rad_t)
|
||
- z_t * torch.sin(altitude_rad_t)
|
||
)
|
||
y_rot = -x_t * torch.sin(azimuth_rad_t) + y_t * torch.cos(azimuth_rad_t)
|
||
z_rot = (
|
||
x_t * torch.cos(azimuth_rad_t) * torch.sin(altitude_rad_t)
|
||
+ y_t * torch.sin(azimuth_rad_t) * torch.sin(altitude_rad_t)
|
||
+ z_t * torch.cos(altitude_rad_t)
|
||
)
|
||
|
||
# 與 generate_cylinder_n_torch 相同的 spacing/單位處理；
|
||
# 長度固定為 butt_mm（mm）→ voxel
|
||
if spacing == [1, 1, 1]:
|
||
radius = diameter / 2.0
|
||
butt_len = butt_mm
|
||
elif spacing == [0.5, 0.5, 0.5]:
|
||
radius = (diameter / 2.0) * 2
|
||
butt_len = butt_mm * 2
|
||
else:
|
||
raise ValueError(f"Unsupported spacing: {spacing}")
|
||
|
||
mask = (
|
||
(x_rot**2 + y_rot**2 <= radius**2)
|
||
& (z_rot >= 0)
|
||
& (z_rot < butt_len)
|
||
)
|
||
|
||
return mask.to(torch.uint8)
|
||
|
||
def generate_cone_torch(
|
||
diameter: float,
|
||
length: float,
|
||
position_z: float,
|
||
position_y: float,
|
||
position_x: float,
|
||
azimuth: float,
|
||
altitude: float,
|
||
shape: tuple[int, int, int],
|
||
spacing: list[float],
|
||
device: torch.device,
|
||
grid: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
|
||
) -> torch.Tensor:
|
||
"""螺絲圖用的錐體 mask（與 generate_cylinder_n_torch 同慣例/同座標系）：
|
||
入口端（z_rot=0）半徑 = diameter/2，沿 +z 線性收細到尖端
|
||
（z_rot=length，半徑=0）。純顯示用（render_bone_figure 的螺絲體），
|
||
評分維持圓柱、不經此 mask。"""
|
||
|
||
if grid is None:
|
||
z_t, y_t, x_t = create_coordinate_grid(shape, device)
|
||
else:
|
||
z_t, y_t, x_t = grid
|
||
|
||
azimuth_rad_t = torch.deg2rad(torch.tensor(azimuth, device=device, dtype=torch.float32))
|
||
altitude_rad_t = torch.deg2rad(torch.tensor(altitude, device=device, dtype=torch.float32))
|
||
|
||
z_t = z_t - position_z
|
||
y_t = y_t - position_y
|
||
x_t = x_t - position_x
|
||
|
||
x_rot = (
|
||
x_t * torch.cos(azimuth_rad_t) * torch.cos(altitude_rad_t)
|
||
+ y_t * torch.sin(azimuth_rad_t) * torch.cos(altitude_rad_t)
|
||
- z_t * torch.sin(altitude_rad_t)
|
||
)
|
||
y_rot = -x_t * torch.sin(azimuth_rad_t) + y_t * torch.cos(azimuth_rad_t)
|
||
z_rot = (
|
||
x_t * torch.cos(azimuth_rad_t) * torch.sin(altitude_rad_t)
|
||
+ y_t * torch.sin(azimuth_rad_t) * torch.sin(altitude_rad_t)
|
||
+ z_t * torch.cos(altitude_rad_t)
|
||
)
|
||
|
||
if spacing == [1, 1, 1]:
|
||
radius = diameter / 2.0
|
||
total_len = length
|
||
elif spacing == [0.5, 0.5, 0.5]:
|
||
radius = (diameter / 2.0) * 2
|
||
total_len = length * 2
|
||
else:
|
||
raise ValueError(f"Unsupported spacing: {spacing}")
|
||
|
||
taper = torch.clamp(z_rot / total_len, min=0.0, max=1.0)
|
||
r = radius * (1.0 - taper)
|
||
|
||
mask = (
|
||
(x_rot**2 + y_rot**2 <= r**2)
|
||
& (z_rot >= 0)
|
||
& (z_rot <= total_len)
|
||
)
|
||
return mask.to(torch.uint8)
|
||
|
||
|
||
def generate_cylinder_tip_torch(
|
||
diameter, length,
|
||
position_z, position_y, position_x,
|
||
azimuth, altitude,
|
||
shape, spacing, device, grid=None,
|
||
tip_ratio=0.2 # 取末端 20% 當尖端
|
||
) -> torch.Tensor:
|
||
"""只生成圓柱末端的 mask"""
|
||
|
||
if grid is None:
|
||
z_t, y_t, x_t = create_coordinate_grid(shape, device)
|
||
else:
|
||
z_t, y_t, x_t = grid
|
||
|
||
azimuth_rad_t = torch.deg2rad(torch.tensor(azimuth, device=device, dtype=torch.float32))
|
||
altitude_rad_t = torch.deg2rad(torch.tensor(altitude, device=device, dtype=torch.float32))
|
||
|
||
z_t = z_t - position_z
|
||
y_t = y_t - position_y
|
||
x_t = x_t - position_x
|
||
|
||
x_rot = (
|
||
x_t * torch.cos(azimuth_rad_t) * torch.cos(altitude_rad_t)
|
||
+ y_t * torch.sin(azimuth_rad_t) * torch.cos(altitude_rad_t)
|
||
- z_t * torch.sin(altitude_rad_t)
|
||
)
|
||
y_rot = -x_t * torch.sin(azimuth_rad_t) + y_t * torch.cos(azimuth_rad_t)
|
||
z_rot = (
|
||
x_t * torch.cos(azimuth_rad_t) * torch.sin(altitude_rad_t)
|
||
+ y_t * torch.sin(azimuth_rad_t) * torch.sin(altitude_rad_t)
|
||
+ z_t * torch.cos(altitude_rad_t)
|
||
)
|
||
|
||
if spacing == [0.5, 0.5, 0.5]:
|
||
radius = (diameter / 2.0) * 2
|
||
total_length = length * 2
|
||
else:
|
||
radius = diameter / 2.0
|
||
total_length = length
|
||
|
||
tip_start = total_length * (1 - tip_ratio) # 末端 20% 開始的位置
|
||
|
||
mask = (
|
||
(x_rot**2 + y_rot**2 <= radius**2)
|
||
& (z_rot >= tip_start) # 只取末端
|
||
& (z_rot <= total_length)
|
||
)
|
||
|
||
return mask.to(torch.uint8) |