refactor(core): improve cylinder parameter snapping and optimization logic

Refactor the cylinder parameter selection and optimization process to
improve accuracy and robustness.

- Implement `snap_to_discrete_values_xfr` using a KDTree for efficient
  mapping of continuous diameter and length values to a predefined
  set of discrete points.
- Update `objective_function_xfr` to use the new snapping mechanism
  and introduce a weighted loss component for diameter and length.
- Adjust PSO optimization bounds and search ranges in `run_pso_torch_xfr`
  to better align with image dimensions and anatomical constraints.
- Refine scoring logic in `cl_score_torch_xfr` with updated penalty
  weights for overlaps and out-of-bone voxels.
- Update `config/constant.py` with new allowed diameter and length
  ranges.
- Improve `imaging/preprocessing.py` by making `PROGRESS_FILE` a
  parameter to allow per-output-directory progress tracking.
- Update `xfr_debug.py` with improved error handling and directory
  paths for batch processing.

.gitignore

@@ -214,3 +214,5 @@ __marimo__/
 
 # Streamlit
 .streamlit/secrets.toml
+
+progress.json

.kilo/agents/data.md

@@ -0,0 +1,28 @@
+---
+mode: primary
+description: Run notebook-first data analysis by appending and executing cells
+  for each request.
+options:
+  displayName: Data
+  id: data
+requirements:
+  skills:
+    - data-investigation
+  vscode_extensions:
+    - name: Jupyter
+      id: ms-toolsai.jupyter
+color: "#2563EB"
+---
+
+You are Kilo, a notebook-first data analysis agent. Use an active Jupyter notebook as the working surface.
+
+Guidelines:
+- If no notebook is active, create a uniquely named, descriptive `<topic>.ipynb` in the current workspace folder
+- Use the dedicated notebook tools to create, read, edit, and execute; prefer these tools over other methods like MCP tools and manual raw JSON editing
+- Confirm Jupyter and kernel readiness through the first requested notebook execution; only notify the user if they need to select or configure a kernel before work can continue
+- For every user request, append at least one focused code cell and execute it
+- Preserve notebook history: do not modify or delete existing cells unless explicitly asked; after failures, append diagnostic or corrected cells
+- Keep substantive data work and supporting evidence in the notebook
+- Avoid changing non-notebook files unless explicitly requested or necessary to complete the task
+- Inspect cell output before answering, and keep notebook outputs and final summaries concise
+- Never claim execution when a notebook cell did not run

.vscode/settings.json

@@ -1,5 +1,7 @@
 {
     "python.analysis.extraPaths": [
         "${workspaceFolder}"
-    ]
+    ],
+    "python-envs.defaultEnvManager": "ms-python.python:conda",
+    "python-envs.defaultPackageManager": "ms-python.python:conda"
 }

config/constant.py

@@ -13,16 +13,26 @@ LABEL_MAP = {
     20: "L1", 21: "L2", 22: "L3", 23: "L4", 24: "L5"
 }
 ALLOWED_DIAMETERS = [
-    3.5, 
+    # 3.5, 
-    4.0, 
+    # 4.0, 
     4.5, 
     5.0,
+    5.5,
+    6.0,
+    6.5,
+    # 7.0,
+    # 7.5,
     ]
 ALLOWED_LENGTHS = [
+    25,
+    30,
     35, 
     40, 
     45, 
-    50,
+    # 50,
+    # 60,
+    # 70,
+    # 80,
     ]
 OVERLAP_THRESH = 0.50
 DEFAULT_SPACING = [0.5, 0.5, 0.5]

core/cylinder.py

@@ -1,5 +1,9 @@
+
+from scipy.spatial import KDTree
+
 import torch
 import numpy as np
+
 from config.constant import ALLOWED_DIAMETERS, ALLOWED_LENGTHS
 
 def create_coordinate_grid(
@@ -38,7 +42,52 @@ def snap_to_discrete_values(diameter_raw, length_raw):
     
     return diameter_discrete, length_discrete
 
-def round_down_to_discrete_values(diameter_raw, length_raw):
+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):
     """
     將連續值映射到小於等於該值的最接近允許離散值 (向下取整)
     

core/objective.py

@@ -5,7 +5,7 @@ from scipy.ndimage import map_coordinates
 import numpy as np
 import torch
 
-from core.cylinder import generate_cylinder_n_torch, generate_cylinder_o_torch, snap_to_discrete_values, generate_cylinder_tip_torch
+from core.cylinder import generate_cylinder_n_torch, generate_cylinder_o_torch, snap_to_discrete_values, generate_cylinder_tip_torch, snap_to_discrete_values_xfr
 from core.intersection import center_line_intersections_torch
 from core.scoring import cl_score_torch, cl_score_torch_xfr
 
@@ -138,11 +138,14 @@ def objective_function_xfr(params: list[float], y_indices) -> float:
 
     # 將連續值轉換為離散值
     # diameter_discrete, length_discrete = snap_to_discrete_values(diameter_raw, length_raw)
-    diameter_discrete, length_discrete = diameter_raw, length_raw
+    diameter_discrete, length_discrete = snap_to_discrete_values_xfr(diameter_raw, length_raw)
+
+    diameter_loss = .9*diameter_discrete + .1*diameter_raw
+    length_loss   = .9*  length_discrete + .1*  length_raw
     
     loss = cylinder_circle_line_intersection_loss_deductions_torch(
-        diameter_discrete,
+        diameter_loss,
-        length_discrete,
+        length_loss,
         position_params,
         image2_shape,
         cortical_tensor,
@@ -164,7 +167,6 @@ def objective_function(params: list[float]) -> float:
     
     # 將連續值轉換為離散值
     diameter_discrete, length_discrete = snap_to_discrete_values(diameter_raw, length_raw)
-    # diameter_discrete, length_discrete = diameter_raw, length_raw
     
     loss = cylinder_circle_line_intersection_loss_deductions_torch(
         diameter_discrete,

core/optimizer.py

@@ -7,13 +7,15 @@ from config.constant import ALLOWED_DIAMETERS, ALLOWED_LENGTHS
 from core.objective import objective_function, objective_function_xfr
 from pyswarm import pso 
 import core.objective  # <--- 加入這行，讓我們可以直接操作 objective 模組
-from core.cylinder import generate_cylinder_n_torch, snap_to_discrete_values, create_coordinate_grid, round_down_to_discrete_values
+from core.cylinder import generate_cylinder_n_torch, snap_to_discrete_values, create_coordinate_grid, snap_to_discrete_values_xfr
 from core.scoring import compute_overlap_ratio_from_cylinder_mask, is_solution_ok
 from config.constant import OVERLAP_THRESH
 from visualization.res_plot_3d import res_plt_2_torch
 
 
 def get_first_nonzero_y(arr):
+    OUTSIDE_VALUE = -100
+
     # 1. Create a boolean mask where elements are non-zero
     mask = arr != 0
     
@@ -25,9 +27,10 @@ def get_first_nonzero_y(arr):
     has_nonzero = np.any(mask, axis=1)
     
     # 4. Replace indices where there were no non-zeros with a sentinel value (e.g., -1)
-    y_indices = np.where(has_nonzero, y_indices, -1)
+    y_indices = np.where(has_nonzero, y_indices, OUTSIDE_VALUE)
     
-    y_indices = np.where(y_indices > arr.shape[1] * .4, -1, y_indices)
+    y_indices = np.where(y_indices < arr.shape[1] * .1, OUTSIDE_VALUE, y_indices)
+    y_indices = np.where(y_indices > arr.shape[1] * .4, OUTSIDE_VALUE, y_indices)
     
     return y_indices.astype(np.float32)
 
@@ -115,13 +118,16 @@ def run_pso_torch_xfr(
     # flat_min_index = np.argmin(y_indices)
     # z_border, x_border = np.unravel_index(flat_min_index, y_indices.shape)
 
-    x_with_nonzero = np.where(np.any(image2_array[:,9,:] != 0, axis=0))[0]
+    x_with_nonzero = np.where(np.any(image2_array[:,image_shape[1]//10,:] != 0, axis=0))[0]
     x1 = x_with_nonzero[0]
     x2 = x_with_nonzero[-1]
 
+    # print(x1,x2)
+    # exit()
+
     x_mid = (x1+x2)/2
-    x1 = x_mid-9
+    x1 = x_mid-image_shape[2]*.1
-    x2 = x_mid+9
+    x2 = x_mid+image_shape[2]*.1
 
     z_sum = np.sum(image2_array, axis=(1, 2))
     z_with_nonzero = np.where(z_sum > 0)[0]
@@ -141,19 +147,20 @@ def run_pso_torch_xfr(
 
     # 設定基本的 bounds
     if CBT == True:    
-        z_bounds = (0, image_shape[0] / 2)
+        # z_bounds = (0, image_shape[0]-1)
         # z_bounds = (z1, (z1+z2)/2)
-        x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1)
+        z_bounds = (.1*image_shape[0], .8*image_shape[0])
-        x_bounds_left  = (0, image_shape[2]/2 - image_shape[2]/10 - 1)
+        # x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1)
+        # x_bounds_left  = (0, image_shape[2]/2 - image_shape[2]/10 - 1)
         x_bounds_right = (x2, image_shape[2]*.9)
         x_bounds_left  = (image_shape[2]*.1, x1)
 
         # azimuth_bounds_l = ((95-azi), (145-azi))
         # azimuth_bounds_r = ((50-azi), (85-azi))
         # altitude_bounds = ((60-alt), (75-alt))
-        azimuth_bounds_l = ((95-azi), (105-azi))
+        azimuth_bounds_l = ((98-azi), (120-azi))
-        azimuth_bounds_r = ((75-azi), (85-azi))
+        azimuth_bounds_r = ((60-azi), (82-azi))
-        altitude_bounds = ((55-alt), (70-alt))
+        altitude_bounds = ((60-alt), (70-alt))
 
     else:      
         z_bounds = (0, image_shape[0] - 1)   
@@ -173,7 +180,7 @@ def run_pso_torch_xfr(
         """
         if optimize_size:
             # d, L = snap_to_discrete_values(pos[5], pos[6])
-            d, L = round_down_to_discrete_values(pos[5], pos[6])
+            d, L = snap_to_discrete_values_xfr(pos[5], pos[6])
             params_5 = pos[:5]
         else:
             d, L = diameter, length
@@ -221,7 +228,7 @@ def run_pso_torch_xfr(
         lb_r = [z_bounds[0], y_bounds[0], x_bounds_right[0], azimuth_bounds_r[0], altitude_bounds[0]]
         ub_r = [z_bounds[1], y_bounds[1], x_bounds_right[1], azimuth_bounds_r[1], altitude_bounds[1]]
 
-    if debug:
+    if True or debug:
         print(lb_l)
         print(ub_l)
         print(lb_r)
@@ -233,7 +240,7 @@ def run_pso_torch_xfr(
     best_position_r = None
 
     # Left side optimization
-    print("\n=== 左側 ===")
+    print(f"\n=== {label_str} 左側 ===")
     kwargs = {'y_indices': y_indices}
     position_l, loss_l = pso(objective_function_xfr, lb_l, ub_l,
     
@@ -299,7 +306,7 @@ def run_pso_torch_xfr(
         # retries += 1
     
     # Right side optimization
-    print("\n=== 右側 ===")
+    print(f"\n=== {label_str} 右側 ===")
     position_r, loss_r = pso(objective_function_xfr, lb_r, ub_r,
         # ieqcons=[constraint_y],
         kwargs=kwargs,    
@@ -318,7 +325,7 @@ def run_pso_torch_xfr(
 
     if optimize_size:
         # diameter_r, length_r = snap_to_discrete_values(position_r[5], position_r[6])
-        diameter_r, length_r = round_down_to_discrete_values(position_r[5], position_r[6])
+        diameter_r, length_r = snap_to_discrete_values_xfr(position_r[5], position_r[6])
         print(f"[RIGHT] Position: {position_r[:5]}")
         print(f"[RIGHT] Diameter: {diameter_r} mm (raw: {position_r[5]:.2f})")
         print(f"[RIGHT] Length: {length_r} mm (raw: {position_r[6]:.2f})")
@@ -376,7 +383,7 @@ def run_pso_torch_xfr(
         final_diameter_r = best_position_r[5]
         final_length_r = best_position_r[6]
         
-        print(f"\n=== 最終結果 ===")
+        print(f"\n=== {label_str} 最終結果 ===")
         print(f"Left - Diameter: {final_diameter_l} mm, Length: {final_length_l} mm")
         print(f"Right - Diameter: {final_diameter_r} mm, Length: {final_length_r} mm")
     else:
@@ -520,7 +527,7 @@ def run_pso_torch(
         """
         if optimize_size:
             # d, L = snap_to_discrete_values(pos[5], pos[6])
-            d, L = round_down_to_discrete_values(pos[5], pos[6])
+            d, L = snap_to_discrete_values_xfr(pos[5], pos[6])
             params_5 = pos[:5]
         else:
             d, L = diameter, length
@@ -643,7 +650,7 @@ def run_pso_torch(
 
     if optimize_size:
         # diameter_r, length_r = snap_to_discrete_values(position_r[5], position_r[6])
-        diameter_r, length_r = round_down_to_discrete_values(position_r[5], position_r[6])
+        diameter_r, length_r = snap_to_discrete_values_xfr(position_r[5], position_r[6])
         print(f"[RIGHT] Position: {position_r[:5]}")
         print(f"[RIGHT] Diameter: {diameter_r} mm (raw: {position_r[5]:.2f})")
         print(f"[RIGHT] Length: {length_r} mm (raw: {position_r[6]:.2f})")
@@ -836,7 +843,7 @@ def run_de_torch(
     def eval_overlap_from_position(pos, side: str, optimize_size: bool, spine_tensor: torch.Tensor, image_shape, spacing):
         if optimize_size:
             # d, L = snap_to_discrete_values(pos[5], pos[6])
-            d, L = round_down_to_discrete_values(pos[5], pos[6])
+            d, L = snap_to_discrete_values_xfr(pos[5], pos[6])
             params_5 = pos[:5]
         else:
             d, L = diameter, length
@@ -1015,7 +1022,7 @@ def run_nm_torch(
     def eval_overlap_from_position(pos, side: str, optimize_size: bool, spine_tensor: torch.Tensor, image_shape, spacing):
         if optimize_size:
             # d, L = snap_to_discrete_values(pos[5], pos[6])
-            d, L = round_down_to_discrete_values(pos[5], pos[6])
+            d, L = snap_to_discrete_values_xfr(pos[5], pos[6])
             params_5 = pos[:5]
         else:
             d, L = diameter, length

core/scoring.py

@@ -23,9 +23,8 @@ def cl_score_torch_xfr(
     in_bone= ((spine_tensor == 1) & (cylinder_torch == 1)).sum().item()
     not_in_bone= ((spine_tensor == 0) & (cylinder_torch == 1)).sum().item()
     
-    # if cyl_total == 0:
+    if cyl_total == 0:
-    #     return float(1000*1000)
+        return float(1e9)  # 極差的情況
-        # return float(1e9)  # 極差的情況
     
     overlap_ratio = overlap / cyl_total
 
@@ -37,10 +36,10 @@ def cl_score_torch_xfr(
     # if in_bone == 0:
     #     return float(not_in_bone*200)
 
-    score += overlap*30
+    score += 10   * in_bone # 10 實在太低
-    score += in_bone*10
+    score += 100  * overlap
-    score -= not_in_bone*1000
+    score -= 2000 * max(0, not_in_bone-10)
-    score -= null_vox2*1000
+    score -= 1000 * max(0, null_vox2-10)
 
     return float(-score)
     

imaging/preprocessing.py

@@ -8,15 +8,15 @@ import glob
 from config.constant import LABEL_MAP
 from imaging.nifti_io import sitk_to_nibabel, nibabel_to_sitk
 
-PROGRESS_FILE = "progress.json"
 
-def load_progress():
+
+def load_progress(PROGRESS_FILE):
     if os.path.exists(PROGRESS_FILE):
         with open(PROGRESS_FILE, "r") as f:
             return json.load(f)
     return {}
 
-def save_progress(progress):
+def save_progress(progress, PROGRESS_FILE):
     with open(PROGRESS_FILE, "w") as f:
         json.dump(progress, f, indent=2)
         
@@ -73,7 +73,9 @@ def process_dataset(image_dir, label_dir, output_dir, labels_to_process=None):
     total_files = len(image_files)
     print(f"Total files: {total_files}")
 
-    progress = load_progress()
+    PROGRESS_FILE = os.path.join(output_dir, "progress.json")
+
+    progress = load_progress(PROGRESS_FILE)
     all_file_summary = []
 
     for idx, image_path in enumerate(image_files, 1):
@@ -119,7 +121,7 @@ def process_dataset(image_dir, label_dir, output_dir, labels_to_process=None):
             "processed_labels": result["processed_labels"],
             "missing_labels": result["missing_labels"]
         }
-        save_progress(progress)
+        save_progress(progress, PROGRESS_FILE)
 
         print(f"[{idx}/{total_files}] Finished: {file_name} | Missing labels: {result['missing_labels'] or 'None'}")
 

imaging/segmentation.py

@@ -35,7 +35,7 @@ def seg_bone(n, name, resampled_sitk_img, resampled_sitk_lbl, output_base=None,
     
     # 2. 獲取所有連通區域
     # 連通區域濾波器會將 binary_mask 中的不同物體標記為 1, 2, 3...
-    cc_image = sitk.ConnectedComponent(binary_mask)
+    cc_image = sitk.ConnectedComponent(binary_mask, True) #fullyConnected
     
     # 3. 根據區域大小（像素/體積）重新標記
     # RelabelComponent 會按大小排序，最大的物體標籤會被設為 1

xfr_debug.py

@@ -9,10 +9,10 @@ from core.objective import set_global_context
 from core.optimizer import run_pso_torch, run_de_torch, run_nm_torch, run_pso_torch_xfr
 from imaging.orientation import azimuth_rotation, analyze_vertebral_tilt_contour
 
-standardized_dir = '/mnt/1248/open/cyrou/CBT/Seg/Resample/standardized-xfr/'
+standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr/'
 
-azimuth_rotation_dir = '/mnt/1248/open/cyrou/azimuth_rotation'
+azimuth_rotation_dir = '/mnt/1248/open2/cyrou/azimuth_rotation'
-tilt_contour_dir = '/mnt/1248/open/cyrou/tilt_contour'
+tilt_contour_dir = '/mnt/1248/open2/cyrou/tilt_contour'
 Output_dir = '/mnt/1248/open/cyrou/Output'
 
 def get_device(gpu_id=None):
@@ -95,7 +95,7 @@ def debug_pso(volume_id, level):
     )    
 
     best_l, loss_l, best_r, loss_r, total_time = run_pso_torch_xfr(
-        label_str=level,
+        label_str=f'{volume_id} {level}',
         image1_path=cortical_path,
         image2_path=binary_path,
         image3_path=roi_path,
@@ -121,21 +121,37 @@ def main():
     # process_volume(volume_id, level)
     # exit()
 
-    for volume_id in (
+    MAX_SUCCESSFUL_VOLUMES = 100
-        '1.3.6.1.4.1.9328.50.4.0001', 
+    number_of_successful_volumes = 0
+
+    # for volume_id in (
+    #     '1.3.6.1.4.1.9328.50.4.0001', 
     #     '1.3.6.1.4.1.9328.50.4.0002', 
     #     '1.3.6.1.4.1.9328.50.4.0003', 
     #     '1.3.6.1.4.1.9328.50.4.0004',
     #     '1.3.6.1.4.1.9328.50.4.0005',
-        # '1.3.6.1.4.1.9328.50.4.0006',
+    #     # '1.3.6.1.4.1.9328.50.4.0006',
-    ):
+    # ):
-    # for volume_id in sorted(os.listdir(standardized_dir)):
-        # debug_orientation(volume_id, level)
-        for level in ('L1', 'L2', 'L3', 'L4', 'L5'):
-            # debug_orientation(volume_id, level)
-            debug_pso(volume_id, level)
 
-        exit()
+    for volume_id in sorted(os.listdir(standardized_dir)):
+        # debug_orientation(volume_id, level)
+        error_flag = False
+        for level in ('L1', 'L2', 'L3', 'L4', 'L5'):
+        # for level in ('L5',):
+            # debug_orientation(volume_id, level)
+            try:
+                debug_pso(volume_id, level)
+            except Exception as e:
+                print(f'Error in {volume_id} {level}: {e}')
+                error_flag = True
+                continue
+        if not error_flag:
+            number_of_successful_volumes += 1
+
+        if number_of_successful_volumes >= MAX_SUCCESSFUL_VOLUMES:
+            break
+
+        # exit()
 
 
 

xfr_preprocess.py

@@ -4,7 +4,7 @@ from imaging.preprocessing import process_dataset
 
 data_root   = '/mnt/1220/Public/dataset/Spine/CTSpine1K/data/'
 label_root  = '/mnt/1220/Public/dataset/Spine/CTSpine1K/label/'
-output_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr/'
+output_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-2/'
 
 label_map = {
     'colon': 'conlon',