def metric_evaluation(dataTest,dataTestGT):
careRes=[]
resultImgs=[]
inputImgs=[]
model = UNet(1,depth=3)
model.load_state_dict(torch.load(MODEL_DIR+'/model.pt'))
model.eval()
## Testing on 5 images for demo purpose
for index in range(5):
im=dataTest[index]
gt=dataTestGT[0] # The ground truth is the same for all images
careResult = prediction.tiledPredict(im, model ,ps=256, overlap=48,device=device, noiseModel=None)
inputImgs.append(im)
rangePSNR=np.max(gt)-np.min(gt)
carePrior=PSNR(gt, careResult, rangePSNR )
careRes.append(carePrior)
print("Avg PSNR CARE:", np.mean(np.array(careRes) ), '+-(2SEM)',2*np.std(np.array(careRes) )/np.sqrt(float(len(careRes)) ) )
#### Logging the metric using mlflow #######
mlflow.log_metric("Avg_PSNR_CARE",np.mean(np.array(careRes)))
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
lab_tensor = im_to_tensor(lab)
lab_tensor = torch.unsqueeze(lab_tensor, 0)
img_tensor = im_trans(img)
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(
net_all=net_all,
net_vertical=net_vertical,
format='%(levelname)s: %(message)s')
args = get_args()
device = torch.device(
'cuda' if torch.cuda.is_available() and False else 'cpu')
logging.info(f'Using device {device}')
start_epoch, model_path = get_epoch_model()
net = UNet(n_channels=n_channels, n_classes=n_classes)
logging.info(
f'Network:\n'
f'\t{net.n_channels} input channels\n'
f'\t{net.n_classes} output channels (classes)\n'
f'\t{"Bilinear" if net.bilinear else "Dilated conv"} upscaling')
if model_path is not None:
net.load_state_dict(torch.load(model_path, map_location=device))
logging.info(f'Model loaded from {model_path}')
if not os.path.exists(dir_checkpoint):
os.mkdir(dir_checkpoint)
net.to(device=device)
try:
train_net(net=net,
start_epoch=start_epoch,
epochs=args.epochs,
batch_size=args.batchsize,
lr=args.lr,
device=device,
img_scale=args.scale,
(options, args) = parser.parse_args()
return options
if __name__ == '__main__':
args = get_args()
# os.environ["CUDA_VISIBLE_DEVICES"] = '0'
net = UNet(input_channels=3, nclasses=1)
# net.cuda()
# import pdb
# from torchsummary import summary
# summary(net, (3,1000,1000))
# pdb.set_trace()
if args.load:
net.load_state_dict(torch.load(args.load))
print('Model loaded from {}'.format(args.load))
if args.gpu:
if torch.cuda.device_count() > 1:
net = nn.DataParallel(net)
net.cuda()
try:
train_net(net=net,
epochs=args.epochs,
batch_size=args.batchsize,
lr=args.lr,
gpu=args.gpu,
img_scale=args.scale)
torch.save(net.state_dict(), 'model_fin.pth')