def generate_results_file(net,
target_txt,
classes_names,
out,
bs,
dim,
is_letterbox=False):
numclass = len(classes_names)
if is_letterbox:
transform = Compose(
[IaaAugmentations([IaaLetterbox(dim)]),
ToTensor()])
else:
transform = Compose([IaaAugmentations([iaa.Scale(dim)]), ToTensor()])
ds = COCOEvalDataset(target_txt, dim, transform)
dl = DataLoader(ds,
batch_size=bs,
num_workers=4,
collate_fn=variable_shape_collate_fn)
with open_json_pred_writer(out, classes_names,
is_letterbox) as pred_writer:
predict_and_process(dl,
net,
num_classes=numclass,
batch_handler=pred_writer)
def __init__(self, content=None):
get = content.get_resource
self.__model = onnxruntime.InferenceSession(
get(content.model_path))
self.line_transform = Compose([
Resize((512, 512)),
ToTensor(),
Normalize([0.5], [0.5]),
Lambda(lambda img: np.expand_dims(img, 0))
])
self.hint_transform = Compose([
# input must RGBA !
Resize((128, 128), Image.NEAREST),
Lambda(lambda img: img.convert(mode='RGB')),
ToTensor(),
Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
Lambda(lambda img: np.expand_dims(img, 0))
])
self.line_draft_transform = Compose([
Resize((128, 128)),
ToTensor(),
Normalize([0.5], [0.5]),
Lambda(lambda img: np.expand_dims(img, 0))
])
self.alpha_transform = Compose([
Lambda(lambda img: self.get_alpha(img)),
])
def main():
'''
Runs if the module is called as a script (eg: python3 dataset.py <dataset_name> <frametime>)
Executes self tests
'''
dataset_name = argv[1] if len(argv) > 1 else "clarissa"
wait = argv[2] if len(argv) > 2 else 10
print(
"Dataset module running as script, executing dataset unit test in {}".
format(dataset_name))
if dataset_name == "adni_slices":
unit_test(image_dataset=False,
adni=True,
hiponly=True,
plt_show=True,
nworkers=4,
e2d=True)
elif dataset_name == "clarissa_slices":
unit_test(image_dataset=False,
adni=False,
hiponly=True,
plt_show=True,
nworkers=4,
e2d=True)
elif dataset_name == "concat":
from transforms import ReturnPatch, Intensity, RandomFlip, Noisify, ToTensor, CenterCrop, RandomAffine
train_transforms = Compose([
ReturnPatch(patch_size=(32, 32)),
RandomAffine(),
Intensity(),
RandomFlip(modes=['horflip']),
Noisify(),
ToTensor()
]) #default is 32 32 patch
data_transforms = {
'train': train_transforms,
'validation': train_transforms,
'test': Compose([CenterCrop(160, 160),
ToTensor()])
}
mode = "train"
data, dsizes = get_data(data_transforms=data_transforms,
db="concat",
e2d=True,
batch_size=50 + 150 * (mode != "test"))
print("Dataset sizes: {}".format(dsizes))
for o in orientations:
batch = next(iter(data[o][mode]))
display_batch(batch, o + " concat " + mode + " data")
plt.show()
else:
view_volumes(dataset_name, wait=10)
print("All tests completed!")
def main():
# get test set and test set loader
test_set = CRC_Dataset(root_dir=os.path.join(os.getcwd(), "data\\test"),
transforms=[
MirrorPad(((6, ), (6, ), (0, ))),
ToTensor(),
Normalize(means=(0.7942, 0.6693, 0.7722),
stds=(0.1998, 0.3008, 0.2037))
])
test_loader = torch.utils.data.DataLoader(
test_set,
batch_size=args.batch_size if args.batch_size else 8,
num_workers=args.workers if args.workers else 1,
pin_memory=use_cuda,
)
model = UNet((512, 512), (500, 500),
32,
64,
128,
256,
512,
droprate=0.5,
Norm=torch.nn.BatchNorm2d)
#model = UNet((512, 512), (500, 500), 32, 64, 128, 256, Norm=torch.nn.BatchNorm2d)
model.load_state_dict(torch.load(chkpt))
model.to(device)
dice, acc = compute_metrics_on_test_set(model, test_loader)
print(dice, acc)
def valid_dataset(root_dir,
normalization=None,
grayscale=False,
square=False,
csv_file='B/validation_data.csv',
scale=1.0):
"""This function loads the training dataset with the desired transformations."""
scaled_width = int(round(scale * constants.default_width))
transformations = []
if scale != 1.0:
transformations.append(Resize(scaled_width))
if square:
transformations.append(Square())
if grayscale:
transformations.append(BlackAndWhite())
transformations.append(ToTensor())
if normalization is not None:
transformations.append(normalization)
transform = transforms.Compose(transformations)
valid_dataset = PostureLandmarksDataset(csv_file=csv_file,
root_dir=root_dir,
transform=transform)
return valid_dataset
def get_transforms():
data_transform = torchvision.transforms.Compose([
ToTensor(),
Normalize(mean=constants.DATA_MEAN, std=constants.DATA_STD),
Resize(constants.TRANSFORMED_IMAGE_SIZE)
])
return data_transform
def main():
class imshowCollate(object):
def __init__(self):
pass
def __call__(self, batch):
images, labels = zip(*batch)
idx = 0
for img in images:
img = img.cpu().numpy().transpose((1, 2, 0)) * 255 #totensor
cv2.imwrite(
'datatest/sev_img/img' + str(idx) + '——' +
str(labels[idx]) + '.jpg', img)
# print(img.shape)
idx += 1
return images, labels
from transforms import Compose, Normalize, RandomResizedCrop, RandomHorizontalFlip, \
ColorJitter, ToTensor,Lighting
batch_size = 16
normalize = Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
dataset = FileListLabeledDataset(
'/workspace/mnt/group/algo/yangdecheng/work/multi_task/pytorch-train/datatest/test.txt',
'/workspace/mnt/group/algo/yangdecheng/work/multi_task/pytorch-train/datatest/pic',
Compose([
RandomResizedCrop((112),
scale=(0.7, 1.2),
ratio=(1. / 1., 4. / 1.)),
RandomHorizontalFlip(),
ColorJitter(brightness=[0.5, 1.5],
contrast=[0.5, 1.5],
saturation=[0.5, 1.5],
hue=0),
ToTensor(),
Lighting(1, [0.2175, 0.0188, 0.0045],
[[-0.5675, 0.7192, 0.4009], [-0.5808, -0.0045, -0.8140],
[-0.5836, -0.6948, 0.4203]]), #0.1
# normalize,
]))
train_loader = torch.utils.data.DataLoader(dataset,
batch_size=100,
shuffle=True,
num_workers=10,
pin_memory=True,
sampler=None,
collate_fn=imshowCollate())
from multiprocessing import Process
p_list = []
for i in range(1):
p_list.append(Process(target=iter_f, args=(train_loader, )))
for p in p_list:
p.start()
for p in p_list:
p.join()
def init_seg(input_sizes,
std,
mean,
dataset,
test_base=None,
test_label_id_map=None,
city_aug=0):
if dataset == 'voc':
transform_test = Compose([
ToTensor(),
ZeroPad(size=input_sizes),
Normalize(mean=mean, std=std)
])
elif dataset == 'city' or dataset == 'gtav' or dataset == 'synthia': # All the same size
if city_aug == 2: # ERFNet and ENet
transform_test = Compose([
ToTensor(),
Resize(size_image=input_sizes, size_label=input_sizes),
LabelMap(test_label_id_map)
])
elif city_aug == 1: # City big
transform_test = Compose([
ToTensor(),
Resize(size_image=input_sizes, size_label=input_sizes),
Normalize(mean=mean, std=std),
LabelMap(test_label_id_map)
])
else:
raise ValueError
# Not the actual test set (i.e. validation set)
test_set = StandardSegmentationDataset(
root=test_base,
image_set='val',
transforms=transform_test,
data_set='city'
if dataset == 'gtav' or dataset == 'synthia' else dataset)
val_loader = torch.utils.data.DataLoader(dataset=test_set,
batch_size=1,
num_workers=0,
shuffle=False)
# Testing
return val_loader
def transform_preset(mode='random_crop'):
transformer = None
if mode == 'random_crop':
transformer = Compose([
ToWoldCoordinateSystem(),
ToTensor()])
return transformer
def load_data(datadir):
# Data loading code
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
base_size = 320
crop_size = 256
min_size = int(0.5 * base_size)
max_size = int(2.0 * base_size)
print("Loading training data")
st = time.time()
dataset = VOCSegmentation(datadir,
image_set='train',
download=True,
transforms=Compose([
RandomResize(min_size, max_size),
RandomCrop(crop_size),
RandomHorizontalFlip(0.5),
SampleTransform(
transforms.ColorJitter(brightness=0.3,
contrast=0.3,
saturation=0.1,
hue=0.02)),
ToTensor(),
SampleTransform(normalize)
]))
print("Took", time.time() - st)
print("Loading validation data")
st = time.time()
dataset_test = VOCSegmentation(datadir,
image_set='val',
download=True,
transforms=Compose([
RandomResize(base_size, base_size),
ToTensor(),
SampleTransform(normalize)
]))
print("Took", time.time() - st)
return dataset, dataset_test
def init(batch_size, state, input_sizes, dataset, mean, std, base, workers=10):
# Return data_loaders
# depending on whether the state is
# 0: training
# 1: fast validation by mean IoU (validation set)
# 2: just testing (test set)
# 3: just testing (validation set)
# Transformations
# ! Can't use torchvision.Transforms.Compose
transforms_test = Compose(
[Resize(size_image=input_sizes[0], size_label=input_sizes[0]),
ToTensor(),
Normalize(mean=mean, std=std)])
transforms_train = Compose(
[Resize(size_image=input_sizes[0], size_label=input_sizes[0]),
RandomRotation(degrees=3),
ToTensor(),
Normalize(mean=mean, std=std)])
if state == 0:
data_set = StandardLaneDetectionDataset(root=base, image_set='train', transforms=transforms_train,
data_set=dataset)
data_loader = torch.utils.data.DataLoader(dataset=data_set, batch_size=batch_size,
num_workers=workers, shuffle=True)
validation_set = StandardLaneDetectionDataset(root=base, image_set='val',
transforms=transforms_test, data_set=dataset)
validation_loader = torch.utils.data.DataLoader(dataset=validation_set, batch_size=batch_size * 4,
num_workers=workers, shuffle=False)
return data_loader, validation_loader
elif state == 1 or state == 2 or state == 3:
image_sets = ['valfast', 'test', 'val']
data_set = StandardLaneDetectionDataset(root=base, image_set=image_sets[state - 1],
transforms=transforms_test, data_set=dataset)
data_loader = torch.utils.data.DataLoader(dataset=data_set, batch_size=batch_size,
num_workers=workers, shuffle=False)
return data_loader
else:
raise ValueError
def __init__(self,
files,
channels,
patch_size,
transform=None,
noise_transform=None):
self.channels = channels
self.patch_size = patch_size
self.transform = transform
self.noise_transforms = noise_transform
self.to_tensor = ToTensor()
self.dataset = {'image': [], 'noisy': []}
self.load_dataset(files)
def get_transform(train=True,fixsize=False,img_size=416,min_size=800,max_size=1333,
image_mean=None,image_std=None,advanced=False):
if image_mean is None:image_mean = [0.485, 0.456, 0.406]
if image_std is None:image_std = [0.229, 0.224, 0.225]
if train:
transforms = Compose(
[
Augment(advanced),
ToTensor(),
ResizeFixSize(img_size) if fixsize else ResizeMinMax(min_size, max_size),
RandomHorizontalFlip(0.5),
Normalize(image_mean,image_std)
])
else:
transforms = Compose(
[
ToTensor(),
ResizeFixSize(img_size) if fixsize else ResizeMinMax(min_size, max_size),
# RandomHorizontalFlip(0.5),
Normalize(image_mean, image_std)
])
return transforms
def unit_test(image_dataset=True,
adni=True,
shuffle=True,
ntoshow=30,
show=True,
plt_show=True,
nworkers=0,
hiponly=True,
volume=False,
e2d=False):
'''
Tests vizualisation of a training batch of the dataset
'''
train_transforms = Compose([
ReturnPatch(patch_size=(32, 32)),
RandomAffine(),
Intensity(),
RandomFlip(modes=['horflip']),
Noisify(),
ToTensor()
])
#train_transforms = Compose((Resize(128, 128), RandomAffine(), ToTensor())) # for testing another stuff
print("Testing all orientations...")
for o in orientations:
if adni:
test = FloatHippocampusDataset(h5path=default_adni,
mode="train",
transform=train_transforms,
data_split=(0.5, 0.1, 0.4),
adni=True,
orientation=o,
hiponly=True,
return_volume=False,
e2d=True)
else:
test = FloatHippocampusDataset(mode="train",
transform=train_transforms,
orientation=o,
hiponly=hiponly,
return_volume=volume,
e2d=e2d)
test_loader = data.DataLoader(test,
batch_size=ntoshow,
shuffle=shuffle,
num_workers=0)
batch = next(iter(test_loader))
if show is True:
display_batch(batch, o + " dataloader")
if plt_show:
plt.show()
def test_data(data_path,img_size,classes,useMosaic,fixsize):
seed = 100
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
torch.manual_seed(seed)
kwargs = {'num_workers': 5, 'pin_memory': True} if use_cuda else {}
transforms = Compose(
[
Augment(True,True),
# Pad(),
# ToTensor(),
# Resize(img_size),
# RandomHorizontalFlip(0.5),
# Augment(False),
# ColorJitter(), SSDCropping(),
ToTensor(),
# ResizeFixSize(img_size) if fixsize else ResizeMinMax(800, 1333),
# RandomHorizontalFlip(0.5),
])
dataset = FruitsNutsDataset(data_path,transforms=transforms,classes=classes,useMosaic=useMosaic,fixsize=fixsize)
# dataset = PascalVOCDataset(data_path,transforms=transforms,classes=classes,useMosaic=useMosaic,fixsize=fixsize)
# for img, target in dataset:
# print()
data_loader = DataLoader(dataset, batch_size=4, shuffle=False, collate_fn=collate_fn, **kwargs)
for datas, targets in data_loader:
# datas = batch(datas)
for data, target in zip(datas, targets):
# from c,h,w ->h,w,c
data = data.permute(1, 2, 0)
# to uint8
data = torch.clamp(data * 255, 0, 255).to("cpu").numpy().astype(np.uint8)
# to BGR
# data = data[...,::-1]
data = cv2.cvtColor(data, cv2.COLOR_RGB2BGR)
boxes = target["boxes"].to("cpu").numpy().astype(np.int)
labels = target["labels"].to("cpu").numpy()
for idx, (box, label) in enumerate(zip(boxes, labels)):
data = vis_rect(data, box, str(label), 0.5, label, useMask=False)
cv2.imshow("test", data)
cv2.waitKey(0)
cv2.destroyAllWindows()
def __init__(self, root_dir, transforms=[ToTensor()]):
'''
Args:\n
:root_dir: string containing the root directory of the dataset.
:transforms: list of transformations to apply to the samples.
'''
self.root_dir = root_dir
self.transforms = transforms
self.composed_trsfm = Compose(transforms)
# save name of files in lists
self.images = glob.glob(os.path.join(root_dir, 'frames\\*.npz'))
self.labels = glob.glob(os.path.join(root_dir, 'masks\\*.npz'))
def get_transform_fixsize(train=True,img_size=416,
image_mean=None,image_std=None,advanced=False):
if image_mean is None:image_mean = [0.485, 0.456, 0.406]
if image_std is None:image_std = [0.229, 0.224, 0.225]
if train:
transforms = Compose(
[
Augment(advanced),
Pad(),
ToTensor(),
Resize(img_size),
RandomHorizontalFlip(0.5),
Normalize(image_mean,image_std)
])
else:
transforms = Compose(
[
Pad(),
ToTensor(),
Resize(img_size),
# RandomHorizontalFlip(0.5),
Normalize(image_mean, image_std)
])
return transforms
def init_lane(input_sizes, dataset, mean, std, base, workers=0):
transforms_test = Compose([
Resize(size_image=input_sizes, size_label=input_sizes),
ToTensor(),
Normalize(mean=mean, std=std)
])
validation_set = StandardLaneDetectionDataset(root=base,
image_set='val',
transforms=transforms_test,
data_set=dataset)
validation_loader = torch.utils.data.DataLoader(dataset=validation_set,
batch_size=1,
num_workers=workers,
shuffle=False)
return validation_loader
def __init__(self, options):
self.options = options
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
test_transform_list = []
if self.options.max_scale > 1:
test_transform_list.append(
RandomRescaleBB(1.0, self.options.max_scale))
test_transform_list.append(
CropAndResize(out_size=(self.options.crop_size,
self.options.crop_size)))
test_transform_list.append(
LocsToHeatmaps(out_size=(self.options.heatmap_size,
self.options.heatmap_size)))
test_transform_list.append(ToTensor())
test_transform_list.append(Normalize())
self.test_ds = RctaDataset(
root_dir=self.options.dataset_dir,
is_train=False,
transform=transforms.Compose(test_transform_list))
self.model = StackedHourglass(self.options.num_keypoints).to(
self.device)
# Only create optimizer because it is required to restore from checkpoint
self.optimizer = torch.optim.RMSprop(params=self.model.parameters(),
lr=0,
momentum=0,
weight_decay=0)
self.models_dict = {'stacked_hg': self.model}
self.optimizers_dict = {'optimizer': self.optimizer}
print("log dir:", options.log_dir)
print("checkpoint dir:", options.checkpoint_dir)
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
print("checkpoint:", self.options.checkpoint)
self.checkpoint = self.saver.load_checkpoint(
self.models_dict,
self.optimizers_dict,
checkpoint_file=self.options.checkpoint)
self.criterion = nn.MSELoss().to(self.device)
self.pose = Pose2DEval(detection_thresh=self.options.detection_thresh,
dist_thresh=self.options.dist_thresh)
def get_testing_loader(img_root, label_root, file_list, batch_size,
img_size, num_class):
transformed_dataset = VOCDataset(
img_root,
label_root,
file_list,
transform=transforms.Compose([
Resize(img_size),
ToTensor(),
Normalize(imagenet_stats['mean'], imagenet_stats['std']),
# GenOneHotLabel(num_class),
]))
loader = DataLoader(
transformed_dataset,
batch_size,
shuffle=False,
num_workers=0,
pin_memory=False,
)
return loader
def get_training_loader(img_root, label_root, file_list, batch_size,
img_height, img_width, num_class):
transformed_dataset = VOCTestDataset(
img_root,
label_root,
file_list,
transform=transforms.Compose([
RandomHorizontalFlip(),
Resize((img_height + 5, img_width + 5)),
RandomCrop((img_height, img_width)),
ToTensor(),
Normalize(imagenet_stats['mean'], imagenet_stats['std']),
# GenOneHotLabel(num_class),
]))
loader = DataLoader(
transformed_dataset,
batch_size,
shuffle=True,
num_workers=0,
pin_memory=False,
)
return loader
def init(batch_size, state, split, input_sizes, sets_id, std, mean, keep_scale, reverse_channels, data_set,
valtiny, no_aug):
# Return data_loaders/data_loader
# depending on whether the split is
# 1: semi-supervised training
# 2: fully-supervised training
# 3: Just testing
# Transformations (compatible with unlabeled data/pseudo labeled data)
# ! Can't use torchvision.Transforms.Compose
if data_set == 'voc':
base = base_voc
workers = 4
transform_train = Compose(
[ToTensor(keep_scale=keep_scale, reverse_channels=reverse_channels),
RandomResize(min_size=input_sizes[0], max_size=input_sizes[1]),
RandomCrop(size=input_sizes[0]),
RandomHorizontalFlip(flip_prob=0.5),
Normalize(mean=mean, std=std)])
if no_aug:
transform_train_pseudo = Compose(
[ToTensor(keep_scale=keep_scale, reverse_channels=reverse_channels),
Resize(size_image=input_sizes[0], size_label=input_sizes[0]),
Normalize(mean=mean, std=std)])
else:
transform_train_pseudo = Compose(
[ToTensor(keep_scale=keep_scale, reverse_channels=reverse_channels),
RandomResize(min_size=input_sizes[0], max_size=input_sizes[1]),
RandomCrop(size=input_sizes[0]),
RandomHorizontalFlip(flip_prob=0.5),
Normalize(mean=mean, std=std)])
transform_pseudo = Compose(
[ToTensor(keep_scale=keep_scale, reverse_channels=reverse_channels),
Resize(size_image=input_sizes[0], size_label=input_sizes[0]),
Normalize(mean=mean, std=std)])
transform_test = Compose(
[ToTensor(keep_scale=keep_scale, reverse_channels=reverse_channels),
ZeroPad(size=input_sizes[2]),
Normalize(mean=mean, std=std)])
elif data_set == 'city': # All the same size (whole set is down-sampled by 2)
base = base_city
workers = 8
transform_train = Compose(
[ToTensor(keep_scale=keep_scale, reverse_channels=reverse_channels),
RandomResize(min_size=input_sizes[0], max_size=input_sizes[1]),
RandomCrop(size=input_sizes[0]),
RandomHorizontalFlip(flip_prob=0.5),
Normalize(mean=mean, std=std),
LabelMap(label_id_map_city)])
if no_aug:
transform_train_pseudo = Compose(
[ToTensor(keep_scale=keep_scale, reverse_channels=reverse_channels),
Resize(size_image=input_sizes[0], size_label=input_sizes[0]),
Normalize(mean=mean, std=std)])
else:
transform_train_pseudo = Compose(
[ToTensor(keep_scale=keep_scale, reverse_channels=reverse_channels),
RandomResize(min_size=input_sizes[0], max_size=input_sizes[1]),
RandomCrop(size=input_sizes[0]),
RandomHorizontalFlip(flip_prob=0.5),
Normalize(mean=mean, std=std)])
transform_pseudo = Compose(
[ToTensor(keep_scale=keep_scale, reverse_channels=reverse_channels),
Resize(size_image=input_sizes[0], size_label=input_sizes[0]),
Normalize(mean=mean, std=std),
LabelMap(label_id_map_city)])
transform_test = Compose(
[ToTensor(keep_scale=keep_scale, reverse_channels=reverse_channels),
Resize(size_image=input_sizes[2], size_label=input_sizes[2]),
Normalize(mean=mean, std=std),
LabelMap(label_id_map_city)])
else:
base = ''
# Not the actual test set (i.e.validation set)
test_set = StandardSegmentationDataset(root=base, image_set='valtiny' if valtiny else 'val',
transforms=transform_test, label_state=0, data_set=data_set)
val_loader = torch.utils.data.DataLoader(dataset=test_set, batch_size=batch_size, num_workers=workers, shuffle=False)
# Testing
if state == 3:
return val_loader
else:
# Fully-supervised training
if state == 2:
labeled_set = StandardSegmentationDataset(root=base, image_set=(str(split) + '_labeled_' + str(sets_id)),
transforms=transform_train, label_state=0, data_set=data_set)
labeled_loader = torch.utils.data.DataLoader(dataset=labeled_set, batch_size=batch_size,
num_workers=workers, shuffle=True)
return labeled_loader, val_loader
# Semi-supervised training
elif state == 1:
pseudo_labeled_set = StandardSegmentationDataset(root=base, data_set=data_set,
image_set=(str(split) + '_unlabeled_' + str(sets_id)),
transforms=transform_train_pseudo, label_state=1)
reference_set = SegmentationLabelsDataset(root=base, image_set=(str(split) + '_unlabeled_' + str(sets_id)),
data_set=data_set)
reference_loader = torch.utils.data.DataLoader(dataset=reference_set, batch_size=batch_size,
num_workers=workers, shuffle=False)
unlabeled_set = StandardSegmentationDataset(root=base, data_set=data_set,
image_set=(str(split) + '_unlabeled_' + str(sets_id)),
transforms=transform_pseudo, label_state=2)
labeled_set = StandardSegmentationDataset(root=base, data_set=data_set,
image_set=(str(split) + '_labeled_' + str(sets_id)),
transforms=transform_train, label_state=0)
unlabeled_loader = torch.utils.data.DataLoader(dataset=unlabeled_set, batch_size=batch_size,
num_workers=workers, shuffle=False)
pseudo_labeled_loader = torch.utils.data.DataLoader(dataset=pseudo_labeled_set,
batch_size=int(batch_size / 2),
num_workers=workers, shuffle=True)
labeled_loader = torch.utils.data.DataLoader(dataset=labeled_set,
batch_size=int(batch_size / 2),
num_workers=workers, shuffle=True)
return labeled_loader, pseudo_labeled_loader, unlabeled_loader, val_loader, reference_loader
else:
# Support unsupervised learning here if that's what you want
raise ValueError
title='info_loss',
ytickmin=0,
ytinkmax=10)),
acc=Scalar(vis,
'Accuracy',
opts=dict(showlegend=True,
title='Accuracy',
ytickmin=0,
ytinkmax=2.0)),
inputs=Image3D(vis, 'inputs'),
outputs=Image3D(vis, 'outputs'))
# create train set, x = image, y=target
x, y, train_idx, test_idx, ratio = fold_split(FG)
#transform=Compose([ToWoldCoordinateSystem(), Normalize((0.5, 0.9)), ToTensor()])
transform = Compose([ToWoldCoordinateSystem(), ToTensor()])
trainset = ADNIDataset(FG, x[train_idx], y[train_idx], transform=transform)
testset = ADNIDataset(FG, x[test_idx], y[test_idx], transform=transform)
trainloader = DataLoader(trainset,
batch_size=FG.batch_size,
shuffle=True,
pin_memory=True,
num_workers=4)
testloader = DataLoader(testset,
batch_size=FG.batch_size,
shuffle=True,
num_workers=4,
pin_memory=True)
def init(batch_size,
state,
input_sizes,
dataset,
mean,
std,
base,
workers=10,
method='baseline'):
# Return data_loaders
# depending on whether the state is
# 0: training
# 1: fast validation by mean IoU (validation set)
# 2: just testing (test set)
# 3: just testing (validation set)
# Transformations
# ! Can't use torchvision.Transforms.Compose
transforms_test = Compose([
Resize(size_image=input_sizes[0], size_label=input_sizes[0]),
ToTensor(),
Normalize(mean=mean, std=std)
])
transforms_train = Compose([
Resize(size_image=input_sizes[0], size_label=input_sizes[0]),
RandomRotation(degrees=3),
ToTensor(),
Normalize(mean=mean,
std=std,
normalize_target=True if method == 'lstr' else False)
])
# Batch builder
if method == 'lstr':
collate_fn = dict_collate_fn
else:
collate_fn = None
if state == 0:
if method == 'lstr':
if dataset == 'tusimple':
data_set = TuSimple(root=base,
image_set='train',
transforms=transforms_train,
padding_mask=True,
process_points=True)
elif dataset == 'culane':
data_set = CULane(root=base,
image_set='train',
transforms=transforms_train,
padding_mask=True,
process_points=True)
else:
raise ValueError
else:
data_set = StandardLaneDetectionDataset(
root=base,
image_set='train',
transforms=transforms_train,
data_set=dataset)
data_loader = torch.utils.data.DataLoader(dataset=data_set,
batch_size=batch_size,
collate_fn=collate_fn,
num_workers=workers,
shuffle=True)
validation_set = StandardLaneDetectionDataset(
root=base,
image_set='val',
transforms=transforms_test,
data_set=dataset)
validation_loader = torch.utils.data.DataLoader(dataset=validation_set,
batch_size=batch_size *
4,
num_workers=workers,
shuffle=False,
collate_fn=collate_fn)
return data_loader, validation_loader
elif state == 1 or state == 2 or state == 3:
image_sets = ['valfast', 'test', 'val']
if method == 'lstr':
if dataset == 'tusimple':
data_set = TuSimple(root=base,
image_set=image_sets[state - 1],
transforms=transforms_test,
padding_mask=False,
process_points=False)
elif dataset == 'culane':
data_set = CULane(root=base,
image_set=image_sets[state - 1],
transforms=transforms_test,
padding_mask=False,
process_points=False)
else:
raise ValueError
else:
data_set = StandardLaneDetectionDataset(
root=base,
image_set=image_sets[state - 1],
transforms=transforms_test,
data_set=dataset)
data_loader = torch.utils.data.DataLoader(dataset=data_set,
batch_size=batch_size,
collate_fn=collate_fn,
num_workers=workers,
shuffle=False)
return data_loader
else:
raise ValueError
def __init__(
self,
data_src,
folds2include=None,
num_folds=5,
samples_per_epoch=2000,
roi_size=96,
scale_int=(0, 255),
norm_mean=0.,
norm_sd=1.,
zoom_range=(0.90, 1.1),
prob_unseeded_patch=0.2,
int_aug_offset=None,
int_aug_expansion=None,
valid_labels=None, # if this is None it will include all the available labels
is_preloaded=False,
max_input_size=2048 #if any image is larger than this it will be splitted (only working with folds)
):
_dum = set(dir(self))
self.data_src = Path(data_src)
if not self.data_src.exists():
raise ValueError(f'`data_src` : `{data_src}` does not exists.')
self.folds2include = folds2include
self.num_folds = num_folds
self.samples_per_epoch = samples_per_epoch
self.roi_size = roi_size
self.scale_int = scale_int
self.norm_mean = norm_mean
self.norm_sd = norm_sd
self.zoom_range = zoom_range
self.prob_unseeded_patch = prob_unseeded_patch
self.int_aug_offset = int_aug_offset
self.int_aug_expansion = int_aug_expansion
self.valid_labels = valid_labels
self.is_preloaded = is_preloaded
self.max_input_size = max_input_size
self._input_names = list(
set(dir(self)) - _dum
) #i want the name of this fields so i can access them if necessary
rotation_pad_size = math.ceil(self.roi_size * (math.sqrt(2) - 1) / 2)
padded_roi_size = roi_size + 2 * rotation_pad_size
transforms_random = [
RandomCropWithSeeds(padded_roi_size, rotation_pad_size,
prob_unseeded_patch),
AffineTransform(zoom_range),
RemovePadding(rotation_pad_size),
RandomVerticalFlip(),
RandomHorizontalFlip(),
NormalizeIntensity(scale_int, norm_mean, norm_sd),
RandomIntensityOffset(int_aug_offset),
RandomIntensityExpansion(int_aug_expansion),
OutContours2Segmask(),
FixDTypes()
#I cannot really pass the ToTensor to the dataloader since it crashes when the batchsize is large (>256)
]
self.transforms_random = Compose(transforms_random)
transforms_full = [
NormalizeIntensity(scale_int),
OutContours2Segmask(),
FixDTypes(),
ToTensor()
]
self.transforms_full = Compose(transforms_full)
self.hard_neg_data = None
if self.data_src.is_dir():
assert self.folds2include is None
self.data = self.load_data_from_dir(self.data_src, padded_roi_size,
self.is_preloaded)
else:
assert self.is_preloaded
self.data = self.load_data_from_file(self.data_src)
self.type_ids = sorted(list(self.data.keys()))
self.types2label = {k: (ii + 1) for ii, k in enumerate(self.type_ids)}
self.num_clases = len(self.type_ids)
#flatten data so i can access the whole list by index
self.data_indexes = [(_type, _fname, ii)
for _type, type_data in self.data.items()
for _fname, file_data in type_data.items()
for ii in range(len(file_data))]
assert len(self.data_indexes) > 0 #makes sure there are valid files
def main(args=None):
global thres
global rel_thresh
global attr_thresh
parser = argparse.ArgumentParser(description='Simple training script for training a RetinaNet network.')
parser.add_argument('--dataset', help='Dataset type, must be one of csv or coco.')
parser.add_argument('--data_path', help='Path to COCO directory')
parser.add_argument('--csv_classes', help='Path to file containing class list (see readme)')
parser.add_argument('--csv_val', help='Path to file containing validation annotations (optional, see readme)')
parser.add_argument('--net', help='Network to use', default='fasterrcnn')
parser.add_argument('--set', help='Set on which evaluation will be performed', default='validation')
parser.add_argument('--store_detections', action='store_true', default=False,
help='Cache all detections with very low threshold in order to enable filtering after extraction')
parser.add_argument('--load_detections', action='store_true', default=False,
help='Load cached detections')
parser.add_argument('--model_rel', help='Path to model (.pt) file for relationships.', default=None)
parser.add_argument('--model_attr', help='Path to model (.pt) file for attributes.', default=None)
parser.add_argument('--model_detector', help='Path to model (.pt) file for the detector.')
parser.add_argument('--depth', help='Resnet depth, must be one of 18, 34, 50, 101, 152', type=int, default=50)
parser = parser.parse_args(args)
assert parser.model_rel is not None and parser.model_attr is not None and parser.model_detector is not None, \
'Models snapshots have to be specified!'
assert not (parser.load_detections and parser.store_detections)
det_output_path = os.path.split(parser.model_rel)[0]
if parser.dataset == 'openimages':
dataset_val = OidDatasetVRD(parser.data_path, subset=parser.set,
transform=Compose([ToTensor()]))
else:
raise ValueError('Dataset type not understood (must be csv or coco), exiting.')
#sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
dataloader_val = DataLoader(dataset_val, num_workers=1, collate_fn=collate_fn, batch_size=1)
# Create the model
detector = create_detection_model(dataset_val.num_classes(), parser, box_score_thresh=thres)
model = VRD(detector, dataset=dataset_val, train_relationships=parser.model_rel is not None,
train_attributes=parser.model_attr is not None, max_objects=max_objects)
# Load the detector
checkpoint = torch.load(parser.model_detector, map_location=lambda storage, loc: storage)
weights = checkpoint['model']
weights = {k.replace('module.', ''): v for k, v in weights.items()}
model.detector.load_state_dict(weights)
print('Detector correctly loaded!')
# Load the attributes, if needed
if parser.model_rel:
checkpoint = torch.load(parser.model_rel, map_location=lambda storage, loc: storage)
weights = checkpoint['model_rel']
weights = {k.replace('module.', ''): v for k, v in weights.items()}
model.relationships_net.load_state_dict(weights)
print('Relationships correctly loaded!')
if parser.model_attr:
checkpoint = torch.load(parser.model_attr, map_location=lambda storage, loc: storage)
weights = checkpoint['model_attr']
weights = {k.replace('module.', ''): v for k, v in weights.items()}
model.attributes_net.load_state_dict(weights)
print('Attributes correctly loaded!')
if use_gpu:
model = model.cuda()
model.eval()
all_detections = []
if parser.load_detections or parser.store_detections:
print('Opening detections database file...')
flag = 'r' if parser.load_detections else 'c'
loaded_detections = shelve.open(os.path.join(det_output_path, 'cached_detections_detthr{}.db'.format(thres)), flag=flag)
for idx, data in enumerate(tqdm.tqdm(dataloader_val)):
if parser.load_detections:
loaded_det = loaded_detections[str(idx)]
scores = loaded_det[0]
classification = loaded_det[1]
boxes = loaded_det[2]
relationships = loaded_det[3]
rel_scores = loaded_det[4]
attributes = loaded_det[5]
attr_scores = loaded_det[6]
else:
with torch.no_grad():
st = time.time()
images, targets = data
# targets = [{k: v.cuda() for k, v in t.items()} for t in targets]
if use_gpu:
input_images = list(image.cuda().float() for image in images)
else:
input_images = list(image.float() for image in images)
# TODO: adapt retinanet output to the one by torchvision 0.3
# scores, classification, transformed_anchors = model(data_img.float())
outputs = model(input_images)
outputs = [{k: v.cpu() for k, v in t.items()} for t in outputs]
output = outputs[0] # take the only batch
scores = output['scores']
classification = output['labels']
boxes = output['boxes']
relationships = output['relationships']
rel_scores = output['relationships_scores']
attributes = output['attributes']
attr_scores = output['attributes_scores']
if parser.store_detections:
loaded_detections[str(idx)] = [scores, classification, boxes, relationships, rel_scores, attributes, attr_scores]
else:
'''if parser.load_detections:
pdb.set_trace()
# filter objects, relationships and attributes
filtered_idxs = np.where(scores > thres)[0]
scores = scores[filtered_idxs]
classification = classification[filtered_idxs]
boxes = boxes[filtered_idxs]
relationships = relationships[np.ix_(filtered_idxs, filtered_idxs)]
rel_scores = rel_scores[np.ix_(filtered_idxs, filtered_idxs)]
attributes = attributes[filtered_idxs]
attr_scores = attr_scores[filtered_idxs]
'''
subj_boxes_out = []
subj_labels_out = []
obj_boxes_out = []
obj_labels_out = []
rel_labels_out = []
rel_scores_out = []
if len(boxes) != 0:
# num_objects = min(boxes.shape[0], max_objects)
# Collect objects and attributes
for j in range(attributes.shape[0]):
bbox = boxes[j, :4]
attr = attributes[j, 0].item() if parser.model_attr is not None and attr_scores[j, 0] > attr_thresh else 0 # TODO: only the top rank attribute is considered, generalize better!
# We add an 'is' relation. 'is' relation is mapped to relation index of -1.
if attr != 0:
subj_boxes_out.append(bbox)
obj_boxes_out.append(bbox)
rel_labels_out.append(-1)
rel_scores_out.append(attr_scores[j, 0])
subj_labels_out.append(int(classification[j]))
obj_labels_out.append(attr)
# Collect relationships
for s_ind in range(relationships.shape[0]):
for o_ind in range(relationships.shape[1]):
subj = boxes[s_ind, :4]
obj = boxes[o_ind, :4]
rel = relationships[s_ind, o_ind].item() if rel_scores[s_ind, o_ind] > rel_thresh else 0
if rel != 0:
subj_boxes_out.append(subj)
obj_boxes_out.append(obj)
rel_labels_out.append(rel)
rel_scores_out.append(rel_scores[s_ind, o_ind])
subj_labels_out.append(int(classification[s_ind]))
obj_labels_out.append(int(classification[o_ind]))
all_detections.append([idx, subj_boxes_out, subj_labels_out, obj_boxes_out, obj_labels_out, rel_labels_out, rel_scores_out])
# if idx == 400:
# break
if not parser.store_detections:
print('Evaluating...')
# TODO: add identification parameter to evaluate so that detections from different checkpoints are not overwritten
dataset_val.evaluate(all_detections, det_output_path, file_identifier='{}_relthr{}_attrthr{}_detthr{}'.format(parser.set, rel_thresh, attr_thresh, thres))
print('DONE!')
showlegend=True, title='DG_z2', ytickmin=0, ytinkmax=2.0)),
D_x = Scalar(vis, 'D_x', opts=dict(
showlegend=True, title='D_x', ytickmin=0, ytinkmax=2.0)),
inputs0 = Image3D(vis, 'inputs0'),
inputs1 = Image3D(vis, 'inputs1'),
fake0 = Image3D(vis, 'fake0'),
fake1 = Image3D(vis, 'fake1'),
outputs0 = Image3D(vis, 'outputs0'),
outputs1 = Image3D(vis, 'outputs1'))
# dataset setting
x, y = Trainset(FG)
# x, y, train_idx, test_idx, ratio = fold_split(FG)
# transform = Compose([ToFloatTensor(), Normalize(0.5,0.5)])
# trainset = ADNIDataset2D(FG, x, y, transform=transform)
transform=Compose([ToWoldCoordinateSystem(), ToTensor(), Pad(1,0,1,0,1,0), Normalize(0.5,0.5)])
trainset = ADNIDataset(FG, x, y, transform=transform)
trainloader = DataLoader(trainset, batch_size=FG.batch_size,
shuffle=True, pin_memory=True)
# trainset = ADNIDataset2D(FG, x[train_idx], y[train_idx], transform=transform)
# testset = ADNIDataset2D(FG, x[test_idx], y[test_idx], transform=transform)
# trainloader = DataLoader(trainset, batch_size=FG.batch_size, shuffle=True,
# pin_memory=True, num_workers=4)
# testloader = DataLoader(testset, batch_size=FG.batch_size, shuffle=True,
# num_workers=4, pin_memory=True)
# models
D = infoDiscriminator3D(FG.c_code).to('cuda:{}'.format(FG.devices[0]))
G = infoGenerator3D(FG.z_dim, FG.c_code).to('cuda:{}'.format(FG.devices[0]))
if len(FG.devices) != 1:
ytinkmax=2.0)),
inputs=Image3D(vis, 'inputs'),
fake=Image3D(vis, 'fake'),
valid=Image3D(vis, 'valid'),
outputs=Image3D(vis, 'outputs'),
outputs2=Image3D(vis, 'outputs2'))
# x, y = Trainset(FG) # x = image, y=target
x, y, train_idx, test_idx, ratio = fold_split(FG)
# transform=Compose([ToFloatTensor(), Normalize(0.5,0.5)])
# trainset = ADNIDataset2D(FG, x[train_idx], y[train_idx], transform=transform)
# testset = ADNIDataset2D(FG, x[test_idx], y[test_idx], transform=transform)
transform = Compose([
ToWoldCoordinateSystem(),
ToTensor(),
Pad(1, 0, 1, 0, 1, 0),
Normalize(0.5, 0.5)
])
trainset = ADNIDataset(FG, x[train_idx], y[train_idx], transform=transform)
testset = ADNIDataset(FG, x[test_idx], y[test_idx], transform=transform)
trainloader = DataLoader(trainset,
batch_size=FG.batch_size,
shuffle=True,
pin_memory=True,
num_workers=4)
testloader = DataLoader(testset,
batch_size=FG.batch_size,
shuffle=True,
num_workers=4,
def main(args):
print(args)
torch.backends.cudnn.benchmark = True
# Data loading
normalize = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
base_size = 320
crop_size = 256
min_size, max_size = int(0.5 * base_size), int(2.0 * base_size)
interpolation_mode = InterpolationMode.BILINEAR
train_loader, val_loader = None, None
if not args.test_only:
st = time.time()
train_set = VOCSegmentation(args.data_path,
image_set='train',
download=True,
transforms=Compose([
RandomResize(min_size, max_size,
interpolation_mode),
RandomCrop(crop_size),
RandomHorizontalFlip(0.5),
ImageTransform(
T.ColorJitter(brightness=0.3,
contrast=0.3,
saturation=0.1,
hue=0.02)),
ToTensor(),
ImageTransform(normalize)
]))
train_loader = torch.utils.data.DataLoader(
train_set,
batch_size=args.batch_size,
drop_last=True,
sampler=RandomSampler(train_set),
num_workers=args.workers,
pin_memory=True,
worker_init_fn=worker_init_fn)
print(f"Training set loaded in {time.time() - st:.2f}s "
f"({len(train_set)} samples in {len(train_loader)} batches)")
if args.show_samples:
x, target = next(iter(train_loader))
plot_samples(x, target, ignore_index=255)
return
if not (args.lr_finder or args.check_setup):
st = time.time()
val_set = VOCSegmentation(args.data_path,
image_set='val',
download=True,
transforms=Compose([
Resize((crop_size, crop_size),
interpolation_mode),
ToTensor(),
ImageTransform(normalize)
]))
val_loader = torch.utils.data.DataLoader(
val_set,
batch_size=args.batch_size,
drop_last=False,
sampler=SequentialSampler(val_set),
num_workers=args.workers,
pin_memory=True,
worker_init_fn=worker_init_fn)
print(
f"Validation set loaded in {time.time() - st:.2f}s ({len(val_set)} samples in {len(val_loader)} batches)"
)
if args.source.lower() == 'holocron':
model = segmentation.__dict__[args.arch](args.pretrained,
num_classes=len(VOC_CLASSES))
elif args.source.lower() == 'torchvision':
model = tv_segmentation.__dict__[args.arch](
args.pretrained, num_classes=len(VOC_CLASSES))
# Loss setup
loss_weight = None
if isinstance(args.bg_factor, float) and args.bg_factor != 1:
loss_weight = torch.ones(len(VOC_CLASSES))
loss_weight[0] = args.bg_factor
if args.loss == 'crossentropy':
criterion = nn.CrossEntropyLoss(weight=loss_weight,
ignore_index=255,
label_smoothing=args.label_smoothing)
elif args.loss == 'focal':
criterion = holocron.nn.FocalLoss(weight=loss_weight, ignore_index=255)
elif args.loss == 'mc':
criterion = holocron.nn.MutualChannelLoss(weight=loss_weight,
ignore_index=255,
xi=3)
# Optimizer setup
model_params = [p for p in model.parameters() if p.requires_grad]
if args.opt == 'sgd':
optimizer = torch.optim.SGD(model_params,
args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
elif args.opt == 'radam':
optimizer = holocron.optim.RAdam(model_params,
args.lr,
betas=(0.95, 0.99),
eps=1e-6,
weight_decay=args.weight_decay)
elif args.opt == 'adamp':
optimizer = holocron.optim.AdamP(model_params,
args.lr,
betas=(0.95, 0.99),
eps=1e-6,
weight_decay=args.weight_decay)
elif args.opt == 'adabelief':
optimizer = holocron.optim.AdaBelief(model_params,
args.lr,
betas=(0.95, 0.99),
eps=1e-6,
weight_decay=args.weight_decay)
log_wb = lambda metrics: wandb.log(metrics) if args.wb else None
trainer = SegmentationTrainer(model,
train_loader,
val_loader,
criterion,
optimizer,
args.device,
args.output_file,
num_classes=len(VOC_CLASSES),
amp=args.amp,
on_epoch_end=log_wb)
if args.resume:
print(f"Resuming {args.resume}")
checkpoint = torch.load(args.resume, map_location='cpu')
trainer.load(checkpoint)
if args.show_preds:
x, target = next(iter(train_loader))
with torch.no_grad():
if isinstance(args.device, int):
x = x.cuda()
trainer.model.eval()
preds = trainer.model(x)
plot_predictions(x.cpu(), preds.cpu(), target, ignore_index=255)
return
if args.test_only:
print("Running evaluation")
eval_metrics = trainer.evaluate()
print(
f"Validation loss: {eval_metrics['val_loss']:.4} (Mean IoU: {eval_metrics['mean_iou']:.2%})"
)
return
if args.lr_finder:
print("Looking for optimal LR")
trainer.lr_find(args.freeze_until,
norm_weight_decay=args.norm_weight_decay,
num_it=min(len(train_loader), 100))
trainer.plot_recorder()
return
if args.check_setup:
print("Checking batch overfitting")
is_ok = trainer.check_setup(args.freeze_until,
args.lr,
norm_weight_decay=args.norm_weight_decay,
num_it=min(len(train_loader), 100))
print(is_ok)
return
# Training monitoring
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
exp_name = f"{args.arch}-{current_time}" if args.name is None else args.name
# W&B
if args.wb:
run = wandb.init(name=exp_name,
project="holocron-semantic-segmentation",
config={
"learning_rate": args.lr,
"scheduler": args.sched,
"weight_decay": args.weight_decay,
"epochs": args.epochs,
"batch_size": args.batch_size,
"architecture": args.arch,
"source": args.source,
"input_size": 256,
"optimizer": args.opt,
"dataset": "Pascal VOC2012 Segmentation",
"loss": args.loss,
})
print("Start training")
start_time = time.time()
trainer.fit_n_epochs(args.epochs,
args.lr,
args.freeze_until,
args.sched,
norm_weight_decay=args.norm_weight_decay)
total_time_str = str(
datetime.timedelta(seconds=int(time.time() - start_time)))
print(f"Training time {total_time_str}")
if args.wb:
run.finish()
def main():
gflags.DEFINE_string('id', None, 'ID for Training')
gflags.DEFINE_string('epoch', 25, 'Number of Epochs')
gflags.DEFINE_string('pretrained', None, 'Pretrained for Resuming Training')
gflags.DEFINE_string('threshold', 0.5, 'Threshold probability for predicting class')
gflags.DEFINE_string('batchsize', 128, 'Batch Size')
gflags.DEFINE_string('gpu', True, 'Use GPU or Not')
gflags.DEFINE_string('lr', 0.001, 'Learning Rate')
gflags.DEFINE_string('class_name', 'None', 'class name')
gflags.FLAGS(sys.argv)
# Directory Path for saving weights of Trained Model
save_path = 'Train_id' + str(gflags.FLAGS.id)
class_name = gflags.FLAGS.class_name
threshold = gflags.FLAGS.threshold
class_name = gflags.FLAGS.class_name
writer = SummaryWriter('./runs/{}'.format(gflags.FLAGS.id))
if not os.path.isdir(save_path):
os.mkdir(save_path)
os.mkdir(save_path + '/Checkpoint')
train_dataset_path = 'data/train'
val_dataset_path = 'data/valid'
train_transform = transforms.Compose([
ImgAugTransform(),
ToTensor()
])
valid_transform = transforms.Compose([
ToTensor()
])
train_dataset = TrainDataset(path=train_dataset_path, transform=valid_transform, class_name=class_name)
val_dataset = TrainDataset(path=val_dataset_path, transform=valid_transform, class_name=class_name)
sampler = WeightedRandomSampler(torch.DoubleTensor(train_dataset.weights), len(train_dataset.weights))
train_dataloader = DataLoader(train_dataset, batch_size=4,
pin_memory=True, num_workers=4)
val_dataloader = DataLoader(val_dataset, batch_size=4, shuffle=False,
pin_memory=True, num_workers=4)
size_train = len(train_dataloader)
size_val = len(val_dataloader)
print('Number of Training Images: {}'.format(size_train))
print('Number of Validation Images: {}'.format(size_val))
# Reads class weights from a Json file
with open('class_weights.json', 'r') as fp:
class_weights = json.load(fp)
weight = torch.tensor([1/class_weights[class_name]])
start_epoch = 0
if class_name in ['Roads', 'Railway']:
model = DinkNet34(num_classes = 1)
else:
model = Unet(n_ch=4, n_classes=1)
if pretrained is not None:
criterion = FocalLoss()
else:
criterion = LogDiceLoss()
criterion1 = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=float(gflags.FLAGS.lr))
if gflags.FLAGS.gpu:
model = model.cuda()
criterion = criterion.cuda()
criterion1 = criterion1.cuda()
if gflags.FLAGS.pretrained is not None:
weight_path = sorted(os.listdir('./weights/' + save_path+ '/Checkpoint/'), key=lambda x:float(x[:-8]))[0]
checkpoint = torch.load('./weights/' + save_path + '/Checkpoint/' + weight_path)
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print('Loaded Checkpoint of Epoch: {}'.format(gflags.FLAGS.weight))
for epoch in range(start_epoch, int(gflags.FLAGS.epoch) + start_epoch):
print("epoch {}".format(epoch))
train(model, train_dataloader, criterion, criterion1, optimizer, epoch, writer, size_train, threshold)
print('')
val_loss = val(model, val_dataloader, criterion, criterion1, epoch, writer, size_val, threshold)
print('')
save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, filename= save_path + '/Checkpoint/' + str(val_loss) + '.pth.tar')
writer.export_scalars_to_json(save_path + 'log.json')