Vincentqyw
fix: roma
c74a070
raw
history blame
28.2 kB
"""
This file implements the synthetic shape dataset object for pytorch
"""
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
import os
import math
import h5py
import pickle
import torch
import numpy as np
import cv2
from tqdm import tqdm
from torchvision import transforms
from torch.utils.data import Dataset
import torch.utils.data.dataloader as torch_loader
from ..config.project_config import Config as cfg
from . import synthetic_util
from .transforms import photometric_transforms as photoaug
from .transforms import homographic_transforms as homoaug
from ..misc.train_utils import parse_h5_data
def synthetic_collate_fn(batch):
"""Customized collate_fn."""
batch_keys = ["image", "junction_map", "heatmap", "valid_mask", "homography"]
list_keys = ["junctions", "line_map", "file_key"]
outputs = {}
for data_key in batch[0].keys():
batch_match = sum([_ in data_key for _ in batch_keys])
list_match = sum([_ in data_key for _ in list_keys])
# print(batch_match, list_match)
if batch_match > 0 and list_match == 0:
outputs[data_key] = torch_loader.default_collate(
[b[data_key] for b in batch]
)
elif batch_match == 0 and list_match > 0:
outputs[data_key] = [b[data_key] for b in batch]
elif batch_match == 0 and list_match == 0:
continue
else:
raise ValueError(
"[Error] A key matches batch keys and list keys simultaneously."
)
return outputs
class SyntheticShapes(Dataset):
"""Dataset of synthetic shapes."""
# Initialize the dataset
def __init__(self, mode="train", config=None):
super(SyntheticShapes, self).__init__()
if not mode in ["train", "val", "test"]:
raise ValueError(
"[Error] Supported dataset modes are 'train', 'val', and 'test'."
)
self.mode = mode
# Get configuration
if config is None:
self.config = self.get_default_config()
else:
self.config = config
# Set all available primitives
self.available_primitives = [
"draw_lines",
"draw_polygon",
"draw_multiple_polygons",
"draw_star",
"draw_checkerboard_multiseg",
"draw_stripes_multiseg",
"draw_cube",
"gaussian_noise",
]
# Some cache setting
self.dataset_name = self.get_dataset_name()
self.cache_name = self.get_cache_name()
self.cache_path = cfg.synthetic_cache_path
# Check if export dataset exists
print("===============================================")
self.filename_dataset, self.datapoints = self.construct_dataset()
self.print_dataset_info()
# Initialize h5 file handle
self.dataset_path = os.path.join(
cfg.synthetic_dataroot, self.dataset_name + ".h5"
)
# Fix the random seed for torch and numpy in testing mode
if (self.mode == "val" or self.mode == "test") and self.config[
"add_augmentation_to_all_splits"
]:
seed = self.config.get("test_augmentation_seed", 200)
np.random.seed(seed)
torch.manual_seed(seed)
# For CuDNN
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
##########################################
## Dataset construction related methods ##
##########################################
def construct_dataset(self):
"""Dataset constructor."""
# Check if the filename cache exists
# If cache exists, load from cache
if self._check_dataset_cache():
print("[Info]: Found filename cache at ...")
print("\t Load filename cache...")
filename_dataset, datapoints = self.get_filename_dataset_from_cache()
print("\t Check if all file exists...")
# If all file exists, continue
if self._check_file_existence(filename_dataset):
print("\t All files exist!")
# If not, need to re-export the synthetic dataset
else:
print(
"\t Some files are missing. Re-export the synthetic shape dataset."
)
self.export_synthetic_shapes()
print("\t Initialize filename dataset")
filename_dataset, datapoints = self.get_filename_dataset()
print("\t Create filename dataset cache...")
self.create_filename_dataset_cache(filename_dataset, datapoints)
# If not, initialize dataset from scratch
else:
print("[Info]: Can't find filename cache ...")
print("\t First check export dataset exists.")
# If export dataset exists, then just update the filename_dataset
if self._check_export_dataset():
print("\t Synthetic dataset exists. Initialize the dataset ...")
# If export dataset does not exist, export from scratch
else:
print(
"\t Synthetic dataset does not exist. Export the synthetic dataset."
)
self.export_synthetic_shapes()
print("\t Initialize filename dataset")
filename_dataset, datapoints = self.get_filename_dataset()
print("\t Create filename dataset cache...")
self.create_filename_dataset_cache(filename_dataset, datapoints)
return filename_dataset, datapoints
def get_cache_name(self):
"""Get cache name from dataset config / default config."""
if self.config["dataset_name"] is None:
dataset_name = self.default_config["dataset_name"] + "_%s" % self.mode
else:
dataset_name = self.config["dataset_name"] + "_%s" % self.mode
# Compose cache name
cache_name = dataset_name + "_cache.pkl"
return cache_name
def get_dataset_name(self):
"""Get dataset name from dataset config / default config."""
if self.config["dataset_name"] is None:
dataset_name = self.default_config["dataset_name"] + "_%s" % self.mode
else:
dataset_name = self.config["dataset_name"] + "_%s" % self.mode
return dataset_name
def get_filename_dataset_from_cache(self):
"""Get filename dataset from cache."""
# Load from the pkl cache
cache_file_path = os.path.join(self.cache_path, self.cache_name)
with open(cache_file_path, "rb") as f:
data = pickle.load(f)
return data["filename_dataset"], data["datapoints"]
def get_filename_dataset(self):
"""Get filename dataset from scratch."""
# Path to the exported dataset
dataset_path = os.path.join(cfg.synthetic_dataroot, self.dataset_name + ".h5")
filename_dataset = {}
datapoints = []
# Open the h5 dataset
with h5py.File(dataset_path, "r") as f:
# Iterate through all the primitives
for prim_name in f.keys():
filenames = sorted(f[prim_name].keys())
filenames_full = [os.path.join(prim_name, _) for _ in filenames]
filename_dataset[prim_name] = filenames_full
datapoints += filenames_full
return filename_dataset, datapoints
def create_filename_dataset_cache(self, filename_dataset, datapoints):
"""Create filename dataset cache for faster initialization."""
# Check cache path exists
if not os.path.exists(self.cache_path):
os.makedirs(self.cache_path)
cache_file_path = os.path.join(self.cache_path, self.cache_name)
data = {"filename_dataset": filename_dataset, "datapoints": datapoints}
with open(cache_file_path, "wb") as f:
pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)
def export_synthetic_shapes(self):
"""Export synthetic shapes to disk."""
# Set the global random state for data generation
synthetic_util.set_random_state(
np.random.RandomState(self.config["generation"]["random_seed"])
)
# Define the export path
dataset_path = os.path.join(cfg.synthetic_dataroot, self.dataset_name + ".h5")
# Open h5py file
with h5py.File(dataset_path, "w", libver="latest") as f:
# Iterate through all types of shape
primitives = self.parse_drawing_primitives(self.config["primitives"])
split_size = self.config["generation"]["split_sizes"][self.mode]
for prim in primitives:
# Create h5 group
group = f.create_group(prim)
# Export single primitive
self.export_single_primitive(prim, split_size, group)
f.swmr_mode = True
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.0
).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")
def get_default_config(self):
"""Get default configuration of the dataset."""
# Initialize the default configuration
self.default_config = {
"dataset_name": "synthetic_shape",
"primitives": "all",
"add_augmentation_to_all_splits": False,
# Shape generation configuration
"generation": {
"split_sizes": {"train": 10000, "val": 400, "test": 500},
"random_seed": 10,
"image_size": [960, 1280],
"min_len": 0.09,
"min_label_len": 0.1,
"params": {
"generate_background": {
"min_kernel_size": 150,
"max_kernel_size": 500,
"min_rad_ratio": 0.02,
"max_rad_ratio": 0.031,
},
"draw_stripes": {"transform_params": (0.1, 0.1)},
"draw_multiple_polygons": {"kernel_boundaries": (50, 100)},
},
},
# Date preprocessing configuration.
"preprocessing": {"resize": [240, 320], "blur_size": 11},
"augmentation": {
"photometric": {
"enable": False,
"primitives": "all",
"params": {},
"random_order": True,
},
"homographic": {
"enable": False,
"params": {},
"valid_border_margin": 0,
},
},
}
return self.default_config
def parse_drawing_primitives(self, names):
"""Parse the primitives in config to list of primitive names."""
if names == "all":
p = self.available_primitives
else:
if isinstance(names, list):
p = names
else:
p = [names]
assert set(p) <= set(self.available_primitives)
return p
@staticmethod
def get_padded_filename(num_pad, idx):
"""Get the padded filename using adaptive padding."""
file_len = len("%d" % (idx))
filename = "0" * (num_pad - file_len) + "%d" % (idx)
return filename
def print_dataset_info(self):
"""Print dataset info."""
print("\t ---------Summary------------------")
print("\t Dataset mode: \t\t %s" % self.mode)
print("\t Number of primitive: \t %d" % len(self.filename_dataset.keys()))
print("\t Number of data: \t %d" % len(self.datapoints))
print("\t ----------------------------------")
#########################
## Pytorch related API ##
#########################
def get_data_from_datapoint(self, datapoint, reader=None):
"""Get data given the datapoint
(keyname of the h5 dataset e.g. "draw_lines/0000.h5")."""
# Check if the datapoint is valid
if not datapoint in self.datapoints:
raise ValueError(
"[Error] The specified datapoint is not in available datapoints."
)
# Get data from h5 dataset
if reader is None:
raise ValueError("[Error] The reader must be provided in __getitem__.")
else:
data = reader[datapoint]
return parse_h5_data(data)
def get_data_from_signature(self, primitive_name, index):
"""Get data given the primitive name and index ("draw_lines", 10)"""
# Check the primitive name and index
self._check_primitive_and_index(primitive_name, index)
# Get the datapoint from filename dataset
datapoint = self.filename_dataset[primitive_name][index]
return self.get_data_from_datapoint(datapoint)
def parse_transforms(self, names, all_transforms):
trans = (
all_transforms
if (names == "all")
else (names if isinstance(names, list) else [names])
)
assert set(trans) <= set(all_transforms)
return trans
def get_photo_transform(self):
"""Get list of photometric transforms (according to the config)."""
# Get the photometric transform config
photo_config = self.config["augmentation"]["photometric"]
if not photo_config["enable"]:
raise ValueError("[Error] Photometric augmentation is not enabled.")
# Parse photometric transforms
trans_lst = self.parse_transforms(
photo_config["primitives"], photoaug.available_augmentations
)
trans_config_lst = [photo_config["params"].get(p, {}) for p in trans_lst]
# List of photometric augmentation
photometric_trans_lst = [
getattr(photoaug, trans)(**conf)
for (trans, conf) in zip(trans_lst, trans_config_lst)
]
return photometric_trans_lst
def get_homo_transform(self):
"""Get homographic transforms (according to the config)."""
# Get homographic transforms for image
homo_config = self.config["augmentation"]["homographic"]["params"]
if not self.config["augmentation"]["homographic"]["enable"]:
raise ValueError("[Error] Homographic augmentation is not enabled")
# Parse the homographic transforms
# ToDo: use the shape from the config
image_shape = self.config["preprocessing"]["resize"]
# Compute the min_label_len from config
try:
min_label_tmp = self.config["generation"]["min_label_len"]
except:
min_label_tmp = None
# float label len => fraction
if isinstance(min_label_tmp, float): # Skip if not provided
min_label_len = min_label_tmp * min(image_shape)
# int label len => length in pixel
elif isinstance(min_label_tmp, int):
scale_ratio = (
self.config["preprocessing"]["resize"]
/ self.config["generation"]["image_size"][0]
)
min_label_len = self.config["generation"]["min_label_len"] * scale_ratio
# if none => no restriction
else:
min_label_len = 0
# Initialize the transform
homographic_trans = homoaug.homography_transform(
image_shape, homo_config, 0, min_label_len
)
return homographic_trans
@staticmethod
def junc_to_junc_map(junctions, image_size):
"""Convert junction points to junction maps."""
junctions = np.round(junctions).astype(np.int)
# Clip the boundary by image size
junctions[:, 0] = np.clip(junctions[:, 0], 0.0, image_size[0] - 1)
junctions[:, 1] = np.clip(junctions[:, 1], 0.0, image_size[1] - 1)
# Create junction map
junc_map = np.zeros([image_size[0], image_size[1]])
junc_map[junctions[:, 0], junctions[:, 1]] = 1
return junc_map[..., None].astype(np.int)
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.0).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["homo"]
# 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 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.0).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 __getitem__(self, index):
datapoint = self.datapoints[index]
# Initialize reader and use it
with h5py.File(self.dataset_path, "r", swmr=True) as reader:
data = self.get_data_from_datapoint(datapoint, reader)
# Apply different transforms in different mod.
if self.mode == "train" or self.config["add_augmentation_to_all_splits"]:
return_type = self.config.get("return_type", "single")
data = self.train_preprocessing(data)
else:
data = self.test_preprocessing(data)
return data
def __len__(self):
return len(self.datapoints)
########################
## Some other methods ##
########################
def _check_dataset_cache(self):
"""Check if dataset cache exists."""
cache_file_path = os.path.join(self.cache_path, self.cache_name)
if os.path.exists(cache_file_path):
return True
else:
return False
def _check_export_dataset(self):
"""Check if exported dataset exists."""
dataset_path = os.path.join(cfg.synthetic_dataroot, self.dataset_name)
if os.path.exists(dataset_path) and len(os.listdir(dataset_path)) > 0:
return True
else:
return False
def _check_file_existence(self, filename_dataset):
"""Check if all exported file exists."""
# Path to the exported dataset
dataset_path = os.path.join(cfg.synthetic_dataroot, self.dataset_name + ".h5")
flag = True
# Open the h5 dataset
with h5py.File(dataset_path, "r") as f:
# Iterate through all the primitives
for prim_name in f.keys():
if len(filename_dataset[prim_name]) != len(f[prim_name].keys()):
flag = False
return flag
def _check_primitive_and_index(self, primitive, index):
"""Check if the primitve and index are valid."""
# Check primitives
if not primitive in self.available_primitives:
raise ValueError("[Error] The primitive is not in available primitives.")
prim_len = len(self.filename_dataset[primitive])
# Check the index
if not index < prim_len:
raise ValueError(
"[Error] The index exceeds the total file counts %d for %s"
% (prim_len, primitive)
)