Source code for scenic.simulators.gta.center_detection

"""
This file contains helper functions
"""

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

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


"""
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.0 theta_minus = theta - np.pi / 2.0 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.0 else: theta = theta_minus tangent_theta = theta + np.pi / 2.0 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.0), 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.0: 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.0: tangent_theta = tangent_theta - np.pi elif mid_loc[0] < x[1] and mid_loc[1] > x[0]: if tangent_theta < 3.0 * np.pi / 2.0 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.0 / 2.0 * 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.0: 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] = float(theta) return point_data