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)