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(