loggerIoU.py
No OneTemporary
Actions

Subscribers

None

File Metadata

	import os
	import torch
	import numpy as np
	import sklearn.metrics

	class LoggerIoU:
	def __init__(self, log_dir, name, nClasses, ignoredIdx, saveBest=False,
	preproc=lambda o,t: (o,t)):
	self.log_dir=log_dir
	self.name=name
	self.log_file=os.path.join(self.log_dir,"logIou"+self.name+".txt")
	text_file = open(self.log_file, "w")
	text_file.close()

	self.nClasses=nClasses
	self.confMat=np.zeros((nClasses,nClasses))
	self.ignoredIdx=ignoredIdx
	self.saveBest=saveBest
	self.bestIoU=0
	self.preproc=preproc

	def add(self,l,output,target):
	output=output.cpu().data
	output,target=self.preproc(output,target)
	output=output.numpy()

	outputClass=np.argmax(output, axis=1)
	oc=outputClass.flatten()
	tc=target.cpu().data.numpy().flatten()
	oc_valid=oc[tc!=self.ignoredIdx]
	tc_valid=tc[tc!=self.ignoredIdx]
	self.confMat+=sklearn.metrics.confusion_matrix(tc_valid,oc_valid,labels=np.array(range(self.nClasses)))

	def logEpoch(self,net):
	sums1=np.sum(self.confMat,axis=0)
	sums2=np.sum(self.confMat,axis=1)
	dg=np.diag(self.confMat)
	iou=np.zeros(dg.shape)
	iou=np.divide(dg.astype(np.float64),(sums1+sums2-dg).astype(np.float64), out=iou, where=(dg!=0))
	text_file = open(self.log_file, "a")
	for i in range(self.nClasses):
	text_file.write('{}\t'.format(iou[i]))
	mean_iou=np.mean(iou)
	text_file.write('{}\n'.format(mean_iou))
	text_file.close()
	self.confMat.fill(0)
	if mean_iou > self.bestIoU:
	self.bestIoU=mean_iou
	if self.saveBest:
	torch.save({'state_dict': net.state_dict()},
	os.path.join(self.log_dir,
	'net_'+self.name+'_best.pth'))