def test_preprocessing(self, data):
""" Test preprocessing. """
# Fetch corresponding entries
image = data["image"]
points = data["points"]
line_map = data["line_map"]
heatmap = data["heatmap"]
image_size = image.shape[:2]
# Resize the image before the photometric and homographic transforms
if not (list(image.shape) == self.config["preprocessing"]["resize"]):
# Resize the image and the point location.
size_old = list(image.shape)
image = cv2.resize(
image,
tuple(self.config['preprocessing']['resize'][::-1]),
interpolation=cv2.INTER_LINEAR)
image = np.array(image, dtype=np.uint8)
points = (
points *
np.array(self.config['preprocessing']['resize'], np.float) /
np.array(size_old, np.float))
# Generate the line heatmap after post-processing
junctions = np.flip(np.round(points).astype(np.int32), axis=1)
heatmap = synthetic_util.get_line_heatmap(junctions,
line_map,
size=image.shape)
heatmap = (heatmap * 255.).astype(np.uint8)
# Update image size
image_size = image.shape[:2]
### image transform ###
image_transform = photoaug.normalize_image()
image = image_transform(image)
### joint transform ###
junction_map = self.junc_to_junc_map(points, image_size)
to_tensor = transforms.ToTensor()
image = to_tensor(image)
junctions = to_tensor(points)
junction_map = to_tensor(junction_map).to(torch.int)
line_map = to_tensor(line_map)
heatmap = to_tensor(heatmap)
valid_mask = to_tensor(np.ones(image_size)).to(torch.int32)
return {
"image": image,
"junctions": junctions,
"junction_map": junction_map,
"line_map": line_map,
"heatmap": heatmap,
"valid_mask": valid_mask
}
def __call__(self,
input_image,
junctions,
line_map,
valid_mask=None,
h**o=None,
scale=None):
# Sample one random homography or use the given one
if h**o is None or scale is None:
h**o, scale = sample_homography(self.image_size,
**self.homo_config)
# Warp the image
warped_image = cv2.warpPerspective(input_image,
h**o,
self.target_size,
flags=cv2.INTER_LINEAR)
valid_mask = compute_valid_mask(self.image_size, h**o,
self.border_margin, valid_mask)
# Convert junctions and line_map back to line segments
line_segments = convert_to_line_segments(junctions, line_map)
# Warp the segments and check the length.
# Adjust the min_label_length
warped_segments = warp_line_segment(line_segments, h**o,
self.image_size)
# Convert back to junctions and line_map
junctions_new = np.concatenate(
(warped_segments[:, :2], warped_segments[:, 2:]), axis=0)
if junctions_new.shape[0] == 0:
junctions_new = np.zeros([0, 2])
line_map = np.zeros([0, 0])
warped_heatmap = np.zeros(self.image_size)
else:
junctions_new = np.unique(junctions_new, axis=0)
# Generate line map from points and segments
line_map = get_line_map(junctions_new,
warped_segments).astype(np.int)
# Compute the heatmap
warped_heatmap = get_line_heatmap(np.flip(junctions_new, axis=1),
line_map, self.image_size)
return {
"junctions": junctions_new,
"warped_image": warped_image,
"valid_mask": valid_mask,
"line_map": line_map,
"warped_heatmap": warped_heatmap,
"h**o": h**o,
"scale": scale
}
def test_preprocessing_exported(self, data, numpy=False, scale=1.):
""" Test preprocessing for the exported labels. """
data = copy.deepcopy(data)
# Fetch the corresponding entries
image = data["image"]
junctions = data["junctions"]
line_map = data["line_map"]
image_size = image.shape[:2]
# Resize the image before photometric and homographical augmentations
if not (list(image_size) == self.config["preprocessing"]["resize"]):
# Resize the image and the point location.
size_old = list(image.shape)[:2] # Only H and W dimensions
image = cv2.resize(
image,
tuple(self.config['preprocessing']['resize'][::-1]),
interpolation=cv2.INTER_LINEAR)
image = np.array(image, dtype=np.uint8)
# In HW format
junctions = (
junctions *
np.array(self.config['preprocessing']['resize'], np.float) /
np.array(size_old, np.float))
# Optionally convert the image to grayscale
if self.config["gray_scale"]:
image = (color.rgb2gray(image) * 255.).astype(np.uint8)
# Still need to normalize image
image_transform = photoaug.normalize_image()
image = image_transform(image)
# Generate the line heatmap after post-processing
junctions_xy = np.flip(np.round(junctions).astype(np.int32), axis=1)
image_size = image.shape[:2]
heatmap = get_line_heatmap(junctions_xy, line_map, image_size)
# Declare default valid mask (all ones)
valid_mask = np.ones(image_size)
junction_map = self.junc_to_junc_map(junctions, image_size)
# Convert to tensor and return the results
to_tensor = transforms.ToTensor()
if not numpy:
outputs = {
"image": to_tensor(image),
"junctions": to_tensor(junctions).to(torch.float32)[0, ...],
"junction_map": to_tensor(junction_map).to(torch.int),
"line_map": to_tensor(line_map).to(torch.int32)[0, ...],
"heatmap": to_tensor(heatmap).to(torch.int32),
"valid_mask": to_tensor(valid_mask).to(torch.int32)
}
else:
outputs = {
"image": image,
"junctions": junctions.astype(np.float32),
"junction_map": junction_map.astype(np.int32),
"line_map": line_map.astype(np.int32),
"heatmap": heatmap.astype(np.int32),
"valid_mask": valid_mask.astype(np.int32)
}
return outputs
def test_preprocessing(self, data, numpy=False):
""" Test preprocessing for GT data. """
data = copy.deepcopy(data)
# Fetch the corresponding entries
image = data["image"]
junctions = data["junc"][:, :2]
line_pos = data["Lpos"]
line_neg = data["Lneg"]
image_size = image.shape[:2]
# Convert junctions to pixel coordinates (from 128x128)
junctions[:, 0] *= image_size[0] / 128
junctions[:, 1] *= image_size[1] / 128
# Resize the image before photometric and homographical augmentations
if not (list(image_size) == self.config["preprocessing"]["resize"]):
# Resize the image and the point location.
size_old = list(image.shape)[:2] # Only H and W dimensions
image = cv2.resize(
image,
tuple(self.config['preprocessing']['resize'][::-1]),
interpolation=cv2.INTER_LINEAR)
image = np.array(image, dtype=np.uint8)
# In HW format
junctions = (
junctions *
np.array(self.config['preprocessing']['resize'], np.float) /
np.array(size_old, np.float))
# Optionally convert the image to grayscale
if self.config["gray_scale"]:
image = (color.rgb2gray(image) * 255.).astype(np.uint8)
# Still need to normalize image
image_transform = photoaug.normalize_image()
image = image_transform(image)
# Convert to positive line map and negative line map (our format)
num_junctions = junctions.shape[0]
line_map_pos = self.convert_line_map(line_pos, num_junctions)
line_map_neg = self.convert_line_map(line_neg, num_junctions)
# Generate the line heatmap after post-processing
junctions_xy = np.flip(np.round(junctions).astype(np.int32), axis=1)
# Update image size
image_size = image.shape[:2]
heatmap_pos = get_line_heatmap(junctions_xy, line_map_pos, image_size)
heatmap_neg = get_line_heatmap(junctions_xy, line_map_neg, image_size)
# Declare default valid mask (all ones)
valid_mask = np.ones(image_size)
junction_map = self.junc_to_junc_map(junctions, image_size)
# Convert to tensor and return the results
to_tensor = transforms.ToTensor()
if not numpy:
return {
"image": to_tensor(image),
"junctions": to_tensor(junctions).to(torch.float32)[0, ...],
"junction_map": to_tensor(junction_map).to(torch.int),
"line_map_pos": to_tensor(line_map_pos).to(torch.int32)[0,
...],
"line_map_neg": to_tensor(line_map_neg).to(torch.int32)[0,
...],
"heatmap_pos": to_tensor(heatmap_pos).to(torch.int32),
"heatmap_neg": to_tensor(heatmap_neg).to(torch.int32),
"valid_mask": to_tensor(valid_mask).to(torch.int32)
}
else:
return {
"image": image,
"junctions": junctions.astype(np.float32),
"junction_map": junction_map.astype(np.int32),
"line_map_pos": line_map_pos.astype(np.int32),
"line_map_neg": line_map_neg.astype(np.int32),
"heatmap_pos": heatmap_pos.astype(np.int32),
"heatmap_neg": heatmap_neg.astype(np.int32),
"valid_mask": valid_mask.astype(np.int32)
}
def train_preprocessing_exported(self,
data,
numpy=False,
disable_homoaug=False,
desc_training=False,
H1=None,
H1_scale=None,
H2=None,
scale=1.,
h_crop=None,
w_crop=None):
""" Train preprocessing for the exported labels. """
data = copy.deepcopy(data)
# Fetch the corresponding entries
image = data["image"]
junctions = data["junctions"]
line_map = data["line_map"]
image_size = image.shape[:2]
# Define the random crop for scaling if necessary
if h_crop is None or w_crop is None:
h_crop, w_crop = 0, 0
if scale > 1:
H, W = self.config["preprocessing"]["resize"]
H_scale, W_scale = round(H * scale), round(W * scale)
if H_scale > H:
h_crop = np.random.randint(H_scale - H)
if W_scale > W:
w_crop = np.random.randint(W_scale - W)
# Resize the image before photometric and homographical augmentations
if not (list(image_size) == self.config["preprocessing"]["resize"]):
# Resize the image and the point location.
size_old = list(image.shape)[:2] # Only H and W dimensions
image = cv2.resize(
image,
tuple(self.config['preprocessing']['resize'][::-1]),
interpolation=cv2.INTER_LINEAR)
image = np.array(image, dtype=np.uint8)
# In HW format
junctions = (
junctions *
np.array(self.config['preprocessing']['resize'], np.float) /
np.array(size_old, np.float))
# Generate the line heatmap after post-processing
junctions_xy = np.flip(np.round(junctions).astype(np.int32), axis=1)
image_size = image.shape[:2]
heatmap = get_line_heatmap(junctions_xy, line_map, image_size)
# Optionally convert the image to grayscale
if self.config["gray_scale"]:
image = (color.rgb2gray(image) * 255.).astype(np.uint8)
# Check if we need to apply augmentations
# In training mode => yes.
# In homography adaptation mode (export mode) => No
if self.config["augmentation"]["photometric"]["enable"]:
photo_trans_lst = self.get_photo_transform()
### Image transform ###
np.random.shuffle(photo_trans_lst)
image_transform = transforms.Compose(photo_trans_lst +
[photoaug.normalize_image()])
else:
image_transform = photoaug.normalize_image()
image = image_transform(image)
# Perform the random scaling
if scale != 1.:
image, junctions, line_map, valid_mask = random_scaling(
image,
junctions,
line_map,
scale,
h_crop=h_crop,
w_crop=w_crop)
else:
# Declare default valid mask (all ones)
valid_mask = np.ones(image_size)
# Initialize the empty output dict
outputs = {}
# Convert to tensor and return the results
to_tensor = transforms.ToTensor()
# Check homographic augmentation
warp = (self.config["augmentation"]["homographic"]["enable"]
and disable_homoaug == False)
if warp:
homo_trans = self.get_homo_transform()
# Perform homographic transform
if H1 is None:
homo_outputs = homo_trans(image,
junctions,
line_map,
valid_mask=valid_mask)
else:
homo_outputs = homo_trans(image,
junctions,
line_map,
h**o=H1,
scale=H1_scale,
valid_mask=valid_mask)
homography_mat = homo_outputs["h**o"]
# Give the warp of the other view
if H1 is None:
H1 = homo_outputs["h**o"]
# Sample points along each line segments for the descriptor
if desc_training:
line_points, line_indices = self.get_line_points(
junctions,
line_map,
H1=H1,
H2=H2,
img_size=image_size,
warp=warp)
# Record the warped results
if warp:
junctions = homo_outputs["junctions"] # Should be HW format
image = homo_outputs["warped_image"]
line_map = homo_outputs["line_map"]
valid_mask = homo_outputs["valid_mask"] # Same for pos and neg
heatmap = homo_outputs["warped_heatmap"]
# Optionally put warping information first.
if not numpy:
outputs["homography_mat"] = to_tensor(homography_mat).to(
torch.float32)[0, ...]
else:
outputs["homography_mat"] = homography_mat.astype(np.float32)
junction_map = self.junc_to_junc_map(junctions, image_size)
if not numpy:
outputs.update({
"image":
to_tensor(image).to(torch.float32),
"junctions":
to_tensor(junctions).to(torch.float32)[0, ...],
"junction_map":
to_tensor(junction_map).to(torch.int),
"line_map":
to_tensor(line_map).to(torch.int32)[0, ...],
"heatmap":
to_tensor(heatmap).to(torch.int32),
"valid_mask":
to_tensor(valid_mask).to(torch.int32)
})
if desc_training:
outputs.update({
"line_points":
to_tensor(line_points).to(torch.float32)[0],
"line_indices":
torch.tensor(line_indices, dtype=torch.int)
})
else:
outputs.update({
"image": image,
"junctions": junctions.astype(np.float32),
"junction_map": junction_map.astype(np.int32),
"line_map": line_map.astype(np.int32),
"heatmap": heatmap.astype(np.int32),
"valid_mask": valid_mask.astype(np.int32)
})
if desc_training:
outputs.update({
"line_points": line_points.astype(np.float32),
"line_indices": line_indices.astype(int)
})
return outputs
def train_preprocessing(self, data, numpy=False):
""" Train preprocessing for GT data. """
# Fetch the corresponding entries
image = data["image"]
junctions = data["junc"][:, :2]
line_pos = data["Lpos"]
line_neg = data["Lneg"]
image_size = image.shape[:2]
# Convert junctions to pixel coordinates (from 128x128)
junctions[:, 0] *= image_size[0] / 128
junctions[:, 1] *= image_size[1] / 128
# Resize the image before photometric and homographical augmentations
if not (list(image_size) == self.config["preprocessing"]["resize"]):
# Resize the image and the point location.
size_old = list(image.shape)[:2] # Only H and W dimensions
image = cv2.resize(
image,
tuple(self.config['preprocessing']['resize'][::-1]),
interpolation=cv2.INTER_LINEAR)
image = np.array(image, dtype=np.uint8)
# In HW format
junctions = (
junctions *
np.array(self.config['preprocessing']['resize'], np.float) /
np.array(size_old, np.float))
# Convert to positive line map and negative line map (our format)
num_junctions = junctions.shape[0]
line_map_pos = self.convert_line_map(line_pos, num_junctions)
line_map_neg = self.convert_line_map(line_neg, num_junctions)
# Generate the line heatmap after post-processing
junctions_xy = np.flip(np.round(junctions).astype(np.int32), axis=1)
# Update image size
image_size = image.shape[:2]
heatmap_pos = get_line_heatmap(junctions_xy, line_map_pos, image_size)
heatmap_neg = get_line_heatmap(junctions_xy, line_map_neg, image_size)
# Declare default valid mask (all ones)
valid_mask = np.ones(image_size)
# Optionally convert the image to grayscale
if self.config["gray_scale"]:
image = (color.rgb2gray(image) * 255.).astype(np.uint8)
# Check if we need to apply augmentations
# In training mode => yes.
# In homography adaptation mode (export mode) => No
if self.config["augmentation"]["photometric"]["enable"]:
photo_trans_lst = self.get_photo_transform()
### Image transform ###
np.random.shuffle(photo_trans_lst)
image_transform = transforms.Compose(photo_trans_lst +
[photoaug.normalize_image()])
else:
image_transform = photoaug.normalize_image()
image = image_transform(image)
# Check homographic augmentation
if self.config["augmentation"]["homographic"]["enable"]:
homo_trans = self.get_homo_transform()
# Perform homographic transform
outputs_pos = homo_trans(image, junctions, line_map_pos)
outputs_neg = homo_trans(image, junctions, line_map_neg)
# record the warped results
junctions = outputs_pos["junctions"] # Should be HW format
image = outputs_pos["warped_image"]
line_map_pos = outputs_pos["line_map"]
line_map_neg = outputs_neg["line_map"]
heatmap_pos = outputs_pos["warped_heatmap"]
heatmap_neg = outputs_neg["warped_heatmap"]
valid_mask = outputs_pos["valid_mask"] # Same for pos and neg
junction_map = self.junc_to_junc_map(junctions, image_size)
# Convert to tensor and return the results
to_tensor = transforms.ToTensor()
if not numpy:
return {
"image": to_tensor(image),
"junctions": to_tensor(junctions).to(torch.float32)[0, ...],
"junction_map": to_tensor(junction_map).to(torch.int),
"line_map_pos": to_tensor(line_map_pos).to(torch.int32)[0,
...],
"line_map_neg": to_tensor(line_map_neg).to(torch.int32)[0,
...],
"heatmap_pos": to_tensor(heatmap_pos).to(torch.int32),
"heatmap_neg": to_tensor(heatmap_neg).to(torch.int32),
"valid_mask": to_tensor(valid_mask).to(torch.int32)
}
else:
return {
"image": image,
"junctions": junctions.astype(np.float32),
"junction_map": junction_map.astype(np.int32),
"line_map_pos": line_map_pos.astype(np.int32),
"line_map_neg": line_map_neg.astype(np.int32),
"heatmap_pos": heatmap_pos.astype(np.int32),
"heatmap_neg": heatmap_neg.astype(np.int32),
"valid_mask": valid_mask.astype(np.int32)
}
def train_preprocessing(self, data, disable_homoaug=False):
""" Training preprocessing. """
# Fetch corresponding entries
image = data["image"]
junctions = data["points"]
line_map = data["line_map"]
heatmap = data["heatmap"]
image_size = image.shape[:2]
# Resize the image before the photometric and homographic transforms
# Check if we need to do the resizing
if not (list(image.shape) == self.config["preprocessing"]["resize"]):
# Resize the image and the point location.
size_old = list(image.shape)
image = cv2.resize(
image,
tuple(self.config['preprocessing']['resize'][::-1]),
interpolation=cv2.INTER_LINEAR)
image = np.array(image, dtype=np.uint8)
junctions = (
junctions *
np.array(self.config['preprocessing']['resize'], np.float) /
np.array(size_old, np.float))
# Generate the line heatmap after post-processing
junctions_xy = np.flip(np.round(junctions).astype(np.int32),
axis=1)
heatmap = synthetic_util.get_line_heatmap(junctions_xy,
line_map,
size=image.shape)
heatmap = (heatmap * 255.).astype(np.uint8)
# Update image size
image_size = image.shape[:2]
# Declare default valid mask (all ones)
valid_mask = np.ones(image_size)
# Check if we need to apply augmentations
# In training mode => yes.
# In homography adaptation mode (export mode) => No
# Check photometric augmentation
if self.config["augmentation"]["photometric"]["enable"]:
photo_trans_lst = self.get_photo_transform()
### Image transform ###
np.random.shuffle(photo_trans_lst)
image_transform = transforms.Compose(photo_trans_lst +
[photoaug.normalize_image()])
else:
image_transform = photoaug.normalize_image()
image = image_transform(image)
# Initialize the empty output dict
outputs = {}
# Convert to tensor and return the results
to_tensor = transforms.ToTensor()
# Check homographic augmentation
if (self.config["augmentation"]["homographic"]["enable"]
and disable_homoaug == False):
homo_trans = self.get_homo_transform()
# Perform homographic transform
homo_outputs = homo_trans(image, junctions, line_map)
# Record the warped results
junctions = homo_outputs["junctions"] # Should be HW format
image = homo_outputs["warped_image"]
line_map = homo_outputs["line_map"]
heatmap = homo_outputs["warped_heatmap"]
valid_mask = homo_outputs["valid_mask"] # Same for pos and neg
homography_mat = homo_outputs["h**o"]
# Optionally put warpping information first.
outputs["homography_mat"] = to_tensor(homography_mat).to(
torch.float32)[0, ...]
junction_map = self.junc_to_junc_map(junctions, image_size)
outputs.update({
"image":
to_tensor(image),
"junctions":
to_tensor(np.ascontiguousarray(junctions).copy()).to(
torch.float32)[0, ...],
"junction_map":
to_tensor(junction_map).to(torch.int),
"line_map":
to_tensor(line_map).to(torch.int32)[0, ...],
"heatmap":
to_tensor(heatmap).to(torch.int32),
"valid_mask":
to_tensor(valid_mask).to(torch.int32),
})
return outputs
def export_single_primitive(self, primitive, split_size, group):
""" Export single primitive. """
# Check if the primitive is valid or not
if primitive not in self.available_primitives:
raise ValueError("[Error]: %s is not a supported primitive" %
primitive)
# Set the random seed
synthetic_util.set_random_state(
np.random.RandomState(self.config["generation"]["random_seed"]))
# Generate shapes
print("\t Generating %s ..." % primitive)
for idx in tqdm(range(split_size), ascii=True):
# Generate background image
image = synthetic_util.generate_background(
self.config['generation']['image_size'],
**self.config['generation']['params']['generate_background'])
# Generate points
drawing_func = getattr(synthetic_util, primitive)
kwarg = self.config["generation"]["params"].get(primitive, {})
# Get min_len and min_label_len
min_len = self.config["generation"]["min_len"]
min_label_len = self.config["generation"]["min_label_len"]
# Some only take min_label_len, and gaussian noises take nothing
if primitive in [
"draw_lines", "draw_polygon", "draw_multiple_polygons",
"draw_star"
]:
data = drawing_func(image,
min_len=min_len,
min_label_len=min_label_len,
**kwarg)
elif primitive in [
"draw_checkerboard_multiseg", "draw_stripes_multiseg",
"draw_cube"
]:
data = drawing_func(image,
min_label_len=min_label_len,
**kwarg)
else:
data = drawing_func(image, **kwarg)
# Convert the data
if data["points"] is not None:
points = np.flip(data["points"], axis=1).astype(np.float)
line_map = data["line_map"].astype(np.int32)
else:
points = np.zeros([0, 2]).astype(np.float)
line_map = np.zeros([0, 0]).astype(np.int32)
# Post-processing
blur_size = self.config["preprocessing"]["blur_size"]
image = cv2.GaussianBlur(image, (blur_size, blur_size), 0)
# Resize the image and the point location.
points = (
points *
np.array(self.config['preprocessing']['resize'], np.float) /
np.array(self.config['generation']['image_size'], np.float))
image = cv2.resize(
image,
tuple(self.config['preprocessing']['resize'][::-1]),
interpolation=cv2.INTER_LINEAR)
image = np.array(image, dtype=np.uint8)
# Generate the line heatmap after post-processing
junctions = np.flip(np.round(points).astype(np.int32), axis=1)
heatmap = (synthetic_util.get_line_heatmap(
junctions, line_map, size=image.shape) * 255.).astype(np.uint8)
# Record the data in group
num_pad = math.ceil(math.log10(split_size)) + 1
file_key_name = self.get_padded_filename(num_pad, idx)
file_group = group.create_group(file_key_name)
# Store data
file_group.create_dataset("points",
data=points,
compression="gzip")
file_group.create_dataset("image", data=image, compression="gzip")
file_group.create_dataset("line_map",
data=line_map,
compression="gzip")
file_group.create_dataset("heatmap",
data=heatmap,
compression="gzip")
