Function predict_folder

inference.py:87–139

(rank, world_size)

resize = opt['Resize']

def predict_folder(rank, world_size):
    
    setup(rank, world_size=world_size, Master_port='12354')
    
    # DEFINE NETWORK, SCHEDULER AND OPTIMIZER
    model, _, _ = create_model(opt['network'], rank=rank)

    model = load_model(model, path_weights = opt['save']['path'])
    # create data
    PATH_IMAGES= args.inp_path
    PATH_RESULTS = './images/results'

    #create folder if it doen't exist
    not os.path.isdir(PATH_RESULTS) and os.mkdir(PATH_RESULTS)

    path_images = [os.path.join(PATH_IMAGES, path) for path in os.listdir(PATH_IMAGES) if path.endswith(('.png', '.PNG', '.jpg', '.JPEG'))]
    path_images = [file for file in path_images if not file.endswith('.csv') and not file.endswith('.txt')]
   
    model.eval()
    if rank==0:
        pbar = tqdm(total = len(path_images))
        
    for path_img in path_images:
        tensor = path_to_tensor(path_img).to(device)
        _, _, H, W = tensor.shape
        
        if resize and (H >=1500 or W>=1500):
            new_size = [int(dim//2) for dim in (H, W)]
            downsample = Resize(new_size)
        else:
            downsample = torch.nn.Identity()
        tensor = downsample(tensor)
        
        tensor = pad_tensor(tensor)

        with torch.no_grad():
            output = model(tensor, side_loss=False)
        if resize:
            upsample = Resize((H, W))
        else: upsample = torch.nn.Identity()
        output = upsample(output)
        output = torch.clamp(output, 0., 1.)
        output = output[:,:, :H, :W]
        save_tensor(output, os.path.join(PATH_RESULTS, os.path.basename(path_img)))


        pbar.update(1)
        pass

    print('Finished inference!')
    if rank == 0:
        pbar.close()   
    cleanup()

def main():
    world_size = 1