Test_mask.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cfff7b3c-d474-4562-8c45-ca88efed2c1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import nibabel as nib\n",
    "import sys\n",
    "import pandas as pd \n",
    "import numpy as np\n",
    "import subprocess\n",
    "from glob import glob\n",
    "from skimage.morphology import dilation\n",
    "from skimage.morphology import ball\n",
    "from scipy.ndimage import gaussian_filter\n",
    "import re\n",
    "import csv\n",
    "from skimage.measure import label\n",
    "import mne"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d7d92cb-48ab-49d4-97e3-d740fb9b0eb8",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0e1cb719-67a7-40fd-b713-acba0d6618d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_nifti(img_path):\n",
    "    test_img = nib.load(img_path)\n",
    "    test_data = np.asarray(test_img.dataobj)\n",
    "    test_aff = test_img.affine\n",
    "    test_shape = test_img.shape\n",
    "\n",
    "    return test_img, test_data, test_aff, test_shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c77437c2-5718-485d-bfab-a9dbff0e33a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def determineFCSVCoordSystem(input_fcsv,overwrite_fcsv=False):\n",
    "    # need to determine if file is in RAS or LPS\n",
    "    # loop through header to find coordinate system\n",
    "    coordFlag = re.compile('# CoordinateSystem')\n",
    "    verFlag = re.compile('# Markups fiducial file version')\n",
    "    headFlag = re.compile('# columns')\n",
    "    coord_sys=None\n",
    "    headFin=None\n",
    "    ver_fin=None\n",
    "\n",
    "    with open(input_fcsv, 'r') as myfile:\n",
    "        firstNlines=myfile.readlines()[0:3]\n",
    "    for row in firstNlines:\n",
    "        row=re.sub(\"[\\s\\,]+[\\,]\",\"\",row).replace(\"\\n\",\"\")\n",
    "        cleaned_dict={row.split('=')[0].strip():row.split('=')[1].strip()}\n",
    "        if None in list(cleaned_dict):\n",
    "            cleaned_dict['# columns'] = cleaned_dict.pop(None)\n",
    "        if any(coordFlag.match(x) for x in list(cleaned_dict)):\n",
    "            coord_sys = list(cleaned_dict.values())[0]\n",
    "        if any(verFlag.match(x) for x in list(cleaned_dict)):\n",
    "            verString = list(filter(verFlag.match,  list(cleaned_dict)))\n",
    "            assert len(verString)==1\n",
    "            ver_fin = verString[0].split('=')[-1].strip()\n",
    "        if any(headFlag.match(x) for x in list(cleaned_dict)):\n",
    "            headFin=list(cleaned_dict.values())[0].split(',')\n",
    "    return coord_sys\n",
    "def extract_coords(file_path):\n",
    "    coord_sys=determineFCSVCoordSystem(file_path)\n",
    "    df = pd.read_csv(file_path, skiprows = 3, header = None)\n",
    "    coord_arr = df[[1,2,3,11]].to_numpy()\n",
    "    if any(x in coord_sys for x in {'LPS','1'}):\n",
    "        coord_arr = coord_arr[0:3] * np.array([-1,-1,1])\n",
    "    return coord_arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "59565a1a-a00b-4de8-9894-dba2b47441fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_points_ct_space(coord_path, ct_t1_path):\n",
    "    t1_coords = extract_coords(coord_path)\n",
    "    ct_coords = np.zeros(t1_coords.shape)\n",
    "    transform = np.loadtxt(ct_t1_path)\n",
    "    transform = np.linalg.inv(transform)\n",
    "\n",
    "    M = transform[:3,:3]\n",
    "    abc = transform[:3,3]\n",
    "\n",
    "\n",
    "    for i in range(len(ct_coords)):\n",
    "        vec = t1_coords[i,:]\n",
    "        tvec = M.dot(vec) + abc\n",
    "        ct_coords[i,:] = tvec[:3]\n",
    "\n",
    "    # ct_coords = np.round(ct_coords).astype(int) #DELETE\n",
    "\n",
    "    return ct_coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ea86a23a-8ac0-4f12-b09d-66435d1b7b42",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_coords_vox_space(planned_coords, img_aff):\n",
    "    inv_affine = np.linalg.inv(img_aff)\n",
    "    M = inv_affine[:3,:3]\n",
    "    abc = inv_affine[:3,3]\n",
    "\n",
    "    transform_coords = np.zeros(planned_coords.shape)\n",
    "\n",
    "    for i in range(len(transform_coords)):\n",
    "        vec = planned_coords[i,:]\n",
    "        tvec = M.dot(vec) + abc\n",
    "        transform_coords[i,:] = tvec[:3]\n",
    "    \n",
    "    transform_coords = np.round(transform_coords).astype(int)\n",
    "\n",
    "    return transform_coords\n",
    "\n",
    "def check_coord_dims(pointa, pointb, img_size):\n",
    "    for i in range(3):\n",
    "        if pointa[i] > img_size[i] or pointb[i] > img_size[i]:\n",
    "            print('review planned_fcsv - dim mismatch')\n",
    "            return False\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "939d643d-ec40-46fb-b2b5-2ad6d6a2a2d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_line_mask(point1, point2, shape):\n",
    "    # Create an empty mask with the specified shape\n",
    "    mask = np.zeros(shape, dtype=bool)\n",
    "    x1, y1, z1 = point1\n",
    "    x2, y2, z2 = point2\n",
    "    \n",
    "    mask[x1, y1, z1] = True\n",
    "    # Get the directions of each axis\n",
    "    dx = abs(x2 - x1)\n",
    "    dy = abs(y2 - y1)\n",
    "    dz = abs(z2 - z1)\n",
    "    if x2 > x1:\n",
    "        xs = 1\n",
    "    else:\n",
    "        xs = -1\n",
    "    if y2 > y1:\n",
    "        ys = 1\n",
    "    else:\n",
    "        ys = -1\n",
    "    if z2 > z1:\n",
    "        zs = 1\n",
    "    else:\n",
    "        zs = -1\n",
    "\n",
    "    # Driving axis is X-axis\n",
    "    if dx >= dy and dx >= dz:\n",
    "        p1 = 2 * dy - dx\n",
    "        p2 = 2 * dz - dx\n",
    "        while x1 != x2:\n",
    "            x1 += xs\n",
    "            if p1 >= 0:\n",
    "                y1 += ys\n",
    "                p1 -= 2 * dx\n",
    "            if p2 >= 0:\n",
    "                z1 += zs\n",
    "                p2 -= 2 * dx\n",
    "            p1 += 2 * dy\n",
    "            p2 += 2 * dz\n",
    "            mask[x1, y1, z1] = True\n",
    "        mask[x1, y1, z1] = True\n",
    "\n",
    "    # Driving axis is Y-axis\"\n",
    "    elif dy >= dx and dy >= dz:\n",
    "        p1 = 2 * dx - dy\n",
    "        p2 = 2 * dz - dy\n",
    "        while y1 != y2:\n",
    "            y1 += ys\n",
    "            if p1 >= 0:\n",
    "                x1 += xs\n",
    "                p1 -= 2 * dy\n",
    "            if p2 >= 0:\n",
    "                z1 += zs\n",
    "                p2 -= 2 * dy\n",
    "            p1 += 2 * dx\n",
    "            p2 += 2 * dz\n",
    "            mask[x1, y1, z1] = True\n",
    "        mask[x1, y1, z1] = True\n",
    "\n",
    "    # Driving axis is Z-axis\"\n",
    "    else:\n",
    "        p1 = 2 * dy - dz\n",
    "        p2 = 2 * dx - dz\n",
    "        while z1 != z2:\n",
    "            z1 += zs\n",
    "            if p1 >= 0:\n",
    "                y1 += ys\n",
    "                p1 -= 2 * dz\n",
    "            if p2 >= 0:\n",
    "                x1 += xs\n",
    "                p2 -= 2 * dz\n",
    "            p1 += 2 * dy\n",
    "            p2 += 2 * dx\n",
    "            mask[x1, y1, z1] = True\n",
    "        mask[x1, y1, z1] = True\n",
    "    return mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df62af53-f0fc-4e3d-93af-cd10936b889f",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Test code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d5cd58e4-9596-4cd5-b7f5-141dc7b06472",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = './data/sub-P020_res-0p4mm_ct.nii.gz'\n",
    "coord_path = './data/sub-P020_actual.fcsv'\n",
    "ct_t1_trans = './data/sub-P020_desc-rigid_from-ct_to-T1w_type-ras_ses-post_xfm.txt'\n",
    "final_path = 'test.nii.gz'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "83c77b19-6384-425e-8ea1-fb0acda5cf44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, read nifti file\n",
    "img, img_data, img_aff, img_shape = read_nifti(img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "de4c6e63-ee70-4372-8d72-41b8990514a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  -0.39936084,    0.        ,    0.        ,  125.04445648],\n",
       "       [   0.        ,    0.39936084,    0.        , -137.95605469],\n",
       "       [   0.        ,    0.        ,    0.39911309,  -10.36294365],\n",
       "       [   0.        ,    0.        ,    0.        ,    1.        ]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img_aff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1d709d01-12fe-4538-9d4c-abd41ec7c02c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(626, 626, 451)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img_shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "dac7315b-2b23-403d-9ea1-f79c14abcac5",
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "shapes (3,3) and (4,) not aligned: 3 (dim 1) != 4 (dim 0)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [11], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Now get coordinates in CT space\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m ct_coord_arr \u001b[38;5;241m=\u001b[39m \u001b[43mtransform_points_ct_space\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcoord_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mct_t1_trans\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      3\u001b[0m ct_coord_arr[\u001b[38;5;241m0\u001b[39m:\u001b[38;5;241m3\u001b[39m,:]\n",
      "Cell \u001b[0;32mIn [4], line 13\u001b[0m, in \u001b[0;36mtransform_points_ct_space\u001b[0;34m(coord_path, ct_t1_path)\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mlen\u001b[39m(ct_coords)):\n\u001b[1;32m     12\u001b[0m     vec \u001b[38;5;241m=\u001b[39m t1_coords[i,:]\n\u001b[0;32m---> 13\u001b[0m     tvec \u001b[38;5;241m=\u001b[39m \u001b[43mM\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdot\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvec\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;241m+\u001b[39m abc\n\u001b[1;32m     14\u001b[0m     ct_coords[i,:] \u001b[38;5;241m=\u001b[39m tvec[:\u001b[38;5;241m3\u001b[39m]\n\u001b[1;32m     16\u001b[0m \u001b[38;5;66;03m# ct_coords = np.round(ct_coords).astype(int) #DELETE\u001b[39;00m\n",
      "\u001b[0;31mValueError\u001b[0m: shapes (3,3) and (4,) not aligned: 3 (dim 1) != 4 (dim 0)"
     ]
    }
   ],
   "source": [
    "# Now get coordinates in CT space\n",
    "ct_coord_arr = transform_points_ct_space(coord_path, ct_t1_trans)\n",
    "ct_coord_arr[0:3,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "769dbb0c-0bd5-435b-b57f-4b414c3478b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transform to voxels\n",
    "transformed_coords = transform_coords_vox_space(ct_coord_arr, img_aff)\n",
    "transformed_coords[0:3,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1749ac6-ae1a-47dd-a52e-7ab1e730f8ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create mask\n",
    "final_mask = np.zeros(img_shape).astype(bool)\n",
    "    \n",
    "for i in range(0,len(transformed_coords),2):\n",
    "    pointa = transformed_coords[i, :]\n",
    "    pointb = transformed_coords[i+1, :]\n",
    "\n",
    "\n",
    "    test_mask = create_line_mask(pointa, pointb, img_shape)\n",
    "\n",
    "    footprint = ball(4)\n",
    "    dilated = dilation(test_mask, footprint)\n",
    "    result = gaussian_filter(dilated.astype(float), sigma=0.6)\n",
    "    final_mask += result>0\n",
    "\n",
    "clipped_img = nib.Nifti1Image(final_mask, img_aff, img.header)\n",
    "nib.save(clipped_img, final_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1236c285-8da0-414d-aede-4145eb778b0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "img, img_data, img_aff, img_shape = read_nifti('./data/sub-P020_res-0p4mm_desc-contacts_mask.nii.gz')\n",
    "img_aff, img_shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b533ab0e-fec3-4404-9db1-d3bdcf42e88b",
   "metadata": {},
   "outputs": [],
   "source": [
    "img, img_data, img_aff, img_shape = read_nifti('./data/sub-P020_res-0p4mm_desc-z_norm_pred_ct.nii.gz')\n",
    "img_aff, img_shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "147d5f26-28a3-4f40-ad17-9953c37fb169",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Post-process"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db6253c8-f738-4878-a877-076b652f7e36",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ff3c866-6ac3-464c-ba66-5eeea86c8c76",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Old algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d52105bb-36c0-4168-8ea4-3565626b0de5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.spatial import distance\n",
    "from sklearn.linear_model import LinearRegression\n",
    "\n",
    "def trace_line(entry_point, target_point):\n",
    "    # Vector from target to entry\n",
    "    return (entry_point - target_point)/np.linalg.norm(target_point - entry_point)\n",
    "\n",
    "def find_next_contact_orig(current_point,\n",
    "                      direction,\n",
    "                      contacts,\n",
    "                      max_distance_target,\n",
    "                      min_distance_prior,\n",
    "                      avg_distance,\n",
    "                      entry_point,\n",
    "                      distance_from_entry=1.5):\n",
    "    next_contact = None\n",
    "    next_contact_id = None\n",
    "    min_distance = float('inf')\n",
    "    # print(current_point, avg_distance, direction)\n",
    "    theory_point = current_point + avg_distance*direction\n",
    "    # Projection of entry on direction vector\n",
    "    entry_proj = np.dot(entry_point, direction) / np.linalg.norm(direction)\n",
    "    \n",
    "    for ind, contact_array in enumerate(contacts):\n",
    "        contact = contact_array[0:3]\n",
    "#         vector_to_contact = contact - current_point\n",
    "#         angle = calculate_angle(vector_to_contact, direction)\n",
    "\n",
    "#         # If the angle is lower than 30 degrees (vectors heading the same direction)\n",
    "#         if dot_product > 0:\n",
    "#             allowed_distance = radius + avg_distance\n",
    "#         else:\n",
    "#             allowed_distance = radius\n",
    "\n",
    "        distance_target = np.linalg.norm(theory_point - contact)\n",
    "        distance_prior = np.linalg.norm(current_point - contact)\n",
    "        distance_entry = np.linalg.norm(entry_point - contact)\n",
    "        contact_proj = np.dot(contact, direction) / np.linalg.norm(direction)\n",
    "        print(ind, distance_target, distance_prior, distance_entry)\n",
    "        print(0, max_distance_target*avg_distance, min_distance_prior*avg_distance, distance_from_entry, '\\n')\n",
    "        distance_metric = np.sqrt(distance_target**2+distance_prior**2) # balance between distance to target and prior\n",
    "        # Check that it's the closest point to the target, \n",
    "        # that it's below a specific distance from the target\n",
    "        # and above a certain distance from the current point,\n",
    "        # that it is not too close to the entry\n",
    "        # and that is not beyond the entry point\n",
    "        if (distance_metric < min_distance and\n",
    "            distance_target< max_distance_target*avg_distance and\n",
    "            distance_prior >= min_distance_prior*avg_distance and\n",
    "            (distance_entry >= distance_from_entry and contact_proj < entry_proj)):\n",
    "            min_distance = distance_metric\n",
    "            next_contact = contact\n",
    "            next_contact_id = ind\n",
    "    if next_contact_id is not None:\n",
    "        print(next_contact_id, '\\n')\n",
    "        contacts = np.delete(contacts, next_contact_id, axis=0) \n",
    "    return next_contact, contacts\n",
    "\n",
    "def adjust_direction_using_all_contacts(contacts, original_direction):\n",
    "    contacts = np.array(contacts)[-3:,:] # only use last 3 coords\n",
    "    X = contacts[:, 0].reshape(-1, 1)\n",
    "    y1 = contacts[:, 1]\n",
    "    y2 = contacts[:, 2]\n",
    "\n",
    "    reg1 = LinearRegression().fit(X, y1)\n",
    "    reg2 = LinearRegression().fit(X, y2)\n",
    "\n",
    "    direction_vector = np.array([1, reg1.coef_[0], reg2.coef_[0]])\n",
    "    direction_vector = direction_vector / np.linalg.norm(direction_vector)\n",
    "\n",
    "    # Ensure the direction vector points toward the initial direction\n",
    "    if np.dot(direction_vector, original_direction) < 0:\n",
    "        direction_vector = -direction_vector\n",
    "        \n",
    "    return direction_vector\n",
    "\n",
    "def run_segmentation_qc_orig(entry_point, target_point, contacts, max_distance_target, min_distance_prior, initial_spacing=4):\n",
    "    original_direction = trace_line(entry_point, target_point)\n",
    "    direction = original_direction\n",
    "    current_point = target_point\n",
    "    avg_distance = 1 # mm\n",
    "    \n",
    "    # Initialize with the closest contact to the entry\n",
    "    current_point, contacts = find_next_contact_orig(current_point, direction, contacts, 2, 0, avg_distance, entry_point)\n",
    "    found_contacts = [current_point]\n",
    "    distances_between_contacts = []\n",
    "    next_contact = current_point\n",
    "    # Find next contact\n",
    "    next_contact, contacts = find_next_contact_orig(current_point, direction, contacts, max_distance_target, min_distance_prior, initial_spacing, entry_point)\n",
    "    # Calculate\n",
    "    while next_contact is not None:   \n",
    "        found_contacts.append(next_contact)\n",
    "        distances_between_contacts.append(np.linalg.norm(next_contact - current_point))\n",
    "        if len(found_contacts) >= 3:\n",
    "            # print('Cambio direccion \\n')\n",
    "            direction = adjust_direction_using_all_contacts(found_contacts, original_direction)\n",
    "            # print(direction)\n",
    "        current_point = next_contact\n",
    "\n",
    "        avg_distance = np.mean(distances_between_contacts)\n",
    "        \n",
    "        # Find next contact\n",
    "        if len(contacts) > 0:\n",
    "            next_contact, contacts = find_next_contact_orig(current_point, direction, contacts, max_distance_target, min_distance_prior, avg_distance, entry_point)\n",
    "        else:\n",
    "            next_contact = None\n",
    "    print(f\"Adjusted Mean Distance: {avg_distance}\")\n",
    "    return found_contacts, avg_distance"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5ac3a34-362f-4560-9076-4e10f8b94e64",
   "metadata": {},
   "source": [
    "#### New algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "c6635f3e-d022-4df0-a1f9-6cf2876f332b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.spatial import distance\n",
    "from sklearn.linear_model import LinearRegression\n",
    "\n",
    "# Function to calculate the angle between two vectors in degrees\n",
    "def calculate_angle(v1, v2):\n",
    "    dot_product = np.dot(v1, v2)\n",
    "    magnitude_v1 = np.linalg.norm(v1)\n",
    "    magnitude_v2 = np.linalg.norm(v2)\n",
    "    cos_theta = dot_product / (magnitude_v1 * magnitude_v2)\n",
    "    angle_rad = np.arccos(cos_theta)\n",
    "    angle_deg = np.degrees(angle_rad)\n",
    "    return angle_deg\n",
    "\n",
    "def trace_line(entry_point, target_point):\n",
    "    # Vector from target to entry\n",
    "    return (entry_point - target_point)/np.linalg.norm(target_point - entry_point)\n",
    "\n",
    "def find_next_contact_new(current_point,\n",
    "                      direction,\n",
    "                      contacts,\n",
    "                      max_angle_target,\n",
    "                      max_distance_prior, # percentage\n",
    "                      min_distance_prior, # percentage\n",
    "                      avg_distance,\n",
    "                      entry_point,\n",
    "                      distance_from_entry=1.5,\n",
    "                        ):\n",
    "    next_contact = None\n",
    "    next_contact_id = None\n",
    "    min_distance = float('inf')\n",
    "    # print(current_point, avg_distance, direction)\n",
    "    theory_point = current_point + avg_distance*direction\n",
    "    # Projection of entry on direction vector\n",
    "    entry_proj = np.dot(entry_point, direction) / np.linalg.norm(direction)\n",
    "    \n",
    "    for ind, contact_array in enumerate(contacts):\n",
    "        contact = contact_array[0:3]\n",
    "        \n",
    "        # Vector last detected contact to possible next contact\n",
    "        v_contact = trace_line(contact, current_point)\n",
    "        # Angle between the predicted direction and the v_contact\n",
    "        angle_contact = calculate_angle(v_contact, direction)\n",
    "        distance_target = np.linalg.norm(theory_point - contact)\n",
    "        distance_prior = np.linalg.norm(current_point - contact)\n",
    "        distance_entry = np.linalg.norm(entry_point - contact)\n",
    "        contact_proj = np.dot(contact, direction) / np.linalg.norm(direction)\n",
    "        # print(ind, angle_contact, distance_prior, distance_entry)\n",
    "        # print(0, max_distance_prior*avg_distance, min_distance_prior*avg_distance, distance_from_entry, '\\n')\n",
    "        distance_metric = np.sqrt(distance_target**2+distance_prior**2) # balance between distance to target and prior\n",
    "        # Check that it's the closest point to the target, \n",
    "        # that it's below a specific distance from the target\n",
    "        # and above a certain distance from the current point,\n",
    "        # that it is not too close to the entry\n",
    "        # and that is not beyond the entry point\n",
    "        if (distance_metric < min_distance and\n",
    "            angle_contact <= max_angle_target and\n",
    "            distance_prior >= min_distance_prior*avg_distance and\n",
    "            distance_prior <= max_distance_prior*avg_distance):\n",
    "            if distance_from_entry is None or (distance_entry >= distance_from_entry and contact_proj < entry_proj):\n",
    "                min_distance = distance_metric\n",
    "                next_contact = contact\n",
    "                next_contact_id = ind\n",
    "    if next_contact_id is not None:\n",
    "        # print(contacts[next_contact_id,-1], '\\n')\n",
    "        # print(next_contact_id, '\\n')\n",
    "        contacts = np.delete(contacts, next_contact_id, axis=0) \n",
    "    return next_contact, contacts\n",
    "\n",
    "def adjust_direction_using_all_contacts(contacts, original_direction):\n",
    "    n_elements_reg = 3 # use last 3 elements\n",
    "    contacts = np.array(contacts)[-n_elements_reg:,:] # only use last 3 coords\n",
    "    # print(contacts)\n",
    "    \n",
    "    # Check dimension where all vals are not equal\n",
    "    dim = None\n",
    "    i = 0\n",
    "    dims = list(range(n_elements_reg))\n",
    "    while dim is None and i<n_elements_reg:\n",
    "        if not np.all(contacts[:,i] == contacts[0,i]):\n",
    "            dim = dims.pop(i)\n",
    "        i += 1\n",
    "    assert dim is not None\n",
    "    \n",
    "    X = contacts[:, dim].reshape(-1, 1)\n",
    "    y1 = contacts[:, dims[0]]\n",
    "    y2 = contacts[:, dims[1]]\n",
    "\n",
    "    reg1 = LinearRegression().fit(X, y1)\n",
    "    reg2 = LinearRegression().fit(X, y2)\n",
    "\n",
    "    direction_vector = np.ones(n_elements_reg)\n",
    "    direction_vector[dims[0]] = reg1.coef_[0]\n",
    "    direction_vector[dims[1]] = reg2.coef_[0]\n",
    "    # Normalize\n",
    "    direction_vector = direction_vector / np.linalg.norm(direction_vector)\n",
    "\n",
    "    # Ensure the direction vector points toward the initial direction\n",
    "    if np.dot(direction_vector, original_direction) < 0:\n",
    "        direction_vector = -direction_vector\n",
    "        \n",
    "    return direction_vector\n",
    "\n",
    "def run_segmentation_qc_new(entry_point,\n",
    "                             target_point,\n",
    "                             contacts, \n",
    "                             max_angle_target, \n",
    "                             max_distance_prior, \n",
    "                             min_distance_prior, \n",
    "                             initial_spacing=4,\n",
    "                             max_n_contacts = -1 # set to -1 to apply distance threshold to entry\n",
    "                           ):\n",
    "    original_direction = trace_line(entry_point, target_point)\n",
    "    direction = original_direction\n",
    "    current_point = target_point\n",
    "    avg_distance = 2 # mm\n",
    "    \n",
    "    # Initialize with the closest contact to the entry\n",
    "    current_point, contacts = find_next_contact_new(current_point,\n",
    "                                                    direction, \n",
    "                                                    contacts, \n",
    "                                                    90, \n",
    "                                                    2, \n",
    "                                                    0, \n",
    "                                                    avg_distance, \n",
    "                                                    entry_point)\n",
    "    found_contacts = [current_point]\n",
    "    distances_between_contacts = []\n",
    "    next_contact = current_point\n",
    "    # Find next contact\n",
    "    next_contact, contacts = find_next_contact_new(current_point, \n",
    "                                                   direction, \n",
    "                                                   contacts, \n",
    "                                                   max_angle_target, \n",
    "                                                   max_distance_prior, \n",
    "                                                   min_distance_prior, \n",
    "                                                   initial_spacing, \n",
    "                                                   entry_point)\n",
    "    # Calculate\n",
    "    while next_contact is not None:   \n",
    "        found_contacts.append(next_contact)\n",
    "        distances_between_contacts.append(np.linalg.norm(next_contact - current_point))\n",
    "        if len(found_contacts) >= 3:\n",
    "            # print('Cambio direccion \\n')\n",
    "            # print(direction)\n",
    "            direction = adjust_direction_using_all_contacts(found_contacts, original_direction)\n",
    "            # print(direction)\n",
    "        current_point = next_contact\n",
    "\n",
    "        avg_distance = np.mean(distances_between_contacts)\n",
    "        \n",
    "        # Find next contact\n",
    "        if len(contacts) > 0:\n",
    "            next_contact, contacts = find_next_contact_new(current_point,\n",
    "                                                           direction, \n",
    "                                                           contacts, \n",
    "                                                           max_angle_target, \n",
    "                                                           max_distance_prior, \n",
    "                                                           min_distance_prior, \n",
    "                                                           avg_distance, \n",
    "                                                           entry_point,\n",
    "                                                          distance_from_entry=1.5 if len(found_contacts)>max_n_contacts else None)\n",
    "        else:\n",
    "            next_contact = None\n",
    "    print(f\"Adjusted Mean Distance: {avg_distance}\")\n",
    "    return found_contacts, avg_distance\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5579ba37-971a-4420-948c-fb104ad6a52c",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Others"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "61c56018-afee-4df9-b992-3e314a3a748c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def df_to_fcsv(input_df, output_fcsv):\n",
    "\twith open(output_fcsv, 'w') as fid:\n",
    "\t\tfid.write(\"# Markups fiducial file version = 4.11\\n\")\n",
    "\t\tfid.write(\"# CoordinateSystem = 0\\n\")\n",
    "\t\tfid.write(\"# columns = id,x,y,z,ow,ox,oy,oz,vis,sel,lock,label,desc,associatedNodeID\\n\")\n",
    "\t\n",
    "\tout_df={'node_id':[],'x':[],'y':[],'z':[],'ow':[],'ox':[],'oy':[],'oz':[],\n",
    "\t\t'vis':[],'sel':[],'lock':[],'label':[],'description':[],'associatedNodeID':[]\n",
    "\t}\n",
    "\t\n",
    "\tfor idx,ifid in input_df.iterrows():\n",
    "\t\tout_df['node_id'].append(idx+1)\n",
    "\t\tout_df['x'].append(ifid.iloc[0])\n",
    "\t\tout_df['y'].append(ifid.iloc[1])\n",
    "\t\tout_df['z'].append(ifid.iloc[2])\n",
    "\t\tout_df['ow'].append(0)\n",
    "\t\tout_df['ox'].append(0)\n",
    "\t\tout_df['oy'].append(0)\n",
    "\t\tout_df['oz'].append(0)\n",
    "\t\tout_df['vis'].append(1)\n",
    "\t\tout_df['sel'].append(1)\n",
    "\t\tout_df['lock'].append(1)\n",
    "\t\tout_df['label'].append(str(ifid.iloc[3]))\n",
    "\t\tout_df['description'].append('')\n",
    "\t\tout_df['associatedNodeID'].append('')\n",
    "\n",
    "\tout_df=pd.DataFrame(out_df)\n",
    "\tout_df.to_csv(output_fcsv, sep=',', index=False, lineterminator=\"\", mode='a', header=False, float_format = '%.3f')\n",
    "\n",
    "def filter_points(points_df, mask_data, mask_affine):\n",
    "    filtered_points = []\n",
    "    for _, row in points_df.iterrows():\n",
    "        point = row[[1, 2, 3]].to_numpy()\n",
    "        if inside_elec_mask(point, mask_data, mask_affine):\n",
    "            filtered_points.append(row)\n",
    "    return pd.DataFrame(filtered_points)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bdbdbc1-5792-40a8-aa17-3caa9b5c0cbf",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### First stage: get contacts position"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "89c49cc8-2022-45f0-90cd-2717c3d2bae5",
   "metadata": {},
   "outputs": [],
   "source": [
    "pnms_dir = f'./data/prob_nms'\n",
    "gt_dir = './data'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8e109387-a9c6-4cdb-a3bb-60cfb93f5efe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['sub-D128', 'sub-D112', 'sub-D129']\n"
     ]
    }
   ],
   "source": [
    "subjects = []\n",
    "\n",
    "for root, dirs, files in os.walk(pnms_dir):\n",
    "    for file in files:\n",
    "        subject = file.split('_')[0]\n",
    "        if subject not in subjects:\n",
    "            subjects.append(subject)\n",
    "\n",
    "print(subjects)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "77d90f48-1bf4-4b68-82ea-43aa7935cc19",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub = subjects[1]\n",
    "orig_pnms = pd.read_csv(f'{pnms_dir}/{sub}_prob_nms.fcsv', skiprows=3, header = None)\n",
    "ct_t1_trans = np.loadtxt(f'{gt_dir}/{sub}/{sub}_desc-rigid_from-ct_to-T1w_type-ras_ses-post_xfm.txt')\n",
    "ct = nib.load(f'./data/{sub}/{sub}_res-0p4mm_ct.nii.gz')\n",
    "entry_target_path = f'./data/{sub}/{sub}_actual.fcsv'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb05da9b-34a4-4949-aa5b-5e5dd54ed8e0",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Second stage: apply masking filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e1862813-f2b8-4e2f-a854-1dbd6fe34464",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-32.849 33.566 -12.016 'LPIn']\n",
      " [-13.9734 -17.8134 55.5648 'LPIn']]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(array([[-21.54201376,  25.53473941,  40.50927297],\n",
       "        [ -8.25821108, -16.6060593 , 115.41583266]]),\n",
       " 'LPIn')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get entry-target coords\n",
    "entry_target = extract_coords(entry_target_path)\n",
    "ind = 14-4\n",
    "entry_target_interest = entry_target[ind:ind+2,:]\n",
    "label = entry_target_interest[0,-1]\n",
    "print(entry_target_interest)\n",
    "# Transform to CT space\n",
    "entry_target_interest = mne.transforms.apply_trans(np.linalg.inv(ct_t1_trans), entry_target_interest[:,:-1].astype(float))\n",
    "entry_target_interest, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a3537dfd-a623-414b-aa6a-8661cec1e9fb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((4, 4), (626, 626, 401))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inv_affine = np.linalg.inv(ct.affine)\n",
    "data = np.asarray(ct.dataobj)\n",
    "inv_affine.shape, data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3df8fc46-5ee7-42c0-9a10-b7878410e602",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[389, 372, 143],\n",
       "       [356, 266, 330]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Transform to voxels\n",
    "transformed_coords = np.round(mne.transforms.apply_trans(inv_affine, entry_target_interest)).astype(int)\n",
    "transformed_coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a3821262-70ed-46f8-877a-07cbf44b467a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create mask\n",
    "img_shape = data.shape\n",
    "final_mask = np.zeros(img_shape).astype(bool)\n",
    "    \n",
    "for i in range(0,len(transformed_coords),2):\n",
    "    pointa = transformed_coords[i, :]\n",
    "    pointb = transformed_coords[i+1, :]\n",
    "\n",
    "\n",
    "    test_mask = create_line_mask(pointa, pointb, img_shape)\n",
    "\n",
    "    footprint = ball(4)\n",
    "    dilated = dilation(test_mask, footprint)\n",
    "    result = gaussian_filter(dilated.astype(float), sigma=0.6)\n",
    "    final_mask += result>0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "2237a02d-8784-43f6-ab05-247e452802a5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([False,  True]), (626, 626, 401))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Mask based on CT intensity\n",
    "mask_intensity = (data>2500).astype(bool)\n",
    "\n",
    "# Merge the masks\n",
    "merged_mask = final_mask*mask_intensity\n",
    "\n",
    "# Dilate a bit\n",
    "footprint = ball(2)\n",
    "merged_mask = dilation(merged_mask, footprint)\n",
    "\n",
    "np.unique(merged_mask), merged_mask.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "760505d5-7245-48ea-9e88-c4161d25ce03",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(109, 3)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coords_interest = orig_pnms.iloc[:,1:4].to_numpy()\n",
    "coords_interest.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e698c286-be48-4072-b96a-568241498005",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[343, 434, 140],\n",
       "       [371, 270, 146],\n",
       "       [399, 432, 166]])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get contacts positions in voxel space\n",
    "transformed_coords = np.round(mne.transforms.apply_trans(inv_affine, coords_interest)).astype(int)\n",
    "transformed_coords[:3,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "7f09c5d8-e7fe-4064-8a1d-dedf263c8bea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((10, 3),\n",
       " array([False, False, False, False, False, False, False, False, False,\n",
       "        False, False, False, False, False, False, False, False, False,\n",
       "        False, False, False, False, False, False, False, False, False,\n",
       "        False, False, False, False, False, False,  True, False, False,\n",
       "        False, False, False, False, False, False, False, False, False,\n",
       "        False, False, False, False, False, False, False, False, False,\n",
       "        False, False, False, False,  True, False, False, False, False,\n",
       "        False, False, False, False, False, False, False, False, False,\n",
       "        False, False,  True,  True, False,  True, False,  True, False,\n",
       "        False, False, False, False, False, False, False, False, False,\n",
       "        False, False,  True,  True,  True,  True, False, False, False,\n",
       "        False, False, False, False, False, False, False, False, False,\n",
       "        False]))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Mask the contacts inside the mask\n",
    "filtered_coords = coords_interest[merged_mask[transformed_coords[:,0], transformed_coords[:,1], transformed_coords[:,2]]]\n",
    "filtered_coords.shape, merged_mask[transformed_coords[:,0], transformed_coords[:,1], transformed_coords[:,2]]\n",
    "# electrode_mask = nib.load(f'./data/{sub}/{sub}_res-0p4mm_desc-electrode_mask.nii.gz')\n",
    "# electrode_mask = electrode_mask.get_fdata()\n",
    "\n",
    "# filtered_coords = coords_interest[electrode_mask[transformed_coords[:,0], transformed_coords[:,1], transformed_coords[:,2]] > 0]\n",
    "# filtered_coords.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "763aa46f-6655-4b92-a827-b9d51852f989",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "        x       y       z label_order\n",
      "0 -20.708  22.811  45.021     LPIn-01\n",
      "1 -21.507  24.408  42.627     LPIn-02\n",
      "2 -17.513  14.025  60.582     LPIn-03\n",
      "3 -16.715  11.229  65.370     LPIn-04\n",
      "4 -18.312  15.622  58.188     LPIn-05\n",
      "5 -19.909  21.213  47.415     LPIn-06\n",
      "6 -17.513  13.226  62.976     LPIn-07\n",
      "7 -19.909  20.414  50.208     LPIn-08\n",
      "8 -18.312  17.220  54.996     LPIn-09\n",
      "9 -19.111  18.817  52.602     LPIn-10\n"
     ]
    }
   ],
   "source": [
    "dict_df = {\n",
    "    'x': filtered_coords[:,0],\n",
    "    'y': filtered_coords[:,1],\n",
    "    'z': filtered_coords[:,2],\n",
    "    'label_order': [f'{label}-{idx+1:02d}' for idx in range(filtered_coords.shape[0])],\n",
    "\n",
    "}\n",
    "labeled_contacts = pd.DataFrame(dict_df)\n",
    "print(labeled_contacts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "4e551b58-1bb5-48bf-9d94-f25017d9fc3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_to_fcsv(labeled_contacts, f'./data/{sub}/test_fiducials.fcsv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0d5aa5a-7a74-4596-b119-c6cefe3337bf",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Third stage: filter with line algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "11464d13-a023-43ba-85d4-e7114ae86009",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 \n",
      "\n",
      "0 \n",
      "\n",
      "3 \n",
      "\n",
      "4 \n",
      "\n",
      "5 \n",
      "\n",
      "4 \n",
      "\n",
      "2 \n",
      "\n",
      "0 \n",
      "\n",
      "1 \n",
      "\n",
      "0 \n",
      "\n",
      "Adjusted Mean Distance: 3.016676015126724\n",
      "Found Contacts: [array([-21.507,  24.408,  42.627]), array([-20.708,  22.811,  45.021]), array([-19.909,  21.213,  47.415]), array([-19.909,  20.414,  50.208]), array([-19.111,  18.817,  52.602]), array([-18.312,  17.22 ,  54.996]), array([-18.312,  15.622,  58.188]), array([-17.513,  14.025,  60.582]), array([-17.513,  13.226,  62.976]), array([-16.715,  11.229,  65.37 ])]\n",
      "Average Distance: 3.016676015126724\n"
     ]
    }
   ],
   "source": [
    "# Example usage\n",
    "target_point = entry_target_interest[0,:]  # Replace with actual entry point coordinates\n",
    "entry_point = entry_target_interest[1,:]  # Replace with actual target point coordinates\n",
    "contacts = filtered_coords  # Replace with your file path\n",
    "min_distance_prior = 0.5 # at least half the average distance between contacts\n",
    "max_angle_target = 35\n",
    "max_distance_prior = 2 # not more than twice the average distance between contacts\n",
    "\n",
    "# found_contacts, avg_distance = run_segmentation_qc(entry_point, target_point, contacts, max_distance_target, min_distance_prior)\n",
    "found_contacts, avg_distance = run_segmentation_qc_new(entry_point,\n",
    "                                                        target_point,\n",
    "                                                        contacts, \n",
    "                                                        max_angle_target, \n",
    "                                                        max_distance_prior, \n",
    "                                                        min_distance_prior, \n",
    "                                                        initial_spacing=4)\n",
    "print(\"Found Contacts:\", found_contacts)\n",
    "print(\"Average Distance:\", avg_distance)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "793281b4-c8b1-4213-ad80-38f2d0a92afb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10, 10)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(found_contacts), len(filtered_coords)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "47b141d9-7332-4d9f-bace-76e5a599e0e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "found_contacts = np.array(found_contacts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "e10b13e1-e3ef-4efe-a6f4-8854e035dcdc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10, 3)"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "found_contacts.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "9902de05-a916-444d-9d32-fee022a42c3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "dict_df = {\n",
    "    'x': found_contacts[:,0],\n",
    "    'y': found_contacts[:,1],\n",
    "    'z': found_contacts[:,2],\n",
    "    'label_order': [f'{label}-{idx+1:02d}' for idx in range(found_contacts.shape[0])],\n",
    "\n",
    "}\n",
    "labeled_contacts = pd.DataFrame(dict_df)\n",
    "df_to_fcsv(labeled_contacts, f'./data/{sub}/test_fiducials.fcsv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7ef053f-2cf7-4aff-ae03-09940f53b73b",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Use in all subjects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "4d35ba2f-dcfa-466b-923e-1917cfc958b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['sub-D128', 'sub-D112', 'sub-D129']\n"
     ]
    }
   ],
   "source": [
    "pnms_dir = f'./data/prob_nms'\n",
    "gt_dir = './data'\n",
    "\n",
    "# Find subjects\n",
    "subjects = []\n",
    "\n",
    "for root, dirs, files in os.walk(pnms_dir):\n",
    "    for file in files:\n",
    "        subject = file.split('_')[0]\n",
    "        if subject not in subjects:\n",
    "            subjects.append(subject)\n",
    "\n",
    "print(subjects)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "58add60d-e263-4507-89bd-6d5488f21589",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sub-D112\n",
      "LAm\n",
      "Adjusted Mean Distance: 4.012370745962754\n",
      "LAHc\n",
      "Adjusted Mean Distance: 4.889503697160149\n",
      "LPHc\n",
      "Adjusted Mean Distance: 4.925406591290107\n",
      "RAIn\n",
      "Adjusted Mean Distance: 3.0548136896637557\n",
      "LAIn\n",
      "Adjusted Mean Distance: 2.9836174317785957\n",
      "LPIn\n",
      "Adjusted Mean Distance: 3.016676015126724\n",
      "LOFr\n",
      "Adjusted Mean Distance: 5.002209551373188\n",
      "LInfTe\n",
      "Adjusted Mean Distance: 4.063290593123217\n",
      "LPTe\n",
      "Adjusted Mean Distance: 4.96340957275441\n",
      "LPTeOc\n",
      "Adjusted Mean Distance: 4.022785432892366\n",
      "[10, 8, 10, 10, 10, 10, 9, 10, 9, 10]\n",
      "Number of contacts for electrode:  10\n",
      "LAm\n",
      "Number of contacts found:  10 \n",
      "\n",
      "LAHc\n",
      "Adjusted Mean Distance: 4.887619487260263\n",
      "Number of contacts found:  9 \n",
      "\n",
      "LPHc\n",
      "Number of contacts found:  10 \n",
      "\n",
      "RAIn\n",
      "Number of contacts found:  10 \n",
      "\n",
      "LAIn\n",
      "Number of contacts found:  10 \n",
      "\n",
      "LPIn\n",
      "Number of contacts found:  10 \n",
      "\n",
      "LOFr\n",
      "Adjusted Mean Distance: 5.002209551373188\n",
      "Number of contacts found:  9 \n",
      "\n",
      "LInfTe\n",
      "Number of contacts found:  10 \n",
      "\n",
      "LPTe\n",
      "Adjusted Mean Distance: 4.96340957275441\n",
      "Number of contacts found:  9 \n",
      "\n",
      "LPTeOc\n",
      "Number of contacts found:  10 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "for sub in subjects[1:2]: #.sort()\n",
    "    print(sub)\n",
    "    orig_pnms = pd.read_csv(f'{pnms_dir}/{sub}_prob_nms.fcsv', skiprows=3, header = None)\n",
    "    ct_t1_trans = np.loadtxt(f'{gt_dir}/{sub}/{sub}_desc-rigid_from-ct_to-T1w_type-ras_ses-post_xfm.txt')\n",
    "    ct = nib.load(f'./data/{sub}/{sub}_res-0p4mm_ct.nii.gz')\n",
    "    entry_target_path = f'./data/{sub}/{sub}_actual.fcsv'\n",
    "    \n",
    "    # Get affine to transform to voxels\n",
    "    inv_affine = np.linalg.inv(ct.affine)\n",
    "    data = np.asarray(ct.dataobj)\n",
    "\n",
    "    # Get entry-target coords\n",
    "    entry_target = extract_coords(entry_target_path)\n",
    "    # Transform to CT space\n",
    "    entry_target_coords = mne.transforms.apply_trans(np.linalg.inv(ct_t1_trans), entry_target[:,:-1].astype(float))\n",
    "    # Transform to voxels\n",
    "    entry_target_vox = np.round(mne.transforms.apply_trans(inv_affine, entry_target_coords)).astype(int)\n",
    "    \n",
    "    # Coordinates outputted by the network\n",
    "    coords_interest = orig_pnms.iloc[:,1:4].to_numpy()\n",
    "    # Get contacts positions in voxel space\n",
    "    transformed_coords = np.round(mne.transforms.apply_trans(inv_affine, coords_interest)).astype(int)\n",
    "    \n",
    "    # Empty final df\n",
    "    df_pos = pd.DataFrame(columns=['x', 'y', 'z', 'label_order'])\n",
    "    \n",
    "    \"\"\"\n",
    "    First stage: Masking filtering\n",
    "    \"\"\"\n",
    "    electrode_mask = nib.load(f'./data/{sub}/{sub}_res-0p4mm_desc-electrode_mask.nii.gz')\n",
    "    electrode_mask = electrode_mask.get_fdata()\n",
    "\n",
    "    filtered_coords = coords_interest[electrode_mask[transformed_coords[:,0], transformed_coords[:,1], transformed_coords[:,2]] > 0]\n",
    "\n",
    "    \"\"\"\n",
    "    Second stage: Line filtering\n",
    "    \"\"\"\n",
    "    # Run a first time to figure out the amount of contacts\n",
    "    # Compute for each pair of entry-target\n",
    "    n_contacts = []\n",
    "    contacts_coords = []\n",
    "    for i in range(0,entry_target.shape[0],2):  \n",
    "        # Label for the contact\n",
    "        label = entry_target[i,-1]\n",
    "        print(label, flush=True)        \n",
    "        found_contacts, _ = run_segmentation_qc_new(entry_point = entry_target_coords[i+1, :],\n",
    "                                                               target_point = entry_target_coords[i, :],\n",
    "                                                               contacts = filtered_coords, \n",
    "                                                               max_angle_target = 35, \n",
    "                                                               max_distance_prior = 2, \n",
    "                                                               min_distance_prior = 0.5, \n",
    "                                                               initial_spacing=4)\n",
    "        n_contacts.append(len(found_contacts))\n",
    "        contacts_coords.append(found_contacts)\n",
    "    # Find the max_n_contacts\n",
    "    print(n_contacts)\n",
    "    elements, count = np.unique(n_contacts, return_counts=True)\n",
    "    max_n_contacts = np.max(elements[count>1])\n",
    "    print('Number of contacts for electrode: ', max_n_contacts)\n",
    "    \n",
    "    # Now repeat to find the actual points\n",
    "    # Compute for each pair of entry-target\n",
    "    for i in range(0,entry_target.shape[0],2):  \n",
    "        # Label for the contact\n",
    "        label = entry_target[i,-1]\n",
    "        print(label, flush=True)        \n",
    "        # Only re-execute if len(found_contacts)<max_n_contacts\n",
    "        found_contacts = contacts_coords[i//2]\n",
    "        if len(found_contacts) < max_n_contacts:\n",
    "            found_contacts, avg_distance = run_segmentation_qc_new(entry_point = entry_target_coords[i+1, :],\n",
    "                                                               target_point = entry_target_coords[i, :],\n",
    "                                                               contacts = filtered_coords, \n",
    "                                                               max_angle_target = 35, \n",
    "                                                               max_distance_prior = 2, \n",
    "                                                               min_distance_prior = 0.5, \n",
    "                                                               initial_spacing=4,\n",
    "                                                               max_n_contacts=max_n_contacts)\n",
    "        \"\"\"\n",
    "        Third stage: Labelling\n",
    "        \"\"\"\n",
    "        found_contacts = np.array(found_contacts)\n",
    "        t1_coords = mne.transforms.apply_trans(ct_t1_trans, found_contacts)\n",
    "        df_tmp = pd.DataFrame(t1_coords, columns=['x','y','z'])\n",
    "        df_tmp['label_order'] = [f'{label}-{idx+1:02d}' for idx in range(found_contacts.shape[0])]\n",
    "        df_pos = pd.concat([df_pos, df_tmp])\n",
    "        print('Number of contacts found: ',  found_contacts.shape[0], '\\n')\n",
    "    \n",
    "\n",
    "\n",
    "    # df_to_fcsv(df_pos, final_fname_t1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "3746c023-9a82-4a2c-a27d-3822ef5c4d10",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([False,  True])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.unique(merged_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "c2975b14-d23b-4ad3-9821-ab5721907780",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_fname_t1 = f'./data/{sub}/test_fiducials_D129.fcsv'\n",
    "df_to_fcsv(df_pos, final_fname_t1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "892152c1-4854-4f8a-a03f-0fb8df45cc6b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>z</th>\n",
       "      <th>label_order</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>-23.279147</td>\n",
       "      <td>42.097845</td>\n",
       "      <td>12.778976</td>\n",
       "      <td>LAIn-10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-32.655666</td>\n",
       "      <td>32.171150</td>\n",
       "      <td>-10.073695</td>\n",
       "      <td>LPIn-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-31.676507</td>\n",
       "      <td>30.288086</td>\n",
       "      <td>-7.972420</td>\n",
       "      <td>LPIn-02</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>-30.697338</td>\n",
       "      <td>28.404030</td>\n",
       "      <td>-5.871277</td>\n",
       "      <td>LPIn-03</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>-29.516781</td>\n",
       "      <td>25.361386</td>\n",
       "      <td>-1.114054</td>\n",
       "      <td>LPIn-04</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>-28.537622</td>\n",
       "      <td>23.478321</td>\n",
       "      <td>0.987222</td>\n",
       "      <td>LPIn-05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>-28.299812</td>\n",
       "      <td>21.474331</td>\n",
       "      <td>3.931696</td>\n",
       "      <td>LPIn-06</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>-27.320653</td>\n",
       "      <td>19.591266</td>\n",
       "      <td>6.032972</td>\n",
       "      <td>LPIn-07</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>-27.146087</td>\n",
       "      <td>18.484227</td>\n",
       "      <td>8.294304</td>\n",
       "      <td>LPIn-08</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>-26.164128</td>\n",
       "      <td>16.204695</td>\n",
       "      <td>10.342604</td>\n",
       "      <td>LPIn-09</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-10.186166</td>\n",
       "      <td>59.066699</td>\n",
       "      <td>-13.727334</td>\n",
       "      <td>LOFr-01</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           x          y          z label_order\n",
       "9 -23.279147  42.097845  12.778976     LAIn-10\n",
       "0 -32.655666  32.171150 -10.073695     LPIn-01\n",
       "1 -31.676507  30.288086  -7.972420     LPIn-02\n",
       "2 -30.697338  28.404030  -5.871277     LPIn-03\n",
       "3 -29.516781  25.361386  -1.114054     LPIn-04\n",
       "4 -28.537622  23.478321   0.987222     LPIn-05\n",
       "5 -28.299812  21.474331   3.931696     LPIn-06\n",
       "6 -27.320653  19.591266   6.032972     LPIn-07\n",
       "7 -27.146087  18.484227   8.294304     LPIn-08\n",
       "8 -26.164128  16.204695  10.342604     LPIn-09\n",
       "0 -10.186166  59.066699 -13.727334     LOFr-01"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_pos[48:59]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3af5379-f790-4f13-901c-16a90afd7a72",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Example Test cases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e3110ad6-220a-439b-892e-5202f6294266",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[14.091, -7.65, 44.963, 'RPHc'],\n",
       "       [62.119, -14.953, 45.604, 'RPHc']], dtype=object)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "entry_target_path = './data/sub-P020_actual.fcsv'\n",
    "entry_target = extract_coords(entry_target_path)\n",
    "ind =8-4\n",
    "entry_target_interest = entry_target[ind:ind+2,:]\n",
    "# RAIn: 18-4\n",
    "# LAIn: 28-4\n",
    "# RPIn: 20-4\n",
    "# RPHc: 8-4\n",
    "entry_target_interest "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "09d955e4-4686-41d9-aed0-211e64ad5dbd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[15.675, -7.837, 44.617, 'RPHc1'],\n",
       "       [20.441, -8.295, 45.072, 'RPHc2'],\n",
       "       [25.245, -9.685, 44.893, 'RPHc3'],\n",
       "       [27.7, -7.364, 44.644, 'RPHc4'],\n",
       "       [30.453, -8.579, 45.087, 'RPHc5'],\n",
       "       [30.397, -10.541, 45.454, 'RPHc6'],\n",
       "       [34.437, -10.744, 44.991, 'RPHc7'],\n",
       "       [35.29, -8.793, 44.687, 'RPHc8'],\n",
       "       [39.988, -11.605, 45.584, 'RPHc9'],\n",
       "       [40.84, -9.654, 45.28, 'RPHc10'],\n",
       "       [43.534, -14.082, 45.921, 'RPHc11'],\n",
       "       [44.825, -11.819, 45.185, 'RPHc12'],\n",
       "       [49.579, -12.669, 45.714, 'RPHc13'],\n",
       "       [50.432, -10.718, 45.41, 'RPHc14'],\n",
       "       [55.19, -13.678, 45.524, 'RPHc15'],\n",
       "       [55.614, -11.647, 45.579, 'RPHc16'],\n",
       "       [59.171, -13.732, 45.842, 'RPHc17']], dtype=object)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coords = './data/sub-P020_space-T1w_desc-unet_pnms.fcsv'\n",
    "coords = extract_coords(coords)\n",
    "ind = 123-4\n",
    "coords_interest = coords[ind:ind+17,:]\n",
    "# RAIn: 69\n",
    "# LAIn: 0\n",
    "# RPIn: 140-4\n",
    "# RPHc: 123-4\n",
    "coords_interest"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3cb9475-c994-4de2-a611-44ebcb121529",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### First apply masking\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7508aeff-c17c-4cdb-9b7d-17c9b56c79b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mne"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5f84591c-a84e-4da9-a9d5-48d6b0ec3f44",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((4, 4), (512, 512, 140))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ct = nib.load('./data/sub-D128_space-T1w_desc-rigid_ses-post_ct.nii.gz')\n",
    "ct = nib.load('./data/sub-P020_desc-rigid_space-T1w_ct.nii.gz')\n",
    "# ct = nib.load('./data/sub-P087_space-T1w_desc-rigid_ses-post_ct.nii.gz')\n",
    "inv_affine = np.linalg.inv(ct.affine)\n",
    "data = np.asarray(ct.dataobj)\n",
    "inv_affine.shape, data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "42730054-8f1a-4b99-974b-4041e5e10133",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[221, 293,  58],\n",
       "       [119, 306,  58]])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Transform to voxels\n",
    "transformed_coords = np.round(mne.transforms.apply_trans(inv_affine, entry_target_interest[:,:-1].astype(float))).astype(int)\n",
    "transformed_coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "49fe53e6-1cbb-4e54-a54d-88367dfb31f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create mask\n",
    "img_shape = data.shape\n",
    "final_mask = np.zeros(img_shape).astype(bool)\n",
    "    \n",
    "for i in range(0,len(transformed_coords),2):\n",
    "    pointa = transformed_coords[i, :]\n",
    "    pointb = transformed_coords[i+1, :]\n",
    "\n",
    "\n",
    "    test_mask = create_line_mask(pointa, pointb, img_shape)\n",
    "\n",
    "    footprint = ball(4)\n",
    "    dilated = dilation(test_mask, footprint)\n",
    "    result = gaussian_filter(dilated.astype(float), sigma=0.6)\n",
    "    final_mask += result>0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "03b34363-b560-4924-a0cd-bed232bfba43",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([False,  True]), (512, 512, 140))"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Mask based on CT intensity\n",
    "mask_intensity = (data>2500).astype(bool)\n",
    "\n",
    "# Merge the masks\n",
    "merged_mask = final_mask*mask_intensity\n",
    "\n",
    "# Dilate a bit\n",
    "footprint = ball(2)\n",
    "merged_mask = dilation(merged_mask, footprint)\n",
    "\n",
    "np.unique(merged_mask), merged_mask.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "4493bf2c-156b-4481-bff5-bcea2ca09c17",
   "metadata": {},
   "outputs": [],
   "source": [
    "clipped_img = nib.Nifti1Image(merged_mask, ct.affine, ct.header)\n",
    "nib.save(clipped_img, 'test_mask_sub-P020.nii.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "dee447e4-f83e-4c25-b717-29ff4cf22e4b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[218, 293,  58],\n",
       "       [208, 294,  58],\n",
       "       [198, 297,  58]])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get contacts positions in voxel space\n",
    "transformed_coords = np.round(mne.transforms.apply_trans(inv_affine, coords_interest[:,:-1].astype(float))).astype(int)\n",
    "transformed_coords[:3,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "33dd3d81-ff74-473c-bec0-1d9957c50924",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((13, 4),\n",
       " array([ True,  True,  True, False,  True,  True,  True, False,  True,\n",
       "        False,  True,  True,  True, False,  True,  True,  True]))"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Mask the contacts inside the mask\n",
    "filtered_coords = coords_interest[merged_mask[transformed_coords[:,0], transformed_coords[:,1], transformed_coords[:,2]]]\n",
    "filtered_coords.shape, merged_mask[transformed_coords[:,0], transformed_coords[:,1], transformed_coords[:,2]]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcb0d4e9-f3e8-48c1-ad9c-97dc12a48622",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Now line algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "ab5320d2-2046-46c8-b6b7-1a9b71a7eff8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[14.091, -7.65, 44.963, 'RPHc'],\n",
       "       [62.119, -14.953, 45.604, 'RPHc']], dtype=object)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "entry_target_interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f9b5d501-76fd-477b-bb2d-1467e8758312",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[15.675, -7.837, 44.617, 'RPHc1'],\n",
       "       [20.441, -8.295, 45.072, 'RPHc2'],\n",
       "       [25.245, -9.685, 44.893, 'RPHc3'],\n",
       "       [30.453, -8.579, 45.087, 'RPHc5'],\n",
       "       [30.397, -10.541, 45.454, 'RPHc6'],\n",
       "       [34.437, -10.744, 44.991, 'RPHc7'],\n",
       "       [39.988, -11.605, 45.584, 'RPHc9'],\n",
       "       [43.534, -14.082, 45.921, 'RPHc11'],\n",
       "       [44.825, -11.819, 45.185, 'RPHc12'],\n",
       "       [49.579, -12.669, 45.714, 'RPHc13'],\n",
       "       [55.19, -13.678, 45.524, 'RPHc15'],\n",
       "       [55.614, -11.647, 45.579, 'RPHc16'],\n",
       "       [59.171, -13.732, 45.842, 'RPHc17']], dtype=object)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "filtered_coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "c636ad5a-c9d3-4c0b-b308-c3047cd771b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RPHc1 \n",
      "\n",
      "RPHc2 \n",
      "\n",
      "RPHc3 \n",
      "\n",
      "RPHc6 \n",
      "\n",
      "RPHc7 \n",
      "\n",
      "RPHc9 \n",
      "\n",
      "RPHc12 \n",
      "\n",
      "RPHc13 \n",
      "\n",
      "RPHc15 \n",
      "\n",
      "RPHc17 \n",
      "\n",
      "Adjusted Mean Distance: 4.911242544044423\n",
      "Found Contacts: [array([15.675, -7.837, 44.617], dtype=object), array([20.441, -8.295, 45.072], dtype=object), array([25.245, -9.685, 44.893], dtype=object), array([30.397, -10.541, 45.454], dtype=object), array([34.437, -10.744, 44.991], dtype=object), array([39.988, -11.605, 45.584], dtype=object), array([44.825, -11.819, 45.185], dtype=object), array([49.579, -12.669, 45.714], dtype=object), array([55.19, -13.678, 45.524], dtype=object), array([59.171, -13.732, 45.842], dtype=object)]\n",
      "Average Distance: 4.911242544044423\n"
     ]
    }
   ],
   "source": [
    "# Example usage\n",
    "target_point = entry_target_interest[0,:-1]  # Replace with actual entry point coordinates\n",
    "entry_point = entry_target_interest[1,:-1]  # Replace with actual target point coordinates\n",
    "contacts = filtered_coords  # Replace with your file path\n",
    "max_distance_target = 0.75  # Percentange from average \n",
    "min_distance_prior = 0.5 # at least half the average distance between contacts\n",
    "max_angle_target = 20\n",
    "max_distance_prior = 2 # not more than twice the average distance between contacts\n",
    "\n",
    "# found_contacts, avg_distance = run_segmentation_qc(entry_point, target_point, contacts, max_distance_target, min_distance_prior)\n",
    "found_contacts, avg_distance = run_segmentation_qc_new(entry_point,\n",
    "                                                        target_point,\n",
    "                                                        contacts, \n",
    "                                                        max_angle_target, \n",
    "                                                        max_distance_prior, \n",
    "                                                        min_distance_prior, \n",
    "                                                        initial_spacing=4)\n",
    "print(\"Found Contacts:\", found_contacts)\n",
    "print(\"Average Distance:\", avg_distance)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "a6df1a76-4657-4dd9-a474-2986f69b5940",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10, 13)"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(found_contacts), len(filtered_coords)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "02f6da75-e26c-4de1-b604-65556ac60524",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([15.675, -7.837, 44.617], dtype=object),\n",
       " array([20.441, -8.295, 45.072], dtype=object),\n",
       " array([25.245, -9.685, 44.893], dtype=object),\n",
       " array([30.397, -10.541, 45.454], dtype=object),\n",
       " array([34.437, -10.744, 44.991], dtype=object),\n",
       " array([39.988, -11.605, 45.584], dtype=object),\n",
       " array([44.825, -11.819, 45.185], dtype=object),\n",
       " array([49.579, -12.669, 45.714], dtype=object),\n",
       " array([55.19, -13.678, 45.524], dtype=object),\n",
       " array([59.171, -13.732, 45.842], dtype=object)]"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "found_contacts"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}