def evaluate(model_test):
model_test.eval()
IOUs = 0
total_correct = 0
data_eval = singleDataset(data_rootpath,
data_name=args.data_name,
test=True)
eval_loader = torch.utils.data.DataLoader(data_eval,
batch_size=2,
shuffle=True,
collate_fn=single_collate)
print("dataset size:", len(eval_loader.dataset))
for batch_idx, (imgs, targets) in enumerate(eval_loader):
targets = targets.view(-1, 1, 4096, 64)
if is_GPU:
imgs = Variable(imgs.cuda())
targets = [
Variable(anno.cuda(), requires_grad=False) for anno in targets
]
else:
imgs = Variable(imgs)
targets = [Variable(anno, requires_grad=False) for anno in targets]
outputs = model_test(imgs)
#outputs=F.softmax(outputs,dim=1)
#occupy = (outputs.data[:,1] > 0.5) ## ByteTensor
occupy = (outputs.data > 0.5)
for idx, target in enumerate(targets):
insect, iou = eval_iou(occupy[idx], target.data)
IOUs += iou
total_correct += insect
#print 'correct num:{}'.format(total_correct)
print('the average correct rate:{}'.format(total_correct * 1.0 /
(len(eval_loader.dataset))))
print('the average iou:{}'.format(IOUs * 1.0 / (len(eval_loader.dataset))))
model_test.train()
with open(logfile, 'a') as f:
f.write('\nthe evaluate average iou:{}'.format(
IOUs * 1.0 / (len(eval_loader.dataset))))
return IOUs * 1.0 / (len(eval_loader.dataset))
metavar='PATH',
help='path to latest checkpoint (default: csg_single_model.pth)')
args = parser.parse_args()
data_rootpath = args.data
resume = './model/' + args.resume
## args.resume: just the name of checkpoint file
logfile = args.data_name + '_single_train.txt'
print('logfile name:{}'.format(logfile))
if is_GPU:
torch.cuda.set_device(args.gpu)
dataset = singleDataset(data_rootpath, data_name=args.data_name)
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=single_collate)
#model=single_UNet()
model = singleNet_verydeep()
if is_GPU:
model.cuda()
#critenrion=softmax_loss()
critenrion = CrossEntropy_loss()
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
def evaluate():
## calculate the average IOU between pred and gt
if os.path.isfile(resume):
print ('load the checkpoint file {}'.format(resume))
if is_GPU:
checkoint = torch.load(resume)
else:
checkoint = torch.load(resume, map_location=lambda storage,loc:storage)
start_epoch = checkoint['epoch']
model.load = model.load_state_dict(checkoint['model'])
#optimizer.load_state_dict(checkoint['optim'])
print ('load the resume checkpoint,train from epoch{}'.format(start_epoch))
else:
print("no resume checkpoint to load")
model.eval()
IOUs=0
total_correct=0
data_eval = singleDataset(data_rootpath,data_name=args.data_name,test=True)
eval_loader = torch.utils.data.DataLoader(data_eval,
batch_size=args.batch_size, shuffle=True, collate_fn=single_collate)
print ("dataset size:",len(eval_loader.dataset))
for batch_idx,(imgs, targets) in enumerate(eval_loader):
if is_GPU:
imgs = Variable(imgs.cuda())
targets = [Variable(anno.cuda(),requires_grad=False) for anno in targets]
else:
imgs = Variable(imgs)
targets = [Variable(anno, requires_grad=False) for anno in targets]
t1=time.time()
outputs=model(imgs)
#occupy=(outputs.data>0.5)if is_GPU else (outputs.data>0.5).type(torch.FloatTensor)
#print occupy.sum()
#print 'c1',torch.sum(outputs.data)
occupy = (outputs.data > 0.5) ## ByteTensor
#print 'c1', torch.sum(occupy)
for idx,target in enumerate(targets):
#correct+=(occupy[idx].eq(target.data)).sum()
#print target.data.type() ##(64L, 64L, 64L) FloatTensor
#print occupy[idx].type() ##(64L, 64L, 64L) ByteTensor
'''
insect=(target.data[occupy[idx]]).sum()
union=target.data.sum()+occupy[idx].sum()-insect
iou=insect*1.0/union
IOUs+=iou
'''
insect,iou=eval_iou(occupy[idx],target.data)
print ('iou:',iou)
IOUs += iou
total_correct += insect
t2=time.time()
print ('in batch{} cost{}s'.format(batch_idx,t2-t1))
#print 'correct num:{}'.format(total_correct)
print ('the average correct rate:{}'.format(total_correct*1.0/(len(eval_loader.dataset))))
print ('the average iou:{}'.format(IOUs*1.0/(len(eval_loader.dataset))))