Spaces:
Running
Running
File size: 5,006 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb |
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 123 124 125 126 127 128 |
"""
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
|