def main():
# 0. env preparation
place = paddle.fluid.CUDAPlace(0)
with fluid.dygraph.guard(place):
# 1. create model
if args.net == "pspnet":
model = PSPNet()
elif args.net == "unet":
model = UNet()
else:
raise NotImplementedError(
f"args.net: {args.net} is not Supported!")
# 2. load pretrained model
para_state_dict, _ = fluid.load_dygraph(args.model_folder + '/' +
args.net)
model.set_dict(para_state_dict)
model.eval()
# 3. read test image list
# 4. create transforms for test image, transform should be same as training
train_augmentation = Augmentation(image_size=256)
dataloader = BasicDataLoader(image_folder=args.image_folder,
image_list_file=args.image_list_file,
transform=train_augmentation,
shuffle=False)
# 5. loop over list of images
counter = 0
for im, _ in dataloader():
# 6. read image and do preprocessing
# 7. image to variable
counter += 1
im = im[np.newaxis, :, :, :]
im = to_variable(im)
# NHWC -> NCHW
im = fluid.layers.transpose(im, (0, 3, 1, 2))
pred = model(im)
pred = fluid.layers.softmax(pred, axis=1)
pred = fluid.layers.argmax(pred, axis=1)
pred = pred.numpy()
pred = np.squeeze(pred).astype('uint8')
# 8. call inference func
# if args.method == 'resize':
# pred = inference_resize(im, pred)
# elif args.method == 'sliding':
# pass
# else:
# raise Exception("Unexpected method '{}'".format(args.method))
# 9. save results
# NCHW -> NHWC
im = fluid.layers.transpose(im, (0, 2, 3, 1))
save_images(im.numpy(), pred, counter)
def main():
# 0. env preparation
with fluid.dygraph.guard():
# 1. create model
model = PSPNet(num_classes=59)
# 2. load pretrained model
pretrain_file = os.path.join(args.checkpoint_folder, f"{args.net}-Epoch-{args.num_epochs}")
if os.path.exists(pretrain_file):
state,_ = fluid.load_dygraph(pretrain_file)
model.set_dict(state)
# 3. read test image list
image_folder = "work/dummy_data"
image_list_file = "work/dummy_data/list.txt"
transform = None
if args.inf_type == 0:
transform = Augmentation0()
elif args.inf_type == 1:
transform = Augmentation1()
elif args.inf_type == 2:
transform = Augmentation2()
elif args.inf_type == 3:
transform = Augmentation3()
elif args.inf_type == 4:
transform = Augmentation4()
data = BasicDataLoader(image_folder,image_list_file,transform=transform)
# TODO: craete fluid.io.DataLoader instance
dataloader = fluid.io.DataLoader.from_generator(capacity=2, return_list=True)
# TODO: set sample generator for fluid dataloader
dataloader.set_sample_generator(data,1)
# 4. create transforms for test image, transform should be same as training
# 5. loop over list of images
for idx,data in enumerate(dataloader):
images,labels = data
#images = transform(images)
# 6. read image and do preprocessing
#print(images.shape)
i_file = f"{idx}_input.png"
input_images = (images*255.0).numpy().astype(np.uint8)
#print(np.max(input_images[0]),np.min(input_images[0]),input_images.shape)
image =input_images[0].astype(np.uint8)
cv2.imwrite(i_file,image)
images = fluid.layers.transpose(images,(0,3,1,2))
# 7. image to variable
images = to_variable(images)
# 8. call inference func
preds = model(images).numpy()
pred = np.squeeze(preds)
result = np.argmax(pred,axis=0)
#print(preds.shape,result.shape)
# 9. save results
result = colorize(result,[1,2,3])
o_file = f"{idx}_pred.png"
cv2.imwrite(o_file,result)
save_blend_image(i_file,o_file)
