Spaces:
Running
Running
| """ | |
| 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.0, h_crop=0, w_crop=0): | |
| H, W = image.shape[:2] | |
| H_scale, W_scale = round(H * scale), round(W * scale) | |
| # Nothing to do if the scale is too close to 1 | |
| if H_scale == H and W_scale == W: | |
| return (image, junctions, line_map, np.ones([H, W], dtype=np.int)) | |
| # Zoom-in => resize and random crop | |
| if scale >= 1.0: | |
| image_big = cv2.resize( | |
| image, (W_scale, H_scale), interpolation=cv2.INTER_LINEAR | |
| ) | |
| # Crop the image | |
| image = image_big[h_crop : h_crop + H, w_crop : w_crop + W, ...] | |
| valid_mask = np.ones([H, W], dtype=np.int) | |
| # Process junctions | |
| junctions, line_map = process_junctions_and_line_map( | |
| h_crop, w_crop, H, W, H_scale, W_scale, junctions, line_map, "zoom-in" | |
| ) | |
| # Zoom-out => resize and pad | |
| else: | |
| image_shape_raw = image.shape | |
| image_small = cv2.resize( | |
| image, (W_scale, H_scale), interpolation=cv2.INTER_AREA | |
| ) | |
| # Decide the pasting location | |
| h_start = round((H - H_scale) / 2) | |
| w_start = round((W - W_scale) / 2) | |
| # Paste the image to the middle | |
| image = np.zeros(image_shape_raw, dtype=np.float) | |
| image[ | |
| h_start : h_start + H_scale, w_start : w_start + W_scale, ... | |
| ] = image_small | |
| valid_mask = np.zeros([H, W], dtype=np.int) | |
| valid_mask[h_start : h_start + H_scale, w_start : w_start + W_scale] = 1 | |
| # Process the junctions | |
| 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) | |
| # Crop segments to the new boundaries | |
| 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]): | |
| # Get the line segment | |
| seg_raw = line_segments[idx, :] # in HW format. | |
| # Convert to shapely line (flip to xy format) | |
| seg = sg.LineString([np.flip(seg_raw[:2]), np.flip(seg_raw[2:])]) | |
| # The line segment is just inside the image. | |
| if seg.intersection(image_poly) == seg: | |
| line_segments_new = np.concatenate( | |
| (line_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 = 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.int) | |
| # Filter segments with 0 length | |
| 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, :] | |
| # Convert back to junctions and line_map | |
| 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) | |
| # Generate line map from points and segments | |
| line_map = get_line_map(junctions_new, line_segments_new).astype(np.int) | |
| junctions_new[:, 0] -= h_start | |
| junctions_new[:, 1] -= w_start | |
| junctions = junctions_new | |
| elif mode == "zoom-out": | |
| # Process the junctions | |
| 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 | |