Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,073 Bytes
4c954ae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
"""
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)
# 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.int32))
# Zoom-in => resize and random crop
if scale >= 1.:
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.int32)
# 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='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
# 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.int32)
# 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.int32)
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
|