def load_dataset_al_train(dataset, args, batchsize):
image_transform = transforms.Compose([
transforms.Resize((args.height_train, args.width_train)),
transforms.ToTensor()
])
label_transform = transforms.Compose([
transforms.Resize((args.height_train, args.width_train),
Image.NEAREST),
ext_transforms.PILToLongTensor()
])
train_al_set = dataset(args.dataset_dir,
mode='train_al',
transform=image_transform,
label_transform=label_transform)
train_al_loader = data.DataLoader(
train_al_set,
batch_size=batchsize,
shuffle=True,
# drop_last=True,
num_workers=args.workers) #args.workers)
class_encoding = train_al_set.color_encoding.copy()
if args.dataset.lower() == 'camvid':
# if 'road_marking' in class_encoding:
del class_encoding['road_marking']
# num_classes = len(class_encoding)
print("Size of data in the dataloader", len(train_al_set))
return (train_al_loader), class_encoding
def predict():
image_transform = transforms.Compose(
[transforms.Resize(target_size),
transforms.ToTensor()])
label_transform = transforms.Compose(
[transforms.Resize(target_size),
ext_transforms.PILToLongTensor()])
# Get selected dataset
# Load the training set as tensors
train_set = Cityscapes(data_dir,
mode='test',
transform=image_transform,
label_transform=label_transform)
class_encoding = train_set.color_encoding
num_classes = len(class_encoding)
model = ENet(num_classes).to(device)
# Initialize a optimizer just so we can retrieve the model from the
# checkpoint
optimizer = optim.Adam(model.parameters())
# Load the previoulsy saved model state to the ENet model
model = utils.load_checkpoint(model, optimizer, 'save',
'ENet_cityscapes_mine.pth')[0]
# print(model)
image = Image.open('images/mainz_000000_008001_leftImg8bit.png')
images = Variable(image_transform(image).to(device).unsqueeze(0))
image = np.array(image)
# Make predictions!
predictions = model(images)
_, predictions = torch.max(predictions.data, 1)
# 0~18
prediction = predictions.cpu().numpy()[0] - 1
mask_color = np.asarray(label_to_color_image(prediction, 'cityscapes'),
dtype=np.uint8)
mask_color = cv2.resize(mask_color, (image.shape[1], image.shape[0]))
print(image.shape)
print(mask_color.shape)
res = cv2.addWeighted(image, 0.3, mask_color, 0.7, 0.6)
# cv2.imshow('rr', mask_color)
cv2.imshow('combined', res)
cv2.waitKey(0)
def load_dataset(dataset):
print("\nLoading dataset...\n")
print("Selected dataset:", args.dataset)
print("Dataset directory:", args.dataset_dir)
print("Save directory:", args.save_dir)
image_transform = transforms.Compose(
[transforms.Resize((args.height, args.width)),
transforms.ToTensor()])
label_transform = transforms.Compose([
transforms.Resize((args.height, args.width), transforms.InterpolationMode.NEAREST),
ext_transforms.PILToLongTensor()
])
# Get selected dataset
# Load the training set as tensors
train_set = dataset(
args.dataset_dir,
transform=image_transform,
label_transform=label_transform)
train_loader = data.DataLoader(
train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
# Load the validation set as tensors
val_set = dataset(
args.dataset_dir,
mode='val',
transform=image_transform,
label_transform=label_transform)
val_loader = data.DataLoader(
val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
# Load the test set as tensors
test_set = dataset(
args.dataset_dir,
mode='test',
transform=image_transform,
label_transform=label_transform)
test_loader = data.DataLoader(
test_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
# Get encoding between pixel valus in label images and RGB colors
class_encoding = train_set.color_encoding
# Remove the road_marking class from the CamVid dataset as it's merged
# with the road class
if args.dataset.lower() == 'camvid':
del class_encoding['road_marking']
# Get number of classes to predict
num_classes = len(class_encoding)
# Print information for debugging
print("Number of classes to predict:", num_classes)
print("Train dataset size:", len(train_set))
print("Validation dataset size:", len(val_set))
# Get a batch of samples to display
if args.mode.lower() == 'test':
images, labels = iter(test_loader).next()
else:
images, labels = iter(train_loader).next()
print("Image size:", images.size())
print("Label size:", labels.size())
print("Class-color encoding:", class_encoding)
# Show a batch of samples and labels
if args.imshow_batch:
print("Close the figure window to continue...")
label_to_rgb = transforms.Compose([
ext_transforms.LongTensorToRGBPIL(class_encoding),
transforms.ToTensor()
])
color_labels = utils.batch_transform(labels, label_to_rgb)
utils.imshow_batch(images, color_labels)
# Get class weights from the selected weighing technique
print("\nWeighing technique:", args.weighing)
print("Computing class weights...")
print("(this can take a while depending on the dataset size)")
class_weights = 0
if args.weighing.lower() == 'enet':
class_weights = enet_weighing(train_loader, num_classes)
elif args.weighing.lower() == 'mfb':
class_weights = median_freq_balancing(train_loader, num_classes)
else:
class_weights = None
if class_weights is not None:
class_weights = torch.from_numpy(class_weights).float().to(device)
# Set the weight of the unlabeled class to 0
if args.ignore_unlabeled:
ignore_index = list(class_encoding).index('unlabeled')
class_weights[ignore_index] = 0
print("Class weights:", class_weights)
return (train_loader, val_loader,
test_loader), class_weights, class_encoding
def load_dataset_activelearning(dataset, args):
print("\nLoading dataset...")
print("Selected Backbone:", args.backbone)
print("Selected dataset:", args.dataset)
print("Dataset directory:", args.dataset_dir)
print("Save directory:", args.save_dir)
image_transform = transforms.Compose(
[transforms.Resize((args.height, args.width)),
transforms.ToTensor()])
label_transform = transforms.Compose([
transforms.Resize((args.height, args.width), Image.NEAREST),
ext_transforms.PILToLongTensor()
])
train_set = dataset(args.dataset_dir,
transform=image_transform,
label_transform=label_transform)
train_loader = data.DataLoader(
train_set,
batch_size=args.batch_size_AL,
shuffle=False,
pin_memory=True,
# drop_last=True,
num_workers=args.workers)
val_set = dataset(args.dataset_dir,
mode='val',
transform=image_transform,
label_transform=label_transform)
val_loader = data.DataLoader(
val_set,
batch_size=args.batch_size_AL,
shuffle=False,
pin_memory=True,
# drop_last=True,
num_workers=args.workers)
# Get encoding between pixel valus in label images and RGB colors
class_encoding = train_set.color_encoding.copy()
# Remove the road_marking class from the CamVid dataset as it's merged
# with the road class
if args.dataset.lower() == 'camvid':
del class_encoding['road_marking']
# Get number of classes to predict
num_classes = len(class_encoding)
# Print information for debuggingFs
print("Number of classes to predict:", num_classes)
print("Train dataset size:", len(train_set))
print("Required steps for each epoch: {}".format(len(train_set) // 1))
print("Validation dataset size:", len(val_set))
# Get a batch of samples to display
images, labels, _, _ = iter(train_loader).next()
print("Image size:", images.size())
print("Label size:", labels.size())
print("\n")
print("Class-color encoding:", class_encoding)
return (train_loader, val_loader), class_encoding
def load_dataset(dataset):
print("\nLoading dataset...\n")
print("Selected dataset:", args.dataset)
print("Dataset directory:", args.dataset_dir)
print('Train file:', args.trainFile)
print('Val file:', args.valFile)
print('Test file:', args.testFile)
print("Save directory:", args.save_dir)
image_transform = transforms.Compose(
[transforms.Resize((args.height, args.width)),
transforms.ToTensor()])
label_transform = transforms.Compose([
transforms.Resize((args.height, args.width)),
ext_transforms.PILToLongTensor()
])
# Get selected dataset
# Load the training set as tensors
train_set = dataset(args.dataset_dir, args.trainFile, mode='train', transform=image_transform, \
label_transform=label_transform, color_mean=color_mean, color_std=color_std)
train_loader = data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
# Load the validation set as tensors
val_set = dataset(args.dataset_dir, args.valFile, mode='val', transform=image_transform, \
label_transform=label_transform, color_mean=color_mean, color_std=color_std)
val_loader = data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
# Load the test set as tensors
test_set = dataset(args.dataset_dir, args.testFile, mode='inference', transform=image_transform, \
label_transform=label_transform, color_mean=color_mean, color_std=color_std)
test_loader = data.DataLoader(test_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
# Get encoding between pixel valus in label images and RGB colors
class_encoding = train_set.color_encoding
# Get number of classes to predict
num_classes = len(class_encoding)
# Print information for debugging
print("Number of classes to predict:", num_classes)
print("Train dataset size:", len(train_set))
print("Validation dataset size:", len(val_set))
# Get a batch of samples to display
if args.mode.lower() == 'test':
images, labels = iter(test_loader).next()
else:
images, labels = iter(train_loader).next()
print("Image size:", images.size())
print("Label size:", labels.size())
print("Class-color encoding:", class_encoding)
# Show a batch of samples and labels
if args.imshow_batch:
print("Close the figure window to continue...")
label_to_rgb = transforms.Compose([
ext_transforms.LongTensorToRGBPIL(class_encoding),
transforms.ToTensor()
])
color_labels = utils.batch_transform(labels, label_to_rgb)
utils.imshow_batch(images, color_labels)
# Get class weights from the selected weighing technique
print("Weighing technique:", args.weighing)
# If a class weight file is provided, try loading weights from in there
class_weights = None
if args.class_weights_file:
print('Trying to load class weights from file...')
try:
class_weights = np.loadtxt(args.class_weights_file)
except Exception as e:
raise e
if class_weights is None:
print("Computing class weights...")
print("(this can take a while depending on the dataset size)")
class_weights = 0
if args.weighing.lower() == 'enet':
class_weights = enet_weighing(train_loader, num_classes)
elif args.weighing.lower() == 'mfb':
class_weights = median_freq_balancing(train_loader, num_classes)
else:
class_weights = None
if class_weights is not None:
class_weights = torch.from_numpy(class_weights).float().to(device)
# Set the weight of the unlabeled class to 0
print("Ignoring unlabeled class: ", args.ignore_unlabeled)
if args.ignore_unlabeled:
ignore_index = list(class_encoding).index('unlabeled')
class_weights[ignore_index] = 0
print("Class weights:", class_weights)
return (train_loader, val_loader,
test_loader), class_weights, class_encoding
def load_dataset(dataset):
print("\nLoading dataset...\n")
print("Selected dataset:", args.dataset)
print("Dataset directory:", args.dataset_dir)
print('Test file:', args.testFile)
print("Save directory:", args.save_dir)
image_transform = transforms.Compose(
[transforms.Resize((args.height, args.width)),
transforms.ToTensor()])
label_transform = transforms.Compose([
transforms.Resize((args.height, args.width)),
ext_transforms.PILToLongTensor()
])
# Load the test set as tensors
test_set = dataset(args.dataset_dir, args.testFile, mode='inference', transform=image_transform, \
label_transform=label_transform, color_mean=color_mean, color_std=color_std, \
load_depth=(args.arch=='rgbd'), seg_classes=args.seg_classes)
test_loader = data.DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=args.workers)
# Get encoding between pixel valus in label images and RGB colors
class_encoding = test_set.color_encoding
# Get number of classes to predict
num_classes = len(class_encoding)
# Print information for debugging
print("Number of classes to predict:", num_classes)
print("Test dataset size:", len(test_set))
# Get a batch of samples to display
if args.arch == 'rgbd':
images, labels, data_path, depth_path, label_path = iter(test_loader).next()
else:
images, labels, data_path, label_path = iter(test_loader).next()
print("Image size:", images.size())
print("Label size:", labels.size())
print("Class-color encoding:", class_encoding)
# Show a batch of samples and labels
if args.imshow_batch:
print("Close the figure window to continue...")
label_to_rgb = transforms.Compose([
ext_transforms.LongTensorToRGBPIL(class_encoding),
transforms.ToTensor()
])
color_labels = utils.batch_transform(labels, label_to_rgb)
utils.imshow_batch(images, color_labels)
# Get class weights
# If a class weight file is provided, try loading weights from in there
class_weights = None
if args.class_weights_file:
print('Trying to load class weights from file...')
try:
class_weights = np.loadtxt(args.class_weights_file)
except Exception as e:
raise e
else:
print('No class weights found...')
if class_weights is not None:
class_weights = torch.from_numpy(class_weights).float().to(device)
# Set the weight of the unlabeled class to 0
if args.ignore_unlabeled:
ignore_index = list(class_encoding).index('unlabeled')
class_weights[ignore_index] = 0
print("Class weights:", class_weights)
return test_loader, class_weights, class_encoding
def main():
assert os.path.isdir(
args.dataset_dir), "The directory \"{0}\" doesn't exist.".format(
args.dataset_dir)
# Fail fast if the saving directory doesn't exist
assert os.path.isdir(
args.save_dir), "The directory \"{0}\" doesn't exist.".format(
args.save_dir)
# Import the requested dataset
if args.dataset.lower() == 'cityscapes':
from data import Cityscapes as dataset
else:
# Should never happen...but just in case it does
raise RuntimeError("\"{0}\" is not a supported dataset.".format(
args.dataset))
print("\nLoading dataset...\n")
print("Selected dataset:", args.dataset)
print("Dataset directory:", args.dataset_dir)
print("Save directory:", args.save_dir)
image_transform = transforms.Compose(
[transforms.Resize((args.height, args.width)),
transforms.ToTensor()])
label_transform = transforms.Compose([
transforms.Resize((args.height, args.width)),
ext_transforms.PILToLongTensor()
])
# Get selected dataset
# Load the training set as tensors
train_set = dataset(args.dataset_dir,
mode='train',
max_iters=args.max_iters,
transform=image_transform,
label_transform=label_transform)
train_loader = data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
trainloader_iter = enumerate(train_loader)
# Load the validation set as tensors
val_set = dataset(args.dataset_dir,
mode='val',
max_iters=args.max_iters,
transform=image_transform,
label_transform=label_transform)
val_loader = data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
# Load the test set as tensors
test_set = dataset(args.dataset_dir,
mode='test',
max_iters=args.max_iters,
transform=image_transform,
label_transform=label_transform)
test_loader = data.DataLoader(test_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
# Get encoding between pixel valus in label images and RGB colors
class_encoding = train_set.color_encoding
# Get number of classes to predict
num_classes = len(class_encoding)
# Print information for debugging
print("Number of classes to predict:", num_classes)
print("Train dataset size:", len(train_set))
print("Validation dataset size:", len(val_set))
# Get the parameters for the validation set
if args.mode.lower() == 'test':
images, labels = iter(test_loader).next()
else:
images, labels = iter(train_loader).next()
print("Image size:", images.size())
print("Label size:", labels.size())
print("Class-color encoding:", class_encoding)
# Show a batch of samples and labels
if args.imshow_batch:
print("Close the figure window to continue...")
label_to_rgb = transforms.Compose([
ext_transforms.LongTensorToRGBPIL(class_encoding),
transforms.ToTensor()
])
color_labels = utils.batch_transform(labels, label_to_rgb)
utils.imshow_batch(images, color_labels)
# Get class weights from the selected weighing technique
print("\nTraining...\n")
num_classes = len(class_encoding)
# Define the model with the encoder and decoder from the deeplabv2
input_encoder = Encoder().to(device)
decoder_t = Decoder(num_classes).to(device)
# Define the entropy loss for the segmentation task
criterion = CrossEntropy2d()
# Set the optimizer function for model
optimizer_g = optim.SGD(itertools.chain(input_encoder.parameters(),
decoder_t.parameters()),
lr=args.learning_rate,
momentum=0.9,
weight_decay=1e-4)
optimizer_g.zero_grad()
# Evaluation metric
if args.ignore_unlabeled:
ignore_index = list(class_encoding).index('unlabeled')
else:
ignore_index = None
metric = IoU(num_classes, ignore_index=ignore_index)
# Optionally resume from a checkpoint
if args.resume:
input_encoder, decoder_t, optimizer_g, start_epoch, best_miou = utils.load_checkpoint(
input_encoder, decoder_t, optimizer_g, args.save_dir, args.name)
print("Resuming from model: Start epoch = {0} "
"| Best mean IoU = {1:.4f}".format(start_epoch, best_miou))
else:
start_epoch = 0
best_miou = 0
# Start Training
print()
metric.reset()
val = Test(input_encoder, decoder_t, val_loader, criterion, metric, device)
for i_iter in range(args.max_iters):
optimizer_g.zero_grad()
adjust_learning_rate(optimizer_g, i_iter)
_, batch_data = trainloader_iter.__next__()
inputs = batch_data[0].to(device)
labels = batch_data[1].to(device)
f_i = input_encoder(inputs)
outputs_i = decoder_t(f_i)
loss_seg = criterion(outputs_i, labels)
loss_g = loss_seg
loss_g.backward()
optimizer_g.step()
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}'.format(
i_iter, args.max_iters, loss_g))
print(">>>> [iter: {0:d}] Validation".format(i_iter))
# Validate the trained model after the weights are saved
loss, (iou, miou) = val.run_epoch(args.print_step)
print(">>>> [iter: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(i_iter, loss, miou))
if miou > best_miou:
for key, class_iou in zip(class_encoding.keys(), iou):
print("{0}: {1:.4f}".format(key, class_iou))
# Save the model if it's the best thus far
if miou > best_miou:
print("\nBest model thus far. Saving...\n")
best_miou = miou
utils.save_checkpoint(input_encoder, decoder_t, optimizer_g,
i_iter + 1, best_miou, args)
def load_dataset(dataset):
print("\n加载数据...\n")
print("选择的数据:", args.dataset)
print("Dataset 目录:", args.dataset_dir)
print("存储目录:", args.save_dir)
# 数据转换和标准化
image_transform = transforms.Compose(
[transforms.Resize((args.height, args.width)),
transforms.ToTensor()])
# 转化:PILToLongTensor,因为是label，所以不能进行标准化
label_transform = transforms.Compose([
transforms.Resize((args.height, args.width)),
ext_transforms.PILToLongTensor() # (H x W x C) 转到 (C x H x W )
])
# 获取选定的数据集
# 加载数据集作为一个tensors
train_set = dataset(args.dataset_dir,
transform=image_transform,
label_transform=label_transform)
train_loader = data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
# 加载验证集作为一个tensors
val_set = dataset(args.dataset_dir,
mode='val',
transform=image_transform,
label_transform=label_transform)
val_loader = data.DataLoader(val_set,
batch_size=3,
shuffle=True,
num_workers=args.workers)
# 加载测试集作为一个tensors
test_set = dataset(args.dataset_dir,
mode='test',
transform=image_transform,
label_transform=label_transform)
test_loader = data.DataLoader(test_set,
batch_size=3,
shuffle=True,
num_workers=args.workers)
# 获取标签图像和RGB颜色中的像素值之间的编码
class_encoding = train_set.color_encoding
# 获取需要预测的类别的数量
num_classes = len(class_encoding)
# 打印调试的信息
print("Number of classes to predict:", num_classes)
print("Train dataset size:", len(train_set))
print("Validation dataset size:", len(val_set))
# 展示一个batch的样本
if args.mode.lower() == 'test':
images, labels = iter(test_loader).next()
else:
images, labels = iter(train_loader).next()
print("Image size:", images.size())
print("Label size:", labels.size())
print("Class-color encoding:", class_encoding)
# 展示一个batch的samples和labels
if args.imshow_batch:
print("Close the figure window to continue...")
label_to_rgb = transforms.Compose([
ext_transforms.LongTensorToRGBPIL(class_encoding),
transforms.ToTensor()
])
color_labels = utils.batch_transform(labels, label_to_rgb)
utils.imshow_batch(images, color_labels)
# 获取类别的权重
print("\nWeighing technique:", args.weighing)
print("Computing class weights...")
print("(this can take a while depending on the dataset size)")
class_weights = 0
if args.weighing.lower() == 'enet':
# 传回的class_weights是一个list
class_weights = np.array([
1.44752114, 33.41317956, 43.89576605, 47.85765692, 48.3393951,
47.18958997, 40.2809274, 46.61960781, 48.28854284
])
# class_weights = enet_weighing(train_loader, num_classes)
else:
class_weights = None
if class_weights is not None:
class_weights = torch.Tensor(class_weights)
# 把没有标记的类别设置为0
# if args.ignore_unlabeled:
# ignore_index = list(class_encoding).index('unlabeled')
# class_weights[ignore_index] = 0
print("Class weights:", class_weights)
return (train_loader, val_loader,
test_loader), class_weights, class_encoding
def get_data_loaders(dataset,
train_batch_size,
test_batch_size,
val_batch_size,
single_sample=False):
print("\nLoading dataset...\n")
print("Selected dataset:", dataset)
print("Dataset directory:", args.dataset_dir)
print("Save directory:", args.save_dir)
image_transform = transforms.Compose(
[transforms.Resize((args.height, args.width)),
transforms.ToTensor()])
label_transform = transforms.Compose([
transforms.Resize((args.height, args.width), Image.NEAREST),
ext_transforms.PILToLongTensor()
])
# Get selected dataset
# Load the training set as tensors
train_set = dataset(args.dataset_dir,
transform=image_transform,
label_transform=label_transform)
if single_sample:
print("Reducing Training set to single batch of size {}".format(
train_batch_size))
train_set_reduced = torch.utils.data.Subset(
train_set, list(range(0, train_batch_size)))
train_loader = data.DataLoader(
train_set_reduced,
batch_size=train_batch_size,
shuffle=False, #Changed this
num_workers=args.workers)
else:
train_loader = data.DataLoader(
train_set,
batch_size=train_batch_size,
shuffle=False, #Changed this
num_workers=args.workers)
# Load the validation set as tensors
val_set = dataset(args.dataset_dir,
mode='val',
transform=image_transform,
label_transform=label_transform)
val_loader = data.DataLoader(val_set,
batch_size=val_batch_size,
shuffle=False,
num_workers=args.workers)
# Load the test set as tensors
test_set = dataset(args.dataset_dir,
mode='test',
transform=image_transform,
label_transform=label_transform)
test_loader = data.DataLoader(test_set,
batch_size=test_batch_size,
shuffle=False,
num_workers=args.workers)
# Get encoding between pixel valus in label images and RGB colors
class_encoding = train_set.color_encoding
# Remove the road_marking class from the CamVid dataset as it's merged
# with the road class
if args.dataset.lower() == 'camvid':
del class_encoding['road_marking']
# Get number of classes to predict
num_classes = len(class_encoding)
# Print information for debugging
print("Number of classes to predict:", num_classes)
print("Train dataset size:", len(train_set))
print("Test dataset size:", len(test_set))
print("Validation dataset size:", len(val_set))
class_weights = 0
# Get class weights from the selected weighing technique
print("Computing class weights...")
print("(this can take a while depending on the dataset size)")
class_weights = 0
class_weights = median_freq_balancing(train_loader, num_classes)
if class_weights is not None:
class_weights = torch.from_numpy(class_weights).float().to(device)
# Set the weight of the unlabeled class to 0
if args.ignore_unlabeled:
ignore_index = list(class_encoding).index('unlabeled')
class_weights[ignore_index] = 0
print("Class weights:", class_weights)
return (train_loader, val_loader,
test_loader), class_weights, class_encoding, (train_set, val_set,
test_set)
# checkpoint
optimizer = optim.Adam(model.parameters())
# Load the previoulsy saved model state to the ENet model
model = utils.load_checkpoint(model, optimizer, args.save_dir,
args.name)[0]
# print(model)
#inference(model, test_loader, w_class, class_encoding)
image_transform = transforms.Compose(
[transforms.Resize((args.height, args.width)),
transforms.ToTensor()])
label_transform = transforms.Compose([
transforms.Resize((args.height, args.width)),
ext_transforms.PILToLongTensor()
])
if args.arch.lower() == 'rgb':
image = scannet_loader(args.data_path)
image = image.unsqueeze(dim=0)
# print(image.size())
elif args.arch.lower() == 'rgbd':
image = scannet_loader_depth(args.data_path, args.depth_path)
image = image.unsqueeze(dim=0)
else:
# This condition will not occur (argparse will fail if an invalid option is specified)
raise RuntimeError(
'Invalid network architecture for dataloader specified.')
print(class_encoding)
predict(model, image, class_encoding)
