from unet.loader import ISBI_Loader, Test_Loader from unet.loader import im_to_tensor, im_trans if __name__ == "__main__": data_path = "/home/nonari/Documentos/tfgdata/tfgoct" isbi_dataset = Test_Loader(data_path, 0) train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset, batch_size=1, shuffle=True) criterion = nn.BCEWithLogitsLoss() # Select the device, if there is cuda use cuda, if not, use cpu device = torch.device('cpu' if torch.cuda.is_available() else 'cpu') # Load the network, the picture is single channel, classified as 1. net = UNet(n_channels=1, n_classes=9) # Copy the network to the deivce net.to(device=device) # Load model parameters net.load_state_dict( torch.load('/home/nonari/Descargas/best_model_v1.pth', map_location=device)) # Test mode net.eval() img = cv2.imread("/home/nonari/Documentos/tfgdata/tfgoct/img_9_20.png") lab = cv2.imread("/home/nonari/Documentos/tfgdata/tfgoct/seg_9_20.png") # Convert to grayscale img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) lab = cv2.cvtColor(lab, cv2.COLOR_BGR2GRAY) # Convert to batch as 1, channel as 1, size 512*512 array # Convert to tensor
if __name__ == "__main__": args = get_args() in_files = os.listdir(args.source) net_all = UNet(n_channels=3, n_classes=1) net_vertical = UNet(n_channels=3, n_classes=1) print(f'Loading model {all_best_model}, {vertical_best_model}') device = torch.device( 'cuda' if False and torch.cuda.is_available() else 'cpu') print(f'Using device {device}') net_all.to(device=device) net_vertical.to(device=device) net_all.load_state_dict(torch.load(all_best_model, map_location=device)) net_vertical.load_state_dict( torch.load(vertical_best_model, map_location=device)) print("Model loaded !") os.makedirs(dir_output_all, exist_ok=True) os.makedirs(dir_output_vertical, exist_ok=True) os.makedirs(dir_output, exist_ok=True) for i, file in enumerate(in_files): print(f"\nPredicting image {file} ...") image = Image.open(f'{args.source}/{file}') _all, _vertical, _combined = predict_img(