def Generate_transform_Dict(origin_width=256,
width=227,
ratio=0.16,
rot=0,
args=None):
std_value = 1.0 / 255.0
if (args is not None) and ("ResNet" in args.net):
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
cc = []
else:
normalize = transforms.Normalize(
mean=[104 / 255.0, 117 / 255.0, 128 / 255.0],
std=[1.0 / 255, 1.0 / 255, 1.0 / 255])
print("bgr init")
cc = [transforms.CovertBGR()]
transform_dict = {}
transform_dict['rand-crop'] = \
transforms.Compose(cc +
[transforms.Resize((origin_width)),
transforms.RandomResizedCrop(scale=(ratio, 1), size=width),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
transform_dict['center-crop'] = \
transforms.Compose(cc +
[
transforms.Resize((origin_width)),
transforms.CenterCrop(width),
transforms.ToTensor(),
normalize,
])
transform_dict['resize'] = \
transforms.Compose(cc +
[ transforms.Resize((width)),
transforms.ToTensor(),
normalize,
])
return transform_dict
def get_transform(train, dataset_name):
base_size = cfg.DATA_TRANSFORM.LOADSIZE
crop_size = cfg.DATA_TRANSFORM.CROPSIZE
ignore_label = cfg.DATASET.IGNORE_LABEL
if dataset_name == cfg.DATASET.SOURCE:
input_size = cfg.DATA_TRANSFORM.INPUT_SIZE_S
else:
input_size = cfg.DATA_TRANSFORM.INPUT_SIZE_T
min_size = int((1.0 if train else 1.0) * base_size)
max_size = int((1.3 if train else 1.0) * base_size)
transforms = []
if cfg.DATA_TRANSFORM.RANDOM_RESIZE_AND_CROP:
if train:
transforms.append(T.RandomResize(min_size, max_size))
transforms.append(T.RandomHorizontalFlip(0.5))
transforms.append(
T.RandomCrop(crop_size, ignore_label=ignore_label))
else:
transforms.append(T.Resize(cfg.DATA_TRANSFORM.INPUT_SIZE_T, True))
else:
if train:
transforms.append(T.Resize(input_size))
transforms.append(T.RandomHorizontalFlip(0.5))
else:
transforms.append(T.Resize(input_size, True))
mapping = get_label_map(cfg.DATASET.SOURCE, cfg.DATASET.TARGET)
transforms.append(T.LabelRemap(mapping[dataset_name]))
transforms.append(T.ToTensor(cfg.DATASET.IMG_MODE))
if cfg.DATASET.IMG_MODE == "BGR":
mean = (104.00698793, 116.66876762, 122.67891434)
std = (1.0, 1.0, 1.0)
else:
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)
transforms.append(T.Normalize(mean, std))
return T.Compose(transforms)
def build_transforms(config):
transform_train = T.Compose([
T.RandomCroping(config.DATA.HEIGHT,
config.DATA.WIDTH,
p=config.AUG.RC_PROB),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
T.RandomErasing(probability=config.AUG.RE_PROB)
])
transform_test = T.Compose([
T.Resize((config.DATA.HEIGHT, config.DATA.WIDTH)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
return transform_train, transform_test
def get_transform(dataset_name):
base_size = cfg.DATA_TRANSFORM.LOADSIZE
ignore_label = cfg.DATASET.IGNORE_LABEL
min_size = base_size
max_size = base_size
transforms = []
transforms.append(T.Resize(cfg.DATA_TRANSFORM.INPUT_SIZE_T, True))
mapping = get_label_map(cfg.DATASET.SOURCE, cfg.DATASET.TARGET)
transforms.append(T.LabelRemap(mapping[dataset_name]))
transforms.append(T.ToTensor(cfg.DATASET.IMG_MODE))
if cfg.DATASET.IMG_MODE == "BGR":
mean = (104.00698793, 116.66876762, 122.67891434)
std = (1.0, 1.0, 1.0)
else:
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)
transforms.append(T.Normalize(mean, std))
return T.Compose(transforms)
def __init__(
self,
# exp params
exp_name="u50_block",
# arch params
backbone="resnet50",
backbone_kwargs={},
dim_embedding=256,
feature_spatial_scale=0.25,
max_junctions=512,
junction_pooling_threshold=0.2,
junc_pooling_size=15,
attention_sigma=1.,
junction_heatmap_criterion="binary_cross_entropy",
block_inference_size=64,
adjacency_matrix_criterion="binary_cross_entropy",
# data params
data_root=r"/home/ziheng/indoorDist_new2",
img_size=416,
junc_sigma=3.,
batch_size=2,
# train params
gpus=[
0,
],
num_workers=5,
resume_epoch="latest",
is_train_junc=True,
is_train_adj=True,
# vis params
vis_junc_th=0.3,
vis_line_th=0.3):
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(c) for c in gpus)
self.is_cuda = bool(gpus)
self.model = LSDModule(
backbone=backbone,
dim_embedding=dim_embedding,
backbone_kwargs=backbone_kwargs,
junction_pooling_threshold=junction_pooling_threshold,
max_junctions=max_junctions,
feature_spatial_scale=feature_spatial_scale,
junction_heatmap_criterion=junction_heatmap_criterion,
junction_pooling_size=junc_pooling_size,
attention_sigma=attention_sigma,
block_inference_size=block_inference_size,
adjacency_matrix_criterion=adjacency_matrix_criterion,
weight_fn=weight_fn,
is_train_adj=is_train_adj,
is_train_junc=is_train_junc)
self.exp_name = exp_name
os.makedirs(os.path.join("log", exp_name), exist_ok=True)
os.makedirs(os.path.join("ckpt", exp_name), exist_ok=True)
self.writer = SummaryWriter(log_dir=os.path.join("log", exp_name))
# checkpoints
self.states = dict(last_epoch=-1, elapsed_time=0, state_dict=None)
if resume_epoch and os.path.isfile(
os.path.join("ckpt", exp_name,
f"train_states_{resume_epoch}.pth")):
states = torch.load(
os.path.join("ckpt", exp_name,
f"train_states_{resume_epoch}.pth"))
print(f"resume traning from epoch {states['last_epoch']}")
self.model.load_state_dict(states["state_dict"])
self.states.update(states)
self.train_data = SISTLine(data_root=data_root,
transforms=tf.Compose(
tf.Resize((img_size, img_size)),
tf.RandomHorizontalFlip(),
tf.RandomColorAug()),
phase="train",
sigma_junction=junc_sigma,
max_junctions=max_junctions)
self.train_loader = DataLoader(self.train_data,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True)
self.eval_data = SISTLine(data_root=data_root,
transforms=tf.Compose(
tf.Resize((img_size, img_size)), ),
phase="val",
sigma_junction=junc_sigma,
max_junctions=max_junctions)
self.eval_loader = DataLoader(self.eval_data,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=True)
self.vis_junc_th = vis_junc_th
self.vis_line_th = vis_line_th
self.block_size = block_inference_size
self.max_junctions = max_junctions
self.is_train_junc = is_train_junc
self.is_train_adj = is_train_adj
torch.multiprocessing.set_start_method('spawn')
# Define Saver
saver = Saver(args)
saver.save_experiment_config()
# Define Tensorboard Summary
summary = TensorboardSummary(saver.experiment_dir)
args.exp = saver.experiment_dir.split('_')[-1]
if args.train_dataset == 'cityscapes':
# Data
train_trans = transforms.Compose([
transforms.ToPILImage(),
# transforms.RandomResizedCrop((args.image_size, args.image_size), scale=(0.2, 2)),
transforms.Resize((args.image_size, args.image_size)),
transforms.RandomHorizontalFlip(),
transforms.RandomAffine(22, scale=(0.75, 1.25)),
transforms.ToTensor(),
transforms.Normalize(mean=[123.675, 116.28, 103.53],
std=[58.395, 57.12, 57.375])
# transforms.NormalizeInstance()
])
val_trans = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((args.image_size, args.image_size),
do_mask=False),
transforms.ToTensor(),
transforms.Normalize(mean=[123.675, 116.28, 103.53],
std=[58.395, 57.12, 57.375])
# transforms.NormalizeInstance()
args = parser.parse_args()
args.checkname = args.arc
# Define Saver
saver = Saver(args)
saver.save_experiment_config()
# Define Tensorboard Summary
summary = TensorboardSummary(saver.experiment_dir)
writer = summary.create_summary()
# Data
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_trans = transforms.Compose([transforms.Resize(321),
transforms.RandomCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
val_trans = transforms.Compose([transforms.Resize(321),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])
train_ds = VOCSBDClassification('/path/to/VOC',
'/path/to/SBD/benchmark_RELEASE/dataset',
transform=train_trans, image_set='train')
train_dl = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=4, drop_last=True)
val_ds = VOCClassification('/path/to/VOC', transform=val_trans, image_set='val')