Spaces:
Running on Zero
Running on Zero
File size: 2,484 Bytes
d43892c d8fe2bd d43892c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import os
import torch
from misc.imutils import save_image
from models.networks import *
class CDEvaluator():
def __init__(self, args):
self.n_class = args.n_class
# define G
self.net_G = define_G(args=args, gpu_ids=args.gpu_ids)
self.device = torch.device("cuda:%s" % args.gpu_ids[0]
if torch.cuda.is_available() and len(args.gpu_ids)>0
else "cpu")
print(self.device)
self.checkpoint_dir = args.checkpoint_dir
self.pred_dir = args.output_folder
os.makedirs(self.pred_dir, exist_ok=True)
def load_checkpoint(self, checkpoint_name='best_ckpt.pt'):
if os.path.exists(os.path.join(self.checkpoint_dir, checkpoint_name)):
try:
checkpoint = torch.load(os.path.join(self.checkpoint_dir, checkpoint_name),
map_location=self.device, weights_only=False)
except TypeError:
checkpoint = torch.load(os.path.join(self.checkpoint_dir, checkpoint_name),
map_location=self.device)
self.net_G.load_state_dict(checkpoint['model_G_state_dict'])
self.net_G.to(self.device)
# update some other states
self.best_val_acc = checkpoint['best_val_acc']
self.best_epoch_id = checkpoint['best_epoch_id']
else:
raise FileNotFoundError('no such checkpoint %s' % checkpoint_name)
return self.net_G
def _visualize_pred(self):
pred = torch.argmax(self.G_pred, dim=1, keepdim=True)
pred_vis = pred * 255
return pred_vis
def _forward_pass(self, batch):
self.batch = batch
img_in1 = batch['A'].to(self.device)
img_in2 = batch['B'].to(self.device)
self.shape_h = img_in1.shape[-2]
self.shape_w = img_in1.shape[-1]
self.G_pred = self.net_G(img_in1, img_in2)
return self._visualize_pred()
def eval(self):
self.net_G.eval()
def _save_predictions(self):
"""
保存模型输出结果，二分类图像
"""
preds = self._visualize_pred()
name = self.batch['name']
for i, pred in enumerate(preds):
file_name = os.path.join(
self.pred_dir, name[i].replace('.jpg', '.png'))
pred = pred[0].cpu().numpy()
save_image(pred, file_name)
|