diff --git a/live_wire.py b/live_wire.py
index 46eb2dd0a8ba0e438d04b7c5a1055c1a4e4bba46..9b49d97f03a4f363c2b32792dd8ae9ba527f8b2f 100644
--- a/live_wire.py
+++ b/live_wire.py
@@ -9,23 +9,44 @@ from skimage.graph import route_through_array
 
 #### Helper functions ####
 
-def load_image(path, type):
-    """
-    Load an image in either gray or color mode (then convert color to gray).
-    """
-    if type == 'gray':
-        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
-        if img is None:
-            raise FileNotFoundError(f"Could not read {path}")
-    elif type == 'color':
-        img = cv2.imread(path, cv2.IMREAD_COLOR)
-        if img is None:
-            raise FileNotFoundError(f"Could not read {path}")
-        else:
-            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    else:
-        raise ValueError("type must be 'gray' or 'color'")
-    return img
+def compute_cost_image(path, sigma=3):
+
+    ### Load image
+    image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+    
+    # Apply histogram equalization
+    image_contrasted = exposure.equalize_adapthist(image, clip_limit=0.01)
+
+    # Apply smoothing
+    smoothed_img = gaussian(image_contrasted, sigma=sigma)
+
+    # Apply Canny edge detection
+    canny_img = canny(smoothed_img)
+
+    # Create cost image
+    cost_img = 1.0 / (canny_img + 1e-5)  # Invert edges: higher cost where edges are stronger
+
+    return cost_img
+
+
+def find_path(cost_image, points):
+
+    if len(points) != 2:
+        raise ValueError("Points should be a list of 2 points: seed and target.")
+    
+    seed_rc, target_rc = points
+
+    path_rc, cost = route_through_array(
+        cost_image, 
+        start=seed_rc, 
+        end=target_rc, 
+        fully_connected=True
+    )
+
+    return path_rc
+
+
+
 
 def downscale(img, points, scale_percent):
     """
@@ -89,68 +110,3 @@ def export_path(path_coords, path_name):
     """
     np.save(path_name, path_coords)
     return None
-
-
-#### Main Script ####
-def main():
-    # Define input parameters
-    image_path = 'agamodon_slice.png'
-    image_type = 'gray'        # 'gray' or 'color'
-    downscale_factor = 100     # % of original size
-    points_path = 'agamodonPoints.npy'
-
-    # Load image
-    image = load_image(image_path, image_type)
-
-    # Load seed and target points
-    points = np.int0(np.round(np.load(points_path)))  # shape: (2, 2), i.e. [[x_seed, y_seed], [x_target, y_target]]
-
-    # Downscale image and points
-    scaled_image, scaled_points = downscale(image, points, downscale_factor)
-    seed, target = scaled_points  # Each is (x, y)
-
-    # Convert to row,col for scikit-image (which uses (row, col) = (y, x))
-    seed_rc = (seed[1], seed[0])
-    target_rc = (target[1], target[0])
-
-    # Compute cost image
-    cost_image, canny_img = compute_cost(scaled_image)
-
-    # Find path using route_through_array
-    # route_through_array expects: route_through_array(image, start, end, fully_connected=True/False)
-    start_time = time.time()
-    path_rc, cost = route_through_array(
-        cost_image, 
-        start=seed_rc, 
-        end=target_rc, 
-        fully_connected=True
-    )
-    end_time = time.time()
-
-    print(f"Elapsed time for pathfinding: {end_time - start_time:.3f} seconds")
-
-    # Convert single-channel image to BGR for coloring
-    color_img = cv2.cvtColor(scaled_image, cv2.COLOR_GRAY2BGR)
-
-    # Draw path. `path_rc` is a list of (row, col).
-    # If you want to mark it in red, do (0,0,255) because OpenCV uses BGR format.
-    color_img = backtrack_pixels_on_image(color_img, path_rc, bgr_color=(0, 0, 255))
-
-
-    # Export path
-    export_path(path_rc, 'agamodonPath.npy')
-
-    # Display results
-    plt.figure(figsize=(20, 8))
-    plt.subplot(1, 2, 1)
-    plt.title("Cost image")
-    plt.imshow(cost_image, cmap='gray')
-
-    plt.subplot(1, 2, 2)
-    plt.title("Path from Seed to Target")
-    # Convert BGR->RGB for pyplot
-    plt.imshow(color_img[..., ::-1])
-    plt.show()
-
-if __name__ == "__main__":
-    main()
\ No newline at end of file