Skip to content
Snippets Groups Projects
live_wire.py 5.38 KiB
Newer Older
  • Learn to ignore specific revisions
  • Christian's avatar
    Christian committed
    import time
    import cv2
    import numpy as np
    import matplotlib.pyplot as plt
    
    from skimage import exposure
    
    Christian's avatar
    Christian committed
    from skimage.filters import gaussian
    from skimage.feature import canny
    
    from skimage.graph import route_through_array
    
    from scipy.signal import convolve2d
    
    '''
    ### Canny Edge cost image
    
    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
    '''
    
    ### Disk live wire cost image
    def compute_cost_image(path, sigma=3, disk_size=15):
    
        ### 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)
    
        # Do disk thing
        binary_img = canny_img
        k_size = 17
        kernel = circle_edge_kernel(k_size=disk_size)
        convolved = convolve2d(binary_img, kernel, mode='same', boundary='fill')
    
        # Create cost image
        cost_img = (convolved.max() - convolved)**4  # 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 circle_edge_kernel(k_size=5, radius=None):
        """
        Create a k_size x k_size array whose values increase linearly
        from 0 at the center to 1 at the circle boundary (radius).
    
        Parameters
        ----------
        k_size : int
            The size (width and height) of the kernel array.
        radius : float, optional
            The circle's radius. By default, set to (k_size-1)/2.
    
        Returns
        -------
        kernel : 2D numpy array of shape (k_size, k_size)
            The circle-edge-weighted kernel.
        """
        if radius is None:
            # By default, let the radius be half the kernel size
            radius = (k_size - 1) / 2
    
        # Create an empty kernel
        kernel = np.zeros((k_size, k_size), dtype=float)
    
        # Coordinates of the center
        center = radius  # same as (k_size-1)/2 if radius is default
    
        # Fill the kernel
        for y in range(k_size):
            for x in range(k_size):
                dist = np.sqrt((x - center)**2 + (y - center)**2)
                if dist <= radius:
                    # Weight = distance / radius => 0 at center, 1 at boundary
                    kernel[y, x] = dist / radius
    
        return kernel
    
    
    
    
    
    
    Christian's avatar
    Christian committed
    def downscale(img, points, scale_percent):
    
        """
        Downsample `img` to `scale_percent` size and scale the given points accordingly.
        Returns (downsampled_img, (scaled_seed, scaled_target)).
        """
    
    Christian's avatar
    Christian committed
        if scale_percent == 100:
            return img, (tuple(points[0]), tuple(points[1]))
        else:
    
            # Compute new dimensions
    
    Christian's avatar
    Christian committed
            width = int(img.shape[1] * scale_percent / 100)
            height = int(img.shape[0] * scale_percent / 100)
            new_dimensions = (width, height)
    
    
            # Downsample
    
    Christian's avatar
    Christian committed
            downsampled_img = cv2.resize(img, new_dimensions, interpolation=cv2.INTER_AREA)
    
            # Scaling factors
    
            scale_x = width / img.shape[1]
            scale_y = height / img.shape[0]
    
            # Scale the points (x, y)
    
    Christian's avatar
    Christian committed
            seed_xy = tuple(points[0])
            target_xy = tuple(points[1])
            scaled_seed_xy = (int(seed_xy[0] * scale_x), int(seed_xy[1] * scale_y))
            scaled_target_xy = (int(target_xy[0] * scale_x), int(target_xy[1] * scale_y))
    
            return downsampled_img, (scaled_seed_xy, scaled_target_xy)
    
    
    def compute_cost(image, sigma=3.0, epsilon=1e-5):
        """
        Smooth the image, run Canny edge detection, then invert the edge map into a cost image.
        """
    
    
        # 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)
    
        cost_img = 1.0 / (canny_img + epsilon)  # Invert edges: higher cost where edges are stronger
    
        return cost_img, canny_img
    
    def backtrack_pixels_on_image(img_color, path_coords, bgr_color=(0, 0, 255)):
        """
        Color the path on the (already converted BGR) image in the specified color.
        `path_coords` should be a list of (row, col) or (y, x).
        """
        for (row, col) in path_coords:
            img_color[row, col] = bgr_color
        return img_color
    
    Christian's avatar
    Christian committed
    def export_path(path_coords, path_name):
        """
        Export the path to a np array.
        """
        np.save(path_name, path_coords)
        return None