Source code for scenic.simulators.gta.center_detection

'''
This file contains helper functions
'''

import itertools
from typing import NamedTuple, Tuple
from copy import deepcopy

from scenic.simulators.gta.img_modf import *

'''
Compute neighbors of a given point/pixel
'''
def generate_neighbors(x):
    delta = np.arange(-1, 2)
    grid_delta = np.array(list(itertools.product(delta, delta)))
    return np.array(x)+ grid_delta


'''
Generate circle of a given radius and transformed by a center and a theta
'''
def generate_circle(radius, num_samples, theta, center):
    transformed_center = transform_center(center, radius, theta)
    thetas = np.linspace(0, 2*np.pi+0.01, num_samples)
    return transformed_center[0] + radius*np.cos(thetas), \
           transformed_center[1] + radius*np.sin(thetas)

'''
Transform the circle center as x+r*cos(90+theta) and y-r*sin(90+theta)
'''
def transform_center(center, radius, theta):
    return [center[0] + radius * np.cos(theta), center[1] - radius * np.sin(
        theta)]


'''
Compute image gradients from sobel edge detector
'''
def compute_gradient_sobel(img_data=None, img_file=None, threshold=100,
                     kernelsize=1, bw_kernelsize=3):
    assert img_data is not None or img_file is not None
    if img_data is None:
        img_data = Image.open(img_file)

    img_copy = img_data.copy()

    # Get the black and white image
    img_bw = convert_black_white(img_data=img_copy, img_file=img_file,
                                 threshold=threshold)
    cv_bw = cv2.cvtColor(np.array(img_bw), cv2.COLOR_RGB2BGR)

    # Gradients along x, y. The laplacian used for edge detection is actually
    #  a combination of these two
    sobelx = cv2.Sobel(cv_bw, cv2.CV_64F, 1, 0, ksize=kernelsize)
    sobely = cv2.Sobel(cv_bw, cv2.CV_64F, 0, 1, ksize=kernelsize)

    angle_sobel = np.arctan(sobelx / (sobely + 1e-6))

    edge_image = get_edges(img_data=img_copy, img_file=img_file,
                           threshold=threshold, kernelsize=bw_kernelsize)

    edge_x, edge_y = np.nonzero(np.asarray(edge_image))

    angle_along_edge = {}
    for x, y in zip(edge_x, edge_y):
        angle_along_edge[(x,y)] = angle_sobel[x][y][0]

    return angle_along_edge

'''
This is very coarse, avoid using this
'''
def compute_gradient_manual(img_data=None, img_file=None, threshold=100,
                     kernelsize=1):
    assert img_data is not None or img_file is not None
    if img_data is None:
        img_data = Image.open(img_file)

    img_copy = img_data.copy()

    # Directly work with the edges
    edge_image = get_edges(img_data=img_copy, img_file=img_file,
                           threshold=threshold, kernelsize=kernelsize)

    edge_x, edge_y = np.nonzero(np.asarray(edge_image))
    grad_x, grad_y = np.gradient(np.asarray(edge_image))
    angle_grad = np.arctan(grad_x/(grad_y+1e-6))

    angle_along_edge = {}
    for x, y in zip(edge_x, edge_y):
        angle_along_edge[(x,y)] = angle_grad[x][y]

    return angle_along_edge

'''
Compute all connected edges
'''

def generate_connected_edges(edge_locs):

    # This is a list of collected edges
    collected_edges = []
    while len(edge_locs) > 0:
        start_point = edge_locs.pop()
        new_edge = set({start_point})
        neighbors = set(tuple(map(tuple, generate_neighbors(start_point))))
        to_visit = neighbors.intersection(edge_locs)
        new_edge.update(neighbors.intersection(edge_locs))
        edge_locs.difference_update(new_edge.intersection(edge_locs))
        while len(to_visit) > 0:
            k = to_visit.pop()
            neighbors = set(tuple(map(tuple, generate_neighbors(k))))
            to_visit.update(neighbors.intersection(edge_locs))
            new_edge.update(neighbors.intersection(edge_locs))
            edge_locs.difference_update(new_edge.intersection(edge_locs))

        collected_edges.append(new_edge)
    return collected_edges

'''
To find the center of the road perpendicular to a given point
num_samples is function which can potentially change with the radius of the
circle
'''

[docs]def find_center(x, theta, collected_edges, all_edges, num_samples, bw_image): ''' Find which edge x lies in ''' for edge in collected_edges: if tuple([x[1], x[0]]) in edge: x_edge = edge break possible_opp_edge = all_edges.difference(x_edge) r = 1 ''' Find direction of computing the growing circles ''' if theta < 0: theta = theta + 2*np.pi theta_plus = theta + np.pi/2. theta_minus = theta - np.pi / 2. center_plus = transform_center(center=x, radius=2, theta=theta_plus) center_plus = np.array(np.around(center_plus), dtype=int) if center_plus[0] >= bw_image.shape[1]: center_plus[0] = bw_image.shape[1] - 1 if center_plus[1] >= bw_image.shape[0]: center_plus[1] = bw_image.shape[0] - 1 if bw_image[center_plus[1]][center_plus[0]] == 0: theta = theta_plus tangent_theta = theta - np.pi/2. else: theta = theta_minus tangent_theta = theta + np.pi / 2. while True: circle_x, circle_y = generate_circle(radius=r, center=x, theta=theta, num_samples=num_samples(r)) int_circle = set(itertools.product(np.array(np.around(circle_y), dtype=int), np.array(np.around(circle_x), dtype=int))) if len(possible_opp_edge.intersection(int_circle)) > 0: opp_edge = possible_opp_edge.intersection(int_circle) break r += 1 if len(opp_edge) > 1: arg_min = np.linalg.norm(np.array(list(opp_edge)) - np.array([x[1], x[0]]), axis=1).argmin() opp_point = np.array(list(opp_edge))[arg_min] else: opp_point = opp_edge.pop() mid_loc = np.array(np.around((np.array([x[1],x[0]]) + np.array(opp_point))/2.), dtype=int) orig_theta = tangent_theta # When there is a clear greater than and less than if mid_loc[0] > x[1] and mid_loc[1] < x[0]: if tangent_theta > np.pi: tangent_theta = tangent_theta - np.pi if tangent_theta < np.pi/2.: tangent_theta = tangent_theta + np.pi elif mid_loc[0] > x[1] and mid_loc[1] > x[0]: if tangent_theta < np.pi: tangent_theta = tangent_theta + np.pi elif mid_loc[0] < x[1] and mid_loc[1] < x[0]: if tangent_theta > np.pi/2.: tangent_theta = tangent_theta - np.pi elif mid_loc[0] < x[1] and mid_loc[1] > x[0]: if tangent_theta < 3.*np.pi/2. and tangent_theta > 0: tangent_theta = tangent_theta + np.pi # When there exists one equality if mid_loc[0] == x[1]: if mid_loc[1] < x[0]: if tangent_theta > np.pi: tangent_theta = tangent_theta - np.pi else: if tangent_theta < 3./2.*np.pi: tangent_theta = tangent_theta + np.pi if mid_loc[1] == x[0]: if mid_loc[0] < x[1]: if tangent_theta > np.pi/2.: tangent_theta = tangent_theta - np.pi else: if tangent_theta < np.pi: tangent_theta = tangent_theta + np.pi if tangent_theta < 0: tangent_theta = tangent_theta + 2 * np.pi return orig_theta, tangent_theta, tuple(opp_point), tuple(mid_loc)
''' Update all edges with the bounding box of the image ''' def compute_bb(img_dimensions): im_x_min = -1 im_y_min = -1 im_x_max = img_dimensions[1] im_y_max = img_dimensions[0] ''' Compute the bounding box ''' bb = set() bb.update({(im_x_min, y) for y in range(im_y_min, im_y_max+1)}) bb.update({(im_x_max, y) for y in range(im_y_min, im_y_max + 1)}) bb.update({(x, im_y_min) for x in range(im_x_min, im_x_max + 1)}) bb.update({(x, im_y_max) for x in range(im_x_min, im_x_max + 1)}) return bb
[docs]class EdgeData(NamedTuple): init_theta: float tangent: float opp_loc: Tuple[float, float] mid_loc: Tuple[float, float]
''' Compute the dictionary with key is the start state and the values are a tuple consisting of the angle, mid point and the the opposite edge ''' def compute_midpoints(img_data, bw_kernelsize=1, kernelsize=3, threshold=100, num_samples = lambda r:16): # First compute the black and white image and store it img_bw = convert_black_white(img_data=img_data) img_bw = img_bw.convert('L') img_bw_int = np.array(img_bw, dtype=int) # Compute the point to the tangent angle dictionary cae = compute_gradient_sobel(img_data=img_data, bw_kernelsize=bw_kernelsize, kernelsize=kernelsize, threshold=threshold) # Collect all edges all_edges = set(cae.keys()) # Collect all roads connected_edges = generate_connected_edges(edge_locs=deepcopy(all_edges)) # Compute bounding box bb = compute_bb(img_dimensions=img_data.size) all_edges.update(bb) # Compute midpoints point_data = {} for edge in connected_edges: for x in edge: init_theta, theta, opp_loc, mid_loc = find_center(x=[x[1], x[0]], theta=cae[x], collected_edges=connected_edges, all_edges=all_edges, num_samples=num_samples, bw_image=img_bw_int) point_data[x] = EdgeData(init_theta=init_theta, tangent=theta, opp_loc=opp_loc, mid_loc=mid_loc \ if opp_loc not in bb else (np.nan, np.nan)) return point_data ''' Compute the heading at a random sample ''' def compute_heading(x, point_data): road_edges = tuple(point_data.keys()) edge_distance = np.linalg.norm(np.array(list(road_edges)) - np.array(x),1, axis=1) closest_edge = list(road_edges)[edge_distance.argmin()] return closest_edge, point_data[closest_edge] ''' Sampling from road ''' def sample_from_road(img_data, num_samples): img_bw = convert_black_white(img_data=img_data) img_bw = img_bw.convert('L') img_bw_int = np.array(img_bw, dtype=int) road_y, road_x = np.where(img_bw_int == 0) all_roads = list(zip(road_y, road_x)) locs = np.random.choice(len(all_roads), num_samples, replace=False) return np.array(all_roads)[locs]