| """ |
| Some useful functions for dataset pre-processing |
| """ |
| import cv2 |
| import numpy as np |
| import shapely.geometry as sg |
|
|
| from ..synthetic_util import get_line_map |
| from . import homographic_transforms as homoaug |
|
|
|
|
| def random_scaling(image, junctions, line_map, scale=1., h_crop=0, w_crop=0): |
| H, W = image.shape[:2] |
| H_scale, W_scale = round(H * scale), round(W * scale) |
|
|
| |
| if H_scale == H and W_scale == W: |
| return (image, junctions, line_map, np.ones([H, W], dtype=np.int32)) |
|
|
| |
| if scale >= 1.: |
| image_big = cv2.resize(image, (W_scale, H_scale), |
| interpolation=cv2.INTER_LINEAR) |
| |
| image = image_big[h_crop:h_crop+H, w_crop:w_crop+W, ...] |
| valid_mask = np.ones([H, W], dtype=np.int32) |
|
|
| |
| junctions, line_map = process_junctions_and_line_map( |
| h_crop, w_crop, H, W, H_scale, W_scale, |
| junctions, line_map, "zoom-in") |
| |
| else: |
| image_shape_raw = image.shape |
| image_small = cv2.resize(image, (W_scale, H_scale), |
| interpolation=cv2.INTER_AREA) |
| |
| h_start = round((H - H_scale) / 2) |
| w_start = round((W - W_scale) / 2) |
| |
| image = np.zeros(image_shape_raw, dtype='float') |
| image[h_start:h_start+H_scale, |
| w_start:w_start+W_scale, ...] = image_small |
| valid_mask = np.zeros([H, W], dtype=np.int32) |
| valid_mask[h_start:h_start+H_scale, w_start:w_start+W_scale] = 1 |
|
|
| |
| junctions, line_map = process_junctions_and_line_map( |
| h_start, w_start, H, W, H_scale, W_scale, |
| junctions, line_map, "zoom-out") |
|
|
| return image, junctions, line_map, valid_mask |
|
|
|
|
| def process_junctions_and_line_map(h_start, w_start, H, W, H_scale, W_scale, |
| junctions, line_map, mode="zoom-in"): |
| if mode == "zoom-in": |
| junctions[:, 0] = junctions[:, 0] * H_scale / H |
| junctions[:, 1] = junctions[:, 1] * W_scale / W |
| line_segments = homoaug.convert_to_line_segments(junctions, line_map) |
| |
| line_segments_new = np.zeros([0, 4]) |
| image_poly = sg.Polygon( |
| [[w_start, h_start], |
| [w_start+W, h_start], |
| [w_start+W, h_start+H], |
| [w_start, h_start+H] |
| ]) |
| for idx in range(line_segments.shape[0]): |
| |
| seg_raw = line_segments[idx, :] |
| |
| seg = sg.LineString([np.flip(seg_raw[:2]), |
| np.flip(seg_raw[2:])]) |
| |
| if seg.intersection(image_poly) == seg: |
| line_segments_new = np.concatenate( |
| (line_segments_new, seg_raw[None, ...]), axis=0) |
| |
| elif seg.intersects(image_poly): |
| |
| try: |
| p = np.array( |
| seg.intersection(image_poly).coords).reshape([-1, 4]) |
| |
| except: |
| continue |
| segment = np.concatenate([np.flip(p[0, :2]), np.flip(p[0, 2:], |
| axis=0)])[None, ...] |
| line_segments_new = np.concatenate( |
| (line_segments_new, segment), axis=0) |
| else: |
| continue |
| line_segments_new = (np.round(line_segments_new)).astype(np.int32) |
| |
| segment_lens = np.linalg.norm( |
| line_segments_new[:, :2] - line_segments_new[:, 2:], axis=-1) |
| seg_mask = segment_lens != 0 |
| line_segments_new = line_segments_new[seg_mask, :] |
| |
| junctions_new = np.concatenate( |
| (line_segments_new[:, :2], line_segments_new[:, 2:]), axis=0) |
| if junctions_new.shape[0] == 0: |
| junctions_new = np.zeros([0, 2]) |
| line_map = np.zeros([0, 0]) |
| else: |
| junctions_new = np.unique(junctions_new, axis=0) |
| |
| line_map = get_line_map(junctions_new, |
| line_segments_new).astype(np.int32) |
| junctions_new[:, 0] -= h_start |
| junctions_new[:, 1] -= w_start |
| junctions = junctions_new |
| elif mode == "zoom-out": |
| |
| junctions[:, 0] = (junctions[:, 0] * H_scale / H) + h_start |
| junctions[:, 1] = (junctions[:, 1] * W_scale / W) + w_start |
| else: |
| raise ValueError("[Error] unknown mode...") |
|
|
| return junctions, line_map |
|
|
