""" Code adapted from https://github.com/rpautrat/SuperPoint Module used to generate geometrical synthetic shapes """ import math import cv2 as cv import numpy as np import shapely.geometry from itertools import combinations random_state = np.random.RandomState(None) def set_random_state(state): global random_state random_state = state def get_random_color(background_color): """Output a random scalar in grayscale with a least a small contrast with the background color.""" color = random_state.randint(256) if abs(color - background_color) < 30: # not enough contrast color = (color + 128) % 256 return color def get_different_color(previous_colors, min_dist=50, max_count=20): """Output a color that contrasts with the previous colors. Parameters: previous_colors: np.array of the previous colors min_dist: the difference between the new color and the previous colors must be at least min_dist max_count: maximal number of iterations """ color = random_state.randint(256) count = 0 while np.any(np.abs(previous_colors - color) < min_dist) and count < max_count: count += 1 color = random_state.randint(256) return color def add_salt_and_pepper(img): """Add salt and pepper noise to an image.""" noise = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8) cv.randu(noise, 0, 255) black = noise < 30 white = noise > 225 img[white > 0] = 255 img[black > 0] = 0 cv.blur(img, (5, 5), img) return np.empty((0, 2), dtype=np.int) def generate_background( size=(960, 1280), nb_blobs=100, min_rad_ratio=0.01, max_rad_ratio=0.05, min_kernel_size=50, max_kernel_size=300, ): """Generate a customized background image. Parameters: size: size of the image nb_blobs: number of circles to draw min_rad_ratio: the radius of blobs is at least min_rad_size * max(size) max_rad_ratio: the radius of blobs is at most max_rad_size * max(size) min_kernel_size: minimal size of the kernel max_kernel_size: maximal size of the kernel """ img = np.zeros(size, dtype=np.uint8) dim = max(size) cv.randu(img, 0, 255) cv.threshold(img, random_state.randint(256), 255, cv.THRESH_BINARY, img) background_color = int(np.mean(img)) blobs = np.concatenate( [ random_state.randint(0, size[1], size=(nb_blobs, 1)), random_state.randint(0, size[0], size=(nb_blobs, 1)), ], axis=1, ) for i in range(nb_blobs): col = get_random_color(background_color) cv.circle( img, (blobs[i][0], blobs[i][1]), np.random.randint(int(dim * min_rad_ratio), int(dim * max_rad_ratio)), col, -1, ) kernel_size = random_state.randint(min_kernel_size, max_kernel_size) cv.blur(img, (kernel_size, kernel_size), img) return img def generate_custom_background( size, background_color, nb_blobs=3000, kernel_boundaries=(50, 100) ): """Generate a customized background to fill the shapes. Parameters: background_color: average color of the background image nb_blobs: number of circles to draw kernel_boundaries: interval of the possible sizes of the kernel """ img = np.zeros(size, dtype=np.uint8) img = img + get_random_color(background_color) blobs = np.concatenate( [ np.random.randint(0, size[1], size=(nb_blobs, 1)), np.random.randint(0, size[0], size=(nb_blobs, 1)), ], axis=1, ) for i in range(nb_blobs): col = get_random_color(background_color) cv.circle(img, (blobs[i][0], blobs[i][1]), np.random.randint(20), col, -1) kernel_size = np.random.randint(kernel_boundaries[0], kernel_boundaries[1]) cv.blur(img, (kernel_size, kernel_size), img) return img def final_blur(img, kernel_size=(5, 5)): """Gaussian blur applied to an image. Parameters: kernel_size: size of the kernel """ cv.GaussianBlur(img, kernel_size, 0, img) def ccw(A, B, C, dim): """Check if the points are listed in counter-clockwise order.""" if dim == 2: # only 2 dimensions return (C[:, 1] - A[:, 1]) * (B[:, 0] - A[:, 0]) > (B[:, 1] - A[:, 1]) * ( C[:, 0] - A[:, 0] ) else: # dim should be equal to 3 return (C[:, 1, :] - A[:, 1, :]) * (B[:, 0, :] - A[:, 0, :]) > ( B[:, 1, :] - A[:, 1, :] ) * (C[:, 0, :] - A[:, 0, :]) def intersect(A, B, C, D, dim): """Return true if line segments AB and CD intersect""" return np.any( (ccw(A, C, D, dim) != ccw(B, C, D, dim)) & (ccw(A, B, C, dim) != ccw(A, B, D, dim)) ) def keep_points_inside(points, size): """Keep only the points whose coordinates are inside the dimensions of the image of size 'size'""" mask = ( (points[:, 0] >= 0) & (points[:, 0] < size[1]) & (points[:, 1] >= 0) & (points[:, 1] < size[0]) ) return points[mask, :] def get_unique_junctions(segments, min_label_len): """Get unique junction points from line segments.""" # Get all junctions from segments junctions_all = np.concatenate((segments[:, :2], segments[:, 2:]), axis=0) if junctions_all.shape[0] == 0: junc_points = None line_map = None # Get all unique junction points else: junc_points = np.unique(junctions_all, axis=0) # Generate line map from points and segments line_map = get_line_map(junc_points, segments) return junc_points, line_map def get_line_map(points: np.ndarray, segments: np.ndarray) -> np.ndarray: """Get line map given the points and segment sets.""" # create empty line map num_point = points.shape[0] line_map = np.zeros([num_point, num_point]) # Iterate through every segment for idx in range(segments.shape[0]): # Get the junctions from a single segement seg = segments[idx, :] junction1 = seg[:2] junction2 = seg[2:] # Get index idx_junction1 = np.where((points == junction1).sum(axis=1) == 2)[0] idx_junction2 = np.where((points == junction2).sum(axis=1) == 2)[0] # label the corresponding entries line_map[idx_junction1, idx_junction2] = 1 line_map[idx_junction2, idx_junction1] = 1 return line_map def get_line_heatmap(junctions, line_map, size=[480, 640], thickness=1): """Get line heat map from junctions and line map.""" # Make sure that the thickness is 1 if not isinstance(thickness, int): thickness = int(thickness) # If the junction points are not int => round them and convert to int if not junctions.dtype == np.int: junctions = (np.round(junctions)).astype(np.int) # Initialize empty map heat_map = np.zeros(size) if junctions.shape[0] > 0: # If empty, just return zero map # Iterate through all the junctions for idx in range(junctions.shape[0]): # if no connectivity, just skip it if line_map[idx, :].sum() == 0: continue # Plot the line segment else: # Iterate through all the connected junctions for idx2 in np.where(line_map[idx, :] == 1)[0]: point1 = junctions[idx, :] point2 = junctions[idx2, :] # Draw line cv.line(heat_map, tuple(point1), tuple(point2), 1.0, thickness) return heat_map def draw_lines(img, nb_lines=10, min_len=32, min_label_len=32): """Draw random lines and output the positions of the pair of junctions and line associativities. Parameters: nb_lines: maximal number of lines """ # Set line number and points placeholder num_lines = random_state.randint(1, nb_lines) segments = np.empty((0, 4), dtype=np.int) points = np.empty((0, 2), dtype=np.int) background_color = int(np.mean(img)) min_dim = min(img.shape) # Convert length constrain to pixel if given float number if isinstance(min_len, float) and min_len <= 1.0: min_len = int(min_dim * min_len) if isinstance(min_label_len, float) and min_label_len <= 1.0: min_label_len = int(min_dim * min_label_len) # Generate lines one by one for i in range(num_lines): x1 = random_state.randint(img.shape[1]) y1 = random_state.randint(img.shape[0]) p1 = np.array([[x1, y1]]) x2 = random_state.randint(img.shape[1]) y2 = random_state.randint(img.shape[0]) p2 = np.array([[x2, y2]]) # Check the length of the line line_length = np.sqrt(np.sum((p1 - p2) ** 2)) if line_length < min_len: continue # Check that there is no overlap if intersect(segments[:, 0:2], segments[:, 2:4], p1, p2, 2): continue col = get_random_color(background_color) thickness = random_state.randint(min_dim * 0.01, min_dim * 0.02) cv.line(img, (x1, y1), (x2, y2), col, thickness) # Only record the segments longer than min_label_len seg_len = math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2) if seg_len >= min_label_len: segments = np.concatenate([segments, np.array([[x1, y1, x2, y2]])], axis=0) points = np.concatenate([points, np.array([[x1, y1], [x2, y2]])], axis=0) # If no line is drawn, recursively call the function if points.shape[0] == 0: return draw_lines(img, nb_lines, min_len, min_label_len) # Get the line associativity map line_map = get_line_map(points, segments) return {"points": points, "line_map": line_map} def check_segment_len(segments, min_len=32): """Check if one of the segments is too short (True means too short).""" point1_vec = segments[:, :2] point2_vec = segments[:, 2:] diff = point1_vec - point2_vec dist = np.sqrt(np.sum(diff**2, axis=1)) if np.any(dist < min_len): return True else: return False def draw_polygon(img, max_sides=8, min_len=32, min_label_len=64): """Draw a polygon with a random number of corners and return the position of the junctions + line map. Parameters: max_sides: maximal number of sides + 1 """ num_corners = random_state.randint(3, max_sides) min_dim = min(img.shape[0], img.shape[1]) rad = max(random_state.rand() * min_dim / 2, min_dim / 10) # Center of a circle x = random_state.randint(rad, img.shape[1] - rad) y = random_state.randint(rad, img.shape[0] - rad) # Convert length constrain to pixel if given float number if isinstance(min_len, float) and min_len <= 1.0: min_len = int(min_dim * min_len) if isinstance(min_label_len, float) and min_label_len <= 1.0: min_label_len = int(min_dim * min_label_len) # Sample num_corners points inside the circle slices = np.linspace(0, 2 * math.pi, num_corners + 1) angles = [ slices[i] + random_state.rand() * (slices[i + 1] - slices[i]) for i in range(num_corners) ] points = np.array( [ [ int(x + max(random_state.rand(), 0.4) * rad * math.cos(a)), int(y + max(random_state.rand(), 0.4) * rad * math.sin(a)), ] for a in angles ] ) # Filter the points that are too close or that have an angle too flat norms = [ np.linalg.norm(points[(i - 1) % num_corners, :] - points[i, :]) for i in range(num_corners) ] mask = np.array(norms) > 0.01 points = points[mask, :] num_corners = points.shape[0] corner_angles = [ angle_between_vectors( points[(i - 1) % num_corners, :] - points[i, :], points[(i + 1) % num_corners, :] - points[i, :], ) for i in range(num_corners) ] mask = np.array(corner_angles) < (2 * math.pi / 3) points = points[mask, :] num_corners = points.shape[0] # Get junction pairs from points segments = np.zeros([0, 4]) # Used to record all the segments no matter we are going to label it or not. segments_raw = np.zeros([0, 4]) for idx in range(num_corners): if idx == (num_corners - 1): p1 = points[idx] p2 = points[0] else: p1 = points[idx] p2 = points[idx + 1] segment = np.concatenate((p1, p2), axis=0) # Only record the segments longer than min_label_len seg_len = np.sqrt(np.sum((p1 - p2) ** 2)) if seg_len >= min_label_len: segments = np.concatenate((segments, segment[None, ...]), axis=0) segments_raw = np.concatenate((segments_raw, segment[None, ...]), axis=0) # If not enough corner, just regenerate one if (num_corners < 3) or check_segment_len(segments_raw, min_len): return draw_polygon(img, max_sides, min_len, min_label_len) # Get junctions from segments junctions_all = np.concatenate((segments[:, :2], segments[:, 2:]), axis=0) if junctions_all.shape[0] == 0: junc_points = None line_map = None else: junc_points = np.unique(junctions_all, axis=0) # Get the line map line_map = get_line_map(junc_points, segments) corners = points.reshape((-1, 1, 2)) col = get_random_color(int(np.mean(img))) cv.fillPoly(img, [corners], col) return {"points": junc_points, "line_map": line_map} def overlap(center, rad, centers, rads): """Check that the circle with (center, rad) doesn't overlap with the other circles.""" flag = False for i in range(len(rads)): if np.linalg.norm(center - centers[i]) < rad + rads[i]: flag = True break return flag def angle_between_vectors(v1, v2): """Compute the angle (in rad) between the two vectors v1 and v2.""" v1_u = v1 / np.linalg.norm(v1) v2_u = v2 / np.linalg.norm(v2) return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0)) def draw_multiple_polygons( img, max_sides=8, nb_polygons=30, min_len=32, min_label_len=64, safe_margin=5, **extra ): """Draw multiple polygons with a random number of corners and return the junction points + line map. Parameters: max_sides: maximal number of sides + 1 nb_polygons: maximal number of polygons """ segments = np.empty((0, 4), dtype=np.int) label_segments = np.empty((0, 4), dtype=np.int) centers = [] rads = [] points = np.empty((0, 2), dtype=np.int) background_color = int(np.mean(img)) min_dim = min(img.shape[0], img.shape[1]) # Convert length constrain to pixel if given float number if isinstance(min_len, float) and min_len <= 1.0: min_len = int(min_dim * min_len) if isinstance(min_label_len, float) and min_label_len <= 1.0: min_label_len = int(min_dim * min_label_len) if isinstance(safe_margin, float) and safe_margin <= 1.0: safe_margin = int(min_dim * safe_margin) # Sequentially generate polygons for i in range(nb_polygons): num_corners = random_state.randint(3, max_sides) min_dim = min(img.shape[0], img.shape[1]) # Also add the real radius rad = max(random_state.rand() * min_dim / 2, min_dim / 9) rad_real = rad - safe_margin # Center of a circle x = random_state.randint(rad, img.shape[1] - rad) y = random_state.randint(rad, img.shape[0] - rad) # Sample num_corners points inside the circle slices = np.linspace(0, 2 * math.pi, num_corners + 1) angles = [ slices[i] + random_state.rand() * (slices[i + 1] - slices[i]) for i in range(num_corners) ] # Sample outer points and inner points new_points = [] new_points_real = [] for a in angles: x_offset = max(random_state.rand(), 0.4) y_offset = max(random_state.rand(), 0.4) new_points.append( [ int(x + x_offset * rad * math.cos(a)), int(y + y_offset * rad * math.sin(a)), ] ) new_points_real.append( [ int(x + x_offset * rad_real * math.cos(a)), int(y + y_offset * rad_real * math.sin(a)), ] ) new_points = np.array(new_points) new_points_real = np.array(new_points_real) # Filter the points that are too close or that have an angle too flat norms = [ np.linalg.norm(new_points[(i - 1) % num_corners, :] - new_points[i, :]) for i in range(num_corners) ] mask = np.array(norms) > 0.01 new_points = new_points[mask, :] new_points_real = new_points_real[mask, :] num_corners = new_points.shape[0] corner_angles = [ angle_between_vectors( new_points[(i - 1) % num_corners, :] - new_points[i, :], new_points[(i + 1) % num_corners, :] - new_points[i, :], ) for i in range(num_corners) ] mask = np.array(corner_angles) < (2 * math.pi / 3) new_points = new_points[mask, :] new_points_real = new_points_real[mask, :] num_corners = new_points.shape[0] # Not enough corners if num_corners < 3: continue # Segments for checking overlap (outer circle) new_segments = np.zeros((1, 4, num_corners)) new_segments[:, 0, :] = [new_points[i][0] for i in range(num_corners)] new_segments[:, 1, :] = [new_points[i][1] for i in range(num_corners)] new_segments[:, 2, :] = [ new_points[(i + 1) % num_corners][0] for i in range(num_corners) ] new_segments[:, 3, :] = [ new_points[(i + 1) % num_corners][1] for i in range(num_corners) ] # Segments to record (inner circle) new_segments_real = np.zeros((1, 4, num_corners)) new_segments_real[:, 0, :] = [new_points_real[i][0] for i in range(num_corners)] new_segments_real[:, 1, :] = [new_points_real[i][1] for i in range(num_corners)] new_segments_real[:, 2, :] = [ new_points_real[(i + 1) % num_corners][0] for i in range(num_corners) ] new_segments_real[:, 3, :] = [ new_points_real[(i + 1) % num_corners][1] for i in range(num_corners) ] # Check that the polygon will not overlap with pre-existing shapes if intersect( segments[:, 0:2, None], segments[:, 2:4, None], new_segments[:, 0:2, :], new_segments[:, 2:4, :], 3, ) or overlap(np.array([x, y]), rad, centers, rads): continue # Check that the the edges of the polygon is not too short if check_segment_len(new_segments_real, min_len): continue # If the polygon is valid, append it to the polygon set centers.append(np.array([x, y])) rads.append(rad) new_segments = np.reshape(np.swapaxes(new_segments, 0, 2), (-1, 4)) segments = np.concatenate([segments, new_segments], axis=0) # Only record the segments longer than min_label_len new_segments_real = np.reshape(np.swapaxes(new_segments_real, 0, 2), (-1, 4)) points1 = new_segments_real[:, :2] points2 = new_segments_real[:, 2:] seg_len = np.sqrt(np.sum((points1 - points2) ** 2, axis=1)) new_label_segment = new_segments_real[seg_len >= min_label_len, :] label_segments = np.concatenate([label_segments, new_label_segment], axis=0) # Color the polygon with a custom background corners = new_points_real.reshape((-1, 1, 2)) mask = np.zeros(img.shape, np.uint8) custom_background = generate_custom_background( img.shape, background_color, **extra ) cv.fillPoly(mask, [corners], 255) locs = np.where(mask != 0) img[locs[0], locs[1]] = custom_background[locs[0], locs[1]] points = np.concatenate([points, new_points], axis=0) # Get all junctions from label segments junctions_all = np.concatenate( (label_segments[:, :2], label_segments[:, 2:]), axis=0 ) if junctions_all.shape[0] == 0: junc_points = None line_map = None else: junc_points = np.unique(junctions_all, axis=0) # Generate line map from points and segments line_map = get_line_map(junc_points, label_segments) return {"points": junc_points, "line_map": line_map} def draw_ellipses(img, nb_ellipses=20): """Draw several ellipses. Parameters: nb_ellipses: maximal number of ellipses """ centers = np.empty((0, 2), dtype=np.int) rads = np.empty((0, 1), dtype=np.int) min_dim = min(img.shape[0], img.shape[1]) / 4 background_color = int(np.mean(img)) for i in range(nb_ellipses): ax = int(max(random_state.rand() * min_dim, min_dim / 5)) ay = int(max(random_state.rand() * min_dim, min_dim / 5)) max_rad = max(ax, ay) x = random_state.randint(max_rad, img.shape[1] - max_rad) # center y = random_state.randint(max_rad, img.shape[0] - max_rad) new_center = np.array([[x, y]]) # Check that the ellipsis will not overlap with pre-existing shapes diff = centers - new_center if np.any(max_rad > (np.sqrt(np.sum(diff * diff, axis=1)) - rads)): continue centers = np.concatenate([centers, new_center], axis=0) rads = np.concatenate([rads, np.array([[max_rad]])], axis=0) col = get_random_color(background_color) angle = random_state.rand() * 90 cv.ellipse(img, (x, y), (ax, ay), angle, 0, 360, col, -1) return np.empty((0, 2), dtype=np.int) def draw_star(img, nb_branches=6, min_len=32, min_label_len=64): """Draw a star and return the junction points + line map. Parameters: nb_branches: number of branches of the star """ num_branches = random_state.randint(3, nb_branches) min_dim = min(img.shape[0], img.shape[1]) # Convert length constrain to pixel if given float number if isinstance(min_len, float) and min_len <= 1.0: min_len = int(min_dim * min_len) if isinstance(min_label_len, float) and min_label_len <= 1.0: min_label_len = int(min_dim * min_label_len) thickness = random_state.randint(min_dim * 0.01, min_dim * 0.025) rad = max(random_state.rand() * min_dim / 2, min_dim / 5) x = random_state.randint(rad, img.shape[1] - rad) y = random_state.randint(rad, img.shape[0] - rad) # Sample num_branches points inside the circle slices = np.linspace(0, 2 * math.pi, num_branches + 1) angles = [ slices[i] + random_state.rand() * (slices[i + 1] - slices[i]) for i in range(num_branches) ] points = np.array( [ [ int(x + max(random_state.rand(), 0.3) * rad * math.cos(a)), int(y + max(random_state.rand(), 0.3) * rad * math.sin(a)), ] for a in angles ] ) points = np.concatenate(([[x, y]], points), axis=0) # Generate segments and check the length segments = np.array([[x, y, _[0], _[1]] for _ in points[1:, :]]) if check_segment_len(segments, min_len): return draw_star(img, nb_branches, min_len, min_label_len) # Only record the segments longer than min_label_len points1 = segments[:, :2] points2 = segments[:, 2:] seg_len = np.sqrt(np.sum((points1 - points2) ** 2, axis=1)) label_segments = segments[seg_len >= min_label_len, :] # Get all junctions from label segments junctions_all = np.concatenate( (label_segments[:, :2], label_segments[:, 2:]), axis=0 ) if junctions_all.shape[0] == 0: junc_points = None line_map = None # Get all unique junction points else: junc_points = np.unique(junctions_all, axis=0) # Generate line map from points and segments line_map = get_line_map(junc_points, label_segments) background_color = int(np.mean(img)) for i in range(1, num_branches + 1): col = get_random_color(background_color) cv.line( img, (points[0][0], points[0][1]), (points[i][0], points[i][1]), col, thickness, ) return {"points": junc_points, "line_map": line_map} def draw_checkerboard_multiseg( img, max_rows=7, max_cols=7, transform_params=(0.05, 0.15), min_label_len=64, seed=None, ): """Draw a checkerboard and output the junctions + line segments Parameters: max_rows: maximal number of rows + 1 max_cols: maximal number of cols + 1 transform_params: set the range of the parameters of the transformations """ if seed is None: global random_state else: random_state = np.random.RandomState(seed) background_color = int(np.mean(img)) min_dim = min(img.shape) if isinstance(min_label_len, float) and min_label_len <= 1.0: min_label_len = int(min_dim * min_label_len) # Create the grid rows = random_state.randint(3, max_rows) # number of rows cols = random_state.randint(3, max_cols) # number of cols s = min((img.shape[1] - 1) // cols, (img.shape[0] - 1) // rows) x_coord = np.tile(range(cols + 1), rows + 1).reshape(((rows + 1) * (cols + 1), 1)) y_coord = np.repeat(range(rows + 1), cols + 1).reshape(((rows + 1) * (cols + 1), 1)) # points are the grid coordinates points = s * np.concatenate([x_coord, y_coord], axis=1) # Warp the grid using an affine transformation and an homography alpha_affine = np.max(img.shape) * ( transform_params[0] + random_state.rand() * transform_params[1] ) center_square = np.float32(img.shape) // 2 min_dim = min(img.shape) square_size = min_dim // 3 pts1 = np.float32( [ center_square + square_size, [center_square[0] + square_size, center_square[1] - square_size], center_square - square_size, [center_square[0] - square_size, center_square[1] + square_size], ] ) pts2 = pts1 + random_state.uniform( -alpha_affine, alpha_affine, size=pts1.shape ).astype(np.float32) affine_transform = cv.getAffineTransform(pts1[:3], pts2[:3]) pts2 = pts1 + random_state.uniform( -alpha_affine / 2, alpha_affine / 2, size=pts1.shape ).astype(np.float32) perspective_transform = cv.getPerspectiveTransform(pts1, pts2) # Apply the affine transformation points = np.transpose( np.concatenate((points, np.ones(((rows + 1) * (cols + 1), 1))), axis=1) ) warped_points = np.transpose(np.dot(affine_transform, points)) # Apply the homography warped_col0 = np.add( np.sum(np.multiply(warped_points, perspective_transform[0, :2]), axis=1), perspective_transform[0, 2], ) warped_col1 = np.add( np.sum(np.multiply(warped_points, perspective_transform[1, :2]), axis=1), perspective_transform[1, 2], ) warped_col2 = np.add( np.sum(np.multiply(warped_points, perspective_transform[2, :2]), axis=1), perspective_transform[2, 2], ) warped_col0 = np.divide(warped_col0, warped_col2) warped_col1 = np.divide(warped_col1, warped_col2) warped_points = np.concatenate([warped_col0[:, None], warped_col1[:, None]], axis=1) warped_points_float = warped_points.copy() warped_points = warped_points.astype(int) # Fill the rectangles colors = np.zeros((rows * cols,), np.int32) for i in range(rows): for j in range(cols): # Get a color that contrast with the neighboring cells if i == 0 and j == 0: col = get_random_color(background_color) else: neighboring_colors = [] if i != 0: neighboring_colors.append(colors[(i - 1) * cols + j]) if j != 0: neighboring_colors.append(colors[i * cols + j - 1]) col = get_different_color(np.array(neighboring_colors)) colors[i * cols + j] = col # Fill the cell cv.fillConvexPoly( img, np.array( [ ( warped_points[i * (cols + 1) + j, 0], warped_points[i * (cols + 1) + j, 1], ), ( warped_points[i * (cols + 1) + j + 1, 0], warped_points[i * (cols + 1) + j + 1, 1], ), ( warped_points[(i + 1) * (cols + 1) + j + 1, 0], warped_points[(i + 1) * (cols + 1) + j + 1, 1], ), ( warped_points[(i + 1) * (cols + 1) + j, 0], warped_points[(i + 1) * (cols + 1) + j, 1], ), ] ), col, ) label_segments = np.empty([0, 4], dtype=np.int) # Iterate through rows for row_idx in range(rows + 1): # Include all the combination of the junctions # Iterate through all the combination of junction index in that row multi_seg_lst = [ np.array( [ warped_points_float[id1, 0], warped_points_float[id1, 1], warped_points_float[id2, 0], warped_points_float[id2, 1], ] )[None, ...] for (id1, id2) in combinations( range(row_idx * (cols + 1), (row_idx + 1) * (cols + 1), 1), 2 ) ] multi_seg = np.concatenate(multi_seg_lst, axis=0) label_segments = np.concatenate((label_segments, multi_seg), axis=0) # Iterate through columns for col_idx in range(cols + 1): # for 5 columns, we will have 5 + 1 edges # Include all the combination of the junctions # Iterate throuhg all the combination of junction index in that column multi_seg_lst = [ np.array( [ warped_points_float[id1, 0], warped_points_float[id1, 1], warped_points_float[id2, 0], warped_points_float[id2, 1], ] )[None, ...] for (id1, id2) in combinations( range(col_idx, col_idx + ((rows + 1) * (cols + 1)), cols + 1), 2 ) ] multi_seg = np.concatenate(multi_seg_lst, axis=0) label_segments = np.concatenate((label_segments, multi_seg), axis=0) label_segments_filtered = np.zeros([0, 4]) # Define image boundary polygon (in x y manner) image_poly = shapely.geometry.Polygon( [ [0, 0], [img.shape[1] - 1, 0], [img.shape[1] - 1, img.shape[0] - 1], [0, img.shape[0] - 1], ] ) for idx in range(label_segments.shape[0]): # Get the line segment seg_raw = label_segments[idx, :] seg = shapely.geometry.LineString([seg_raw[:2], seg_raw[2:]]) # The line segment is just inside the image. if seg.intersection(image_poly) == seg: label_segments_filtered = np.concatenate( (label_segments_filtered, seg_raw[None, ...]), axis=0 ) # Intersect with the image. elif seg.intersects(image_poly): # Check intersection try: p = np.array(seg.intersection(image_poly).coords).reshape([-1, 4]) # If intersect with eact one point except: continue segment = p label_segments_filtered = np.concatenate( (label_segments_filtered, segment), axis=0 ) else: continue label_segments = np.round(label_segments_filtered).astype(np.int) # Only record the segments longer than min_label_len points1 = label_segments[:, :2] points2 = label_segments[:, 2:] seg_len = np.sqrt(np.sum((points1 - points2) ** 2, axis=1)) label_segments = label_segments[seg_len >= min_label_len, :] # Get all junctions from label segments junc_points, line_map = get_unique_junctions(label_segments, min_label_len) # Draw lines on the boundaries of the board at random nb_rows = random_state.randint(2, rows + 2) nb_cols = random_state.randint(2, cols + 2) thickness = random_state.randint(min_dim * 0.01, min_dim * 0.015) for _ in range(nb_rows): row_idx = random_state.randint(rows + 1) col_idx1 = random_state.randint(cols + 1) col_idx2 = random_state.randint(cols + 1) col = get_random_color(background_color) cv.line( img, ( warped_points[row_idx * (cols + 1) + col_idx1, 0], warped_points[row_idx * (cols + 1) + col_idx1, 1], ), ( warped_points[row_idx * (cols + 1) + col_idx2, 0], warped_points[row_idx * (cols + 1) + col_idx2, 1], ), col, thickness, ) for _ in range(nb_cols): col_idx = random_state.randint(cols + 1) row_idx1 = random_state.randint(rows + 1) row_idx2 = random_state.randint(rows + 1) col = get_random_color(background_color) cv.line( img, ( warped_points[row_idx1 * (cols + 1) + col_idx, 0], warped_points[row_idx1 * (cols + 1) + col_idx, 1], ), ( warped_points[row_idx2 * (cols + 1) + col_idx, 0], warped_points[row_idx2 * (cols + 1) + col_idx, 1], ), col, thickness, ) # Keep only the points inside the image points = keep_points_inside(warped_points, img.shape[:2]) return {"points": junc_points, "line_map": line_map} def draw_stripes_multiseg( img, max_nb_cols=13, min_len=0.04, min_label_len=64, transform_params=(0.05, 0.15), seed=None, ): """Draw stripes in a distorted rectangle and output the junctions points + line map. Parameters: max_nb_cols: maximal number of stripes to be drawn min_width_ratio: the minimal width of a stripe is min_width_ratio * smallest dimension of the image transform_params: set the range of the parameters of the transformations """ # Set the optional random seed (most for debugging) if seed is None: global random_state else: random_state = np.random.RandomState(seed) background_color = int(np.mean(img)) # Create the grid board_size = ( int(img.shape[0] * (1 + random_state.rand())), int(img.shape[1] * (1 + random_state.rand())), ) # Number of cols col = random_state.randint(5, max_nb_cols) cols = np.concatenate( [board_size[1] * random_state.rand(col - 1), np.array([0, board_size[1] - 1])], axis=0, ) cols = np.unique(cols.astype(int)) # Remove the indices that are too close min_dim = min(img.shape) # Convert length constrain to pixel if given float number if isinstance(min_len, float) and min_len <= 1.0: min_len = int(min_dim * min_len) if isinstance(min_label_len, float) and min_label_len <= 1.0: min_label_len = int(min_dim * min_label_len) cols = cols[ (np.concatenate([cols[1:], np.array([board_size[1] + min_len])], axis=0) - cols) >= min_len ] # Update the number of cols col = cols.shape[0] - 1 cols = np.reshape(cols, (col + 1, 1)) cols1 = np.concatenate([cols, np.zeros((col + 1, 1), np.int32)], axis=1) cols2 = np.concatenate( [cols, (board_size[0] - 1) * np.ones((col + 1, 1), np.int32)], axis=1 ) points = np.concatenate([cols1, cols2], axis=0) # Warp the grid using an affine transformation and a homography alpha_affine = np.max(img.shape) * ( transform_params[0] + random_state.rand() * transform_params[1] ) center_square = np.float32(img.shape) // 2 square_size = min(img.shape) // 3 pts1 = np.float32( [ center_square + square_size, [center_square[0] + square_size, center_square[1] - square_size], center_square - square_size, [center_square[0] - square_size, center_square[1] + square_size], ] ) pts2 = pts1 + random_state.uniform( -alpha_affine, alpha_affine, size=pts1.shape ).astype(np.float32) affine_transform = cv.getAffineTransform(pts1[:3], pts2[:3]) pts2 = pts1 + random_state.uniform( -alpha_affine / 2, alpha_affine / 2, size=pts1.shape ).astype(np.float32) perspective_transform = cv.getPerspectiveTransform(pts1, pts2) # Apply the affine transformation points = np.transpose(np.concatenate((points, np.ones((2 * (col + 1), 1))), axis=1)) warped_points = np.transpose(np.dot(affine_transform, points)) # Apply the homography warped_col0 = np.add( np.sum(np.multiply(warped_points, perspective_transform[0, :2]), axis=1), perspective_transform[0, 2], ) warped_col1 = np.add( np.sum(np.multiply(warped_points, perspective_transform[1, :2]), axis=1), perspective_transform[1, 2], ) warped_col2 = np.add( np.sum(np.multiply(warped_points, perspective_transform[2, :2]), axis=1), perspective_transform[2, 2], ) warped_col0 = np.divide(warped_col0, warped_col2) warped_col1 = np.divide(warped_col1, warped_col2) warped_points = np.concatenate([warped_col0[:, None], warped_col1[:, None]], axis=1) warped_points_float = warped_points.copy() warped_points = warped_points.astype(int) # Fill the rectangles and get the segments color = get_random_color(background_color) # segments_debug = np.zeros([0, 4]) for i in range(col): # Fill the color color = (color + 128 + random_state.randint(-30, 30)) % 256 cv.fillConvexPoly( img, np.array( [ (warped_points[i, 0], warped_points[i, 1]), (warped_points[i + 1, 0], warped_points[i + 1, 1]), (warped_points[i + col + 2, 0], warped_points[i + col + 2, 1]), (warped_points[i + col + 1, 0], warped_points[i + col + 1, 1]), ] ), color, ) segments = np.zeros([0, 4]) row = 1 # in stripes case # Iterate through rows for row_idx in range(row + 1): # Include all the combination of the junctions # Iterate through all the combination of junction index in that row multi_seg_lst = [ np.array( [ warped_points_float[id1, 0], warped_points_float[id1, 1], warped_points_float[id2, 0], warped_points_float[id2, 1], ] )[None, ...] for (id1, id2) in combinations( range(row_idx * (col + 1), (row_idx + 1) * (col + 1), 1), 2 ) ] multi_seg = np.concatenate(multi_seg_lst, axis=0) segments = np.concatenate((segments, multi_seg), axis=0) # Iterate through columns for col_idx in range(col + 1): # for 5 columns, we will have 5 + 1 edges. # Include all the combination of the junctions # Iterate throuhg all the combination of junction index in that column multi_seg_lst = [ np.array( [ warped_points_float[id1, 0], warped_points_float[id1, 1], warped_points_float[id2, 0], warped_points_float[id2, 1], ] )[None, ...] for (id1, id2) in combinations( range(col_idx, col_idx + (row * col) + 2, col + 1), 2 ) ] multi_seg = np.concatenate(multi_seg_lst, axis=0) segments = np.concatenate((segments, multi_seg), axis=0) # Select and refine the segments segments_new = np.zeros([0, 4]) # Define image boundary polygon (in x y manner) image_poly = shapely.geometry.Polygon( [ [0, 0], [img.shape[1] - 1, 0], [img.shape[1] - 1, img.shape[0] - 1], [0, img.shape[0] - 1], ] ) for idx in range(segments.shape[0]): # Get the line segment seg_raw = segments[idx, :] seg = shapely.geometry.LineString([seg_raw[:2], seg_raw[2:]]) # The line segment is just inside the image. if seg.intersection(image_poly) == seg: segments_new = np.concatenate((segments_new, seg_raw[None, ...]), axis=0) # Intersect with the image. elif seg.intersects(image_poly): # Check intersection try: p = np.array(seg.intersection(image_poly).coords).reshape([-1, 4]) # If intersect at exact one point, just continue. except: continue segment = p segments_new = np.concatenate((segments_new, segment), axis=0) else: continue segments = (np.round(segments_new)).astype(np.int) # Only record the segments longer than min_label_len points1 = segments[:, :2] points2 = segments[:, 2:] seg_len = np.sqrt(np.sum((points1 - points2) ** 2, axis=1)) label_segments = segments[seg_len >= min_label_len, :] # Get all junctions from label segments junctions_all = np.concatenate( (label_segments[:, :2], label_segments[:, 2:]), axis=0 ) if junctions_all.shape[0] == 0: junc_points = None line_map = None # Get all unique junction points else: junc_points = np.unique(junctions_all, axis=0) # Generate line map from points and segments line_map = get_line_map(junc_points, label_segments) # Draw lines on the boundaries of the stripes at random nb_rows = random_state.randint(2, 5) nb_cols = random_state.randint(2, col + 2) thickness = random_state.randint(min_dim * 0.01, min_dim * 0.011) for _ in range(nb_rows): row_idx = random_state.choice([0, col + 1]) col_idx1 = random_state.randint(col + 1) col_idx2 = random_state.randint(col + 1) color = get_random_color(background_color) cv.line( img, ( warped_points[row_idx + col_idx1, 0], warped_points[row_idx + col_idx1, 1], ), ( warped_points[row_idx + col_idx2, 0], warped_points[row_idx + col_idx2, 1], ), color, thickness, ) for _ in range(nb_cols): col_idx = random_state.randint(col + 1) color = get_random_color(background_color) cv.line( img, (warped_points[col_idx, 0], warped_points[col_idx, 1]), (warped_points[col_idx + col + 1, 0], warped_points[col_idx + col + 1, 1]), color, thickness, ) # Keep only the points inside the image # points = keep_points_inside(warped_points, img.shape[:2]) return {"points": junc_points, "line_map": line_map} def draw_cube( img, min_size_ratio=0.2, min_label_len=64, scale_interval=(0.4, 0.6), trans_interval=(0.5, 0.2), ): """Draw a 2D projection of a cube and output the visible juntions. Parameters: min_size_ratio: min(img.shape) * min_size_ratio is the smallest achievable cube side size scale_interval: the scale is between scale_interval[0] and scale_interval[0]+scale_interval[1] trans_interval: the translation is between img.shape*trans_interval[0] and img.shape*(trans_interval[0] + trans_interval[1]) """ # Generate a cube and apply to it an affine transformation # The order matters! # The indices of two adjacent vertices differ only of one bit (Gray code) background_color = int(np.mean(img)) min_dim = min(img.shape[:2]) min_side = min_dim * min_size_ratio lx = min_side + random_state.rand() * 2 * min_dim / 3 # dims of the cube ly = min_side + random_state.rand() * 2 * min_dim / 3 lz = min_side + random_state.rand() * 2 * min_dim / 3 cube = np.array( [ [0, 0, 0], [lx, 0, 0], [0, ly, 0], [lx, ly, 0], [0, 0, lz], [lx, 0, lz], [0, ly, lz], [lx, ly, lz], ] ) rot_angles = random_state.rand(3) * 3 * math.pi / 10.0 + math.pi / 10.0 rotation_1 = np.array( [ [math.cos(rot_angles[0]), -math.sin(rot_angles[0]), 0], [math.sin(rot_angles[0]), math.cos(rot_angles[0]), 0], [0, 0, 1], ] ) rotation_2 = np.array( [ [1, 0, 0], [0, math.cos(rot_angles[1]), -math.sin(rot_angles[1])], [0, math.sin(rot_angles[1]), math.cos(rot_angles[1])], ] ) rotation_3 = np.array( [ [math.cos(rot_angles[2]), 0, -math.sin(rot_angles[2])], [0, 1, 0], [math.sin(rot_angles[2]), 0, math.cos(rot_angles[2])], ] ) scaling = np.array( [ [scale_interval[0] + random_state.rand() * scale_interval[1], 0, 0], [0, scale_interval[0] + random_state.rand() * scale_interval[1], 0], [0, 0, scale_interval[0] + random_state.rand() * scale_interval[1]], ] ) trans = np.array( [ img.shape[1] * trans_interval[0] + random_state.randint( -img.shape[1] * trans_interval[1], img.shape[1] * trans_interval[1] ), img.shape[0] * trans_interval[0] + random_state.randint( -img.shape[0] * trans_interval[1], img.shape[0] * trans_interval[1] ), 0, ] ) cube = trans + np.transpose( np.dot( scaling, np.dot( rotation_1, np.dot(rotation_2, np.dot(rotation_3, np.transpose(cube))) ), ) ) # The hidden corner is 0 by construction # The front one is 7 cube = cube[:, :2] # project on the plane z=0 cube = cube.astype(int) points = cube[1:, :] # get rid of the hidden corner # Get the three visible faces faces = np.array([[7, 3, 1, 5], [7, 5, 4, 6], [7, 6, 2, 3]]) # Get all visible line segments segments = np.zeros([0, 4]) # Iterate through all the faces for face_idx in range(faces.shape[0]): face = faces[face_idx, :] # Brute-forcely expand all the segments segment = np.array( [ np.concatenate((cube[face[0]], cube[face[1]]), axis=0), np.concatenate((cube[face[1]], cube[face[2]]), axis=0), np.concatenate((cube[face[2]], cube[face[3]]), axis=0), np.concatenate((cube[face[3]], cube[face[0]]), axis=0), ] ) segments = np.concatenate((segments, segment), axis=0) # Select and refine the segments segments_new = np.zeros([0, 4]) # Define image boundary polygon (in x y manner) image_poly = shapely.geometry.Polygon( [ [0, 0], [img.shape[1] - 1, 0], [img.shape[1] - 1, img.shape[0] - 1], [0, img.shape[0] - 1], ] ) for idx in range(segments.shape[0]): # Get the line segment seg_raw = segments[idx, :] seg = shapely.geometry.LineString([seg_raw[:2], seg_raw[2:]]) # The line segment is just inside the image. if seg.intersection(image_poly) == seg: segments_new = np.concatenate((segments_new, seg_raw[None, ...]), axis=0) # Intersect with the image. elif seg.intersects(image_poly): try: p = np.array(seg.intersection(image_poly).coords).reshape([-1, 4]) except: continue segment = p segments_new = np.concatenate((segments_new, segment), axis=0) else: continue segments = (np.round(segments_new)).astype(np.int) # Only record the segments longer than min_label_len points1 = segments[:, :2] points2 = segments[:, 2:] seg_len = np.sqrt(np.sum((points1 - points2) ** 2, axis=1)) label_segments = segments[seg_len >= min_label_len, :] # Get all junctions from label segments junctions_all = np.concatenate( (label_segments[:, :2], label_segments[:, 2:]), axis=0 ) if junctions_all.shape[0] == 0: junc_points = None line_map = None # Get all unique junction points else: junc_points = np.unique(junctions_all, axis=0) # Generate line map from points and segments line_map = get_line_map(junc_points, label_segments) # Fill the faces and draw the contours col_face = get_random_color(background_color) for i in [0, 1, 2]: cv.fillPoly(img, [cube[faces[i]].reshape((-1, 1, 2))], col_face) thickness = random_state.randint(min_dim * 0.003, min_dim * 0.015) for i in [0, 1, 2]: for j in [0, 1, 2, 3]: col_edge = ( col_face + 128 + random_state.randint(-64, 64) ) % 256 # color that constrats with the face color cv.line( img, (cube[faces[i][j], 0], cube[faces[i][j], 1]), (cube[faces[i][(j + 1) % 4], 0], cube[faces[i][(j + 1) % 4], 1]), col_edge, thickness, ) return {"points": junc_points, "line_map": line_map} def gaussian_noise(img): """Apply random noise to the image.""" cv.randu(img, 0, 255) return {"points": None, "line_map": None}