def run_pipeline(root_dir, model_path, img_size, batch_size, use_gpu):
images_path = os.path.join(root_dir, "images")
masks_path = os.path.join(root_dir, "masks")
outputs_path = os.path.join(root_dir, "outputs")
sizes = []
file_names = []
for f in os.listdir(images_path):
im = Image.open(os.path.join(images_path, f))
sizes.append(im.size)
file_names.append(f)
model = UNet(num_channels=1, num_classes=2)
use_gpu = use_gpu and torch.cuda.is_available()
if use_gpu:
model.cuda()
model.load_state_dict(torch.load(model_path, map_location='cpu'))
test_dataset = TestDataset(images_path,
im_size=[img_size, img_size],
transform=tr.ToTensor())
test_loader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=1)
print("Test Data : ", len(test_loader.dataset))
start = timer()
predict(model, test_loader, batch_size, sizes, file_names, masks_path,
use_gpu)
end = timer()
print("Prediction completed in {:0.2f}s".format(end - start))
generate_bbox(images_path, masks_path, outputs_path)
end2 = timer()
print("Bbox generation completed in {:0.2f}s".format(end2 - end))
def setup(opts):
model = UNet(backbone=opts["backbone"], num_classes=2)
if torch.cuda.is_available():
print("Using CUDA")
trained_dict = torch.load(opts["checkpoint"])['state_dict']
model.load_state_dict(trained_dict, strict=False)
model.cuda()
else:
print("Using CPU")
trained_dict = torch.load(opts["checkpoint"], map_location="cpu")['state_dict']
model.load_state_dict(trained_dict, strict=False)
return model
def get_model(model_path, model_type, num_classes):
"""
:param model_path:
:param model_type: 'UNet', 'UNet16', 'UNet11', 'LinkNet34',
:param problem_type: 'binary', 'parts', 'instruments'
:return:
"""
if model_type == 'UNet':
model = UNet(num_classes=num_classes)
else:
model_name = model_list[model_type]
model = model_name(num_classes=num_classes)
# print(model)
state = torch.load(str(model_path))
state = {
key.replace('module.', ''): value
for key, value in state['model'].items()
}
model.load_state_dict(state)
if torch.cuda.is_available():
return model.cuda()
model.eval()
return model
def get_model(model_path, model_type):
"""
:param model_path:
:param model_type: 'UNet', 'UNet11', 'UNet16', 'AlbuNet34'
:return:
"""
num_classes = 1
if model_type == 'UNet11':
model = UNet11(num_classes=num_classes)
elif model_type == 'UNet16':
model = UNet16(num_classes=num_classes)
elif model_type == 'AlbuNet34':
model = AlbuNet34(num_classes=num_classes)
elif model_type == 'UNet':
model = UNet(num_classes=num_classes)
else:
model = UNet(num_classes=num_classes)
state = torch.load(str(model_path))
state = {
key.replace('module.', ''): value
for key, value in state['model'].items()
}
model.load_state_dict(state)
if torch.cuda.is_available():
return model.cuda()
model.eval()
return model
def build(config):
global train_dataloader
global val_dataloader
global test_dataloader
global model
global loss_func
global optimizer
# ========= Build Data ==============
if base_config['dataset'] == 'kaggle':
from data import build_kaggle_dataset
train_dataloader, val_dataloader, test_dataloader = build_kaggle_dataset(
base_config)
elif base_config['dataset'] == 'drive':
from data import build_drive_dataset
train_dataloader, val_dataloader, test_dataloader = build_drive_dataset(
base_config)
else:
_logger.error('{} dataset is not supported now'.format(
base_config['dataset']))
# ======== Build Model
if config['model'] == 'resnet101':
from torchvision.models import resnet101
model = resnet101(num_classes=base_config['n_classes'])
elif config['model'] == 'resnext101':
from torchvision.models import resnext101_32x8d
model = resnext101_32x8d(num_classes=base_config['n_classes'])
elif config['model'] == 'densenet':
from torchvision.models import densenet121
model = densenet121(num_classes=base_config['n_classes'])
elif config['model'] == 'unet':
from models import UNet
model = UNet(num_classes=base_config['n_classes'])
else:
_logger.error('{} model is not supported'.format(config['model']))
model = torch.nn.DataParallel(model.cuda())
# Build optimizer
if base_config['loss'] == 'ce':
loss_func = torch.nn.CrossEntropyLoss().cuda()
elif base_config['loss'] == 'bce':
loss_func = torch.nn.BCELoss().cuda()
elif base_config['loss'] == 'MSE':
loss_func = torch.nn.MSELoss().cuda()
else:
_logger.error('{} loss is not supported'.format(config['loss']))
if config['optimizer'] == 'SGD':
optimizer = torch.optim.SGD(model.parameters(),
lr=config['lr'],
momentum=0.9,
weight_decay=5e-4)
if config['optimizer'] == 'Adadelta':
optimizer = torch.optim.Adadelta(model.parameters(), lr=config['lr'])
if config['optimizer'] == 'Adagrad':
optimizer = torch.optim.Adagrad(model.parameters(), lr=config['lr'])
if config['optimizer'] == 'Adam':
optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])
if config['optimizer'] == 'Adamax':
optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])
def get_model(model_path, model_type='unet11', problem_type='binary'):
"""
:param model_path:
:param model_type: 'UNet', 'UNet16', 'UNet11', 'LinkNet34'
:param problem_type: 'binary', 'parts', 'instruments'
:return:
"""
if problem_type == 'binary':
num_classes = 1
elif problem_type == 'parts':
num_classes = 4
elif problem_type == 'instruments':
num_classes = 8
if model_type == 'UNet16':
model = UNet16(num_classes=num_classes)
elif model_type == 'UNet11':
model = UNet11(num_classes=num_classes)
elif model_type == 'LinkNet34':
model = LinkNet34(num_classes=num_classes)
elif model_type == 'UNet':
model = UNet(num_classes=num_classes)
elif model_type == 'DLinkNet':
model = D_LinkNet34(num_classes=num_classes, pretrained=True)
state = torch.load(str(model_path))
state = {key.replace('module.', ''): value for key, value in state['model'].items()}
model.load_state_dict(state)
if torch.cuda.is_available():
return model.cuda()
model.eval()
return model
def start():
parser = argparse.ArgumentParser(
description='UNet + BDCLSTM for BraTS Dataset')
parser.add_argument('--batch-size',
type=int,
default=4,
metavar='N',
help='input batch size for training (default: 4)')
parser.add_argument('--test-batch-size',
type=int,
default=4,
metavar='N',
help='input batch size for testing (default: 4)')
parser.add_argument('--train',
action='store_true',
default=False,
help='Argument to train model (default: False)')
parser.add_argument('--epochs',
type=int,
default=2,
metavar='N',
help='number of epochs to train (default: 10)')
parser.add_argument('--lr',
type=float,
default=0.001,
metavar='LR',
help='learning rate (default: 0.01)')
parser.add_argument('--cuda',
action='store_true',
default=False,
help='enables CUDA training (default: False)')
parser.add_argument('--log-interval',
type=int,
default=1,
metavar='N',
help='batches to wait before logging training status')
parser.add_argument('--size',
type=int,
default=128,
metavar='N',
help='imsize')
parser.add_argument('--load',
type=str,
default=None,
metavar='str',
help='weight file to load (default: None)')
parser.add_argument('--data',
type=str,
default='./Data/',
metavar='str',
help='folder that contains data')
parser.add_argument('--save',
type=str,
default='OutMasks',
metavar='str',
help='Identifier to save npy arrays with')
parser.add_argument('--modality',
type=str,
default='flair',
metavar='str',
help='Modality to use for training (default: flair)')
parser.add_argument('--optimizer',
type=str,
default='SGD',
metavar='str',
help='Optimizer (default: SGD)')
args = parser.parse_args()
args.cuda = args.cuda and torch.cuda.is_available()
DATA_FOLDER = args.data
# %% Loading in the model
# Binary
# model = UNet(num_channels=1, num_classes=2)
# Multiclass
model = UNet(num_channels=1, num_classes=3)
if args.cuda:
model.cuda()
if args.optimizer == 'SGD':
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.99)
if args.optimizer == 'ADAM':
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
betas=(args.beta1, args.beta2))
# Defining Loss Function
criterion = DICELossMultiClass()
if args.train:
# %% Loading in the Dataset
full_dataset = BraTSDatasetUnet(DATA_FOLDER,
im_size=[args.size, args.size],
transform=tr.ToTensor())
#dset_test = BraTSDatasetUnet(DATA_FOLDER, train=False,
# keywords=[args.modality], im_size=[args.size,args.size], transform=tr.ToTensor())
train_size = int(0.9 * len(full_dataset))
test_size = len(full_dataset) - train_size
train_dataset, validation_dataset = torch.utils.data.random_split(
full_dataset, [train_size, test_size])
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=1)
validation_loader = DataLoader(validation_dataset,
batch_size=args.test_batch_size,
shuffle=False,
num_workers=1)
#test_loader = DataLoader(full_dataset, batch_size=args.test_batch_size, shuffle=False, num_workers=1)
print("Training Data : ", len(train_loader.dataset))
print("Validaion Data : ", len(validation_loader.dataset))
#print("Test Data : ", len(test_loader.dataset))
loss_list = []
start = timer()
for i in tqdm(range(args.epochs)):
train(model, i, loss_list, train_loader, optimizer, criterion,
args)
test(model, validation_loader, criterion, args, validation=True)
end = timer()
print("Training completed in {:0.2f}s".format(end - start))
plt.plot(loss_list)
plt.title("UNet bs={}, ep={}, lr={}".format(args.batch_size,
args.epochs, args.lr))
plt.xlabel("Number of iterations")
plt.ylabel("Average DICE loss per batch")
plt.savefig("./plots/{}-UNet_Loss_bs={}_ep={}_lr={}.png".format(
args.save, args.batch_size, args.epochs, args.lr))
np.save(
'./npy-files/loss-files/{}-UNet_Loss_bs={}_ep={}_lr={}.npy'.format(
args.save, args.batch_size, args.epochs, args.lr),
np.asarray(loss_list))
print("Testing Validation")
test(model, validation_loader, criterion, args, save_output=True)
torch.save(
model.state_dict(),
'unet-multiclass-model-{}-{}-{}'.format(args.batch_size,
args.epochs, args.lr))
print("Testing PDF images")
test_dataset = TestDataset('./pdf_data/',
im_size=[args.size, args.size],
transform=tr.ToTensor())
test_loader = DataLoader(test_dataset,
batch_size=args.test_batch_size,
shuffle=False,
num_workers=1)
print("Test Data : ", len(test_loader.dataset))
test_only(model, test_loader, criterion, args)
elif args.load is not None:
test_dataset = TestDataset(DATA_FOLDER,
im_size=[args.size, args.size],
transform=tr.ToTensor())
test_loader = DataLoader(test_dataset,
batch_size=args.test_batch_size,
shuffle=False,
num_workers=1)
print("Test Data : ", len(test_loader.dataset))
model.load_state_dict(torch.load(args.load))
test_only(model, test_loader, criterion, args)
# Alpha transperency
else:
COLOR1 = [255, 0, 0]
COLOR2 = [0, 0, 255]
#------------------------------------------------------------------------------
# Create model and load weights
#------------------------------------------------------------------------------
model = UNet(
backbone="mobilenetv2",
num_classes=2,
pretrained_backbone=None
)
if args.use_cuda:
model = model.cuda()
trained_dict = torch.load(args.checkpoint, map_location="cpu")['state_dict']
model.load_state_dict(trained_dict, strict=False)
model.eval()
#------------------------------------------------------------------------------
# Predict frames
#------------------------------------------------------------------------------
i = 0
while(cap.isOpened()):
# Read frame from camera
start_time = time()
_, frame = cap.read()
image = cv2.transpose(frame[...,::-1])
h, w = image.shape[:2]
def train():
t.cuda.set_device(1)
# n_channels：医学影像为一通道灰度图 n_classes：二分类
net = UNet(n_channels=1, n_classes=1)
optimizer = t.optim.SGD(net.parameters(),
lr=opt.learning_rate,
momentum=0.9,
weight_decay=0.0005)
criterion = t.nn.BCELoss() # 二进制交叉熵(适合mask占据图像面积较大的场景)
start_epoch = 0
if opt.load_model_path:
checkpoint = t.load(opt.load_model_path)
# 加载多GPU模型参数到 单模型上
state_dict = checkpoint['net']
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
net.load_state_dict(new_state_dict) # 加载模型
optimizer.load_state_dict(checkpoint['optimizer']) # 加载优化器
start_epoch = checkpoint['epoch'] # 加载训练批次
# 学习率每当到达milestones值则更新参数
if start_epoch == 0:
scheduler = t.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=opt.milestones,
gamma=0.1,
last_epoch=-1) # 默认为-1
print('从头训练 ，学习率为{}'.format(optimizer.param_groups[0]['lr']))
else:
scheduler = t.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=opt.milestones,
gamma=0.1,
last_epoch=start_epoch)
print('加载预训练模型{}并从{}轮开始训练,学习率为{}'.format(
opt.load_model_path, start_epoch, optimizer.param_groups[0]['lr']))
# 网络转移到GPU上
if opt.use_gpu:
net = t.nn.DataParallel(net, device_ids=opt.device_ids) # 模型转为GPU并行
net.cuda()
cudnn.benchmark = True
# 定义可视化对象
vis = Visualizer(opt.env)
train_data = NodeDataSet(train=True)
val_data = NodeDataSet(val=True)
test_data = NodeDataSet(test=True)
# 数据集加载器
train_dataloader = DataLoader(train_data,
opt.batch_size,
shuffle=True,
num_workers=opt.num_workers)
val_dataloader = DataLoader(val_data,
opt.batch_size,
shuffle=True,
num_workers=opt.num_workers)
test_dataloader = DataLoader(test_data,
opt.test_batch_size,
shuffle=False,
num_workers=opt.num_workers)
for epoch in range(opt.max_epoch - start_epoch):
print('开始 epoch {}/{}.'.format(start_epoch + epoch + 1, opt.max_epoch))
epoch_loss = 0
# 每轮判断是否更新学习率
scheduler.step()
# 迭代数据集加载器
for ii, (img, mask) in enumerate(
train_dataloader): # pytorch0.4写法，不再将tensor封装为Variable
# 将数据转到GPU
if opt.use_gpu:
img = img.cuda()
true_masks = mask.cuda()
masks_pred = net(img)
# 经过sigmoid
masks_probs = t.sigmoid(masks_pred)
# 损失 = 二进制交叉熵损失 + dice损失
loss = criterion(masks_probs.view(-1), true_masks.view(-1))
# 加入dice损失
if opt.use_dice_loss:
loss += dice_loss(masks_probs, true_masks)
epoch_loss += loss.item()
if ii % 2 == 0:
vis.plot('训练集loss', loss.item())
# 优化器梯度清零
optimizer.zero_grad()
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
# 当前时刻的一些信息
vis.log("epoch:{epoch},lr:{lr},loss:{loss}".format(
epoch=epoch, loss=loss.item(), lr=optimizer.param_groups[0]['lr']))
vis.plot('每轮epoch的loss均值', epoch_loss / ii)
# 保存模型、优化器、当前轮次等
state = {
'net': net.state_dict(),
'optimizer': optimizer.state_dict(),
'epoch': epoch
}
t.save(state, opt.checkpoint_root + '{}_unet.pth'.format(epoch))
# ============验证===================
net.eval()
# 评价函数：Dice系数 Dice距离用于度量两个集合的相似性
tot = 0
for jj, (img_val, mask_val) in enumerate(val_dataloader):
img_val = img_val
true_mask_val = mask_val
if opt.use_gpu:
img_val = img_val.cuda()
true_mask_val = true_mask_val.cuda()
mask_pred = net(img_val)
mask_pred = (t.sigmoid(mask_pred) > 0.5).float() # 阈值为0.5
# 评价函数：Dice系数 Dice距离用于度量两个集合的相似性
tot += dice_loss(mask_pred, true_mask_val).item()
val_dice = tot / jj
vis.plot('验证集 Dice损失', val_dice)
# ============验证召回率===================
# 每10轮验证一次测试集召回率
if epoch % 10 == 0:
result_test = []
for kk, (img_test, mask_test) in enumerate(test_dataloader):
# 测试 unet分割能力，故 不使用真值mask
if opt.use_gpu:
img_test = img_test.cuda()
mask_pred_test = net(img_test) # [1,1,512,512]
probs = t.sigmoid(mask_pred_test).squeeze().squeeze().cpu(
).detach().numpy() # [512,512]
mask = probs > opt.out_threshold
result_test.append(mask)
# 得到 测试集所有预测掩码,计算二维召回率
vis.plot('测试集二维召回率', getRecall(result_test).getResult())
net.train()
def predict():
net = UNet(n_channels=1, n_classes=1)
net.eval()
# 将多GPU模型加载为CPU模型
if opt.load_model_path:
checkpoint = t.load(opt.load_model_path)
state_dict = checkpoint['net']
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
net.load_state_dict(new_state_dict) # 加载模型
print('加载预训练模型{}'.format(opt.load_model_path))
if opt.use_gpu:
net.cuda()
test_data = NodeDataSet(test=True)
test_dataloader = DataLoader(test_data,
opt.test_batch_size,
shuffle=False,
num_workers=opt.num_workers)
for ii, full_img in enumerate(test_dataloader):
img_test = full_img[0][0].unsqueeze(
0) # 第一个[0] 取 原图像的一个batch，第二个[0]指batch为1
if opt.use_gpu:
img_test = img_test.cuda()
with t.no_grad(): # pytorch0.4版本写法
output = net(img_test)
probs = t.sigmoid(output).squeeze(0)
full_mask = probs.squeeze().cpu().numpy()
# ===========================================下面方法可能未考虑 一通道图像
# if opt.use_dense_crf:
# full_mask = dense_crf(np.array(full_img).astype(np.uint8), full_mask)
mask = full_mask > opt.out_threshold # 预测mask值都太小,最大0.01
# # 可视化1
# plt.imsave(opt.save_test_dir+str(10000+ii)+'full_img.jpg', full_img[0][0].squeeze(0),cmap = cm.gray) #保存原图
# plt.imsave(opt.save_test_dir+str(10000+ii)+'mask.jpg', mask,cmap = cm.gray) #保存mask
# plt.imsave(opt.save_test_dir+str(10000+ii)+'full_mask.jpg', full_img[0][0].squeeze(0).squeeze(0).numpy() * mask,cmap = cm.gray) #保存mask之后的原图
# 可视化2
# # 多子图显示原图和mask
# plt.subplot(1,3,1)
# plt.title('origin')
# plt.imshow(full_img[0][0].squeeze(0),cmap='Greys_r')
#
# plt.subplot(1, 3, 2)
# plt.title('mask')
# plt.imshow(mask,cmap='Greys_r')
#
# plt.subplot(1, 3, 3)
# plt.title('origin_after_mask')
# plt.imshow( full_img[0][0].squeeze(0).squeeze(0).numpy() * mask,cmap='Greys_r')
#
# plt.show()
# 保存mask为npy
np.save('/home/bobo/data/test/test8/' + str(10000 + ii) + '_mask.npy',
mask)
print('测试完毕')
class Anonymizer:
@classmethod
def check_for_numpy(cls, tensor):
if isinstance(tensor, torch.Tensor):
tensor = tensor.detach().cpu().numpy()
return tensor
@classmethod
def get_number_of_batches(cls, image_paths, batch_size):
batches = len(image_paths) / batch_size
if not batches.is_integer():
batches = math.floor(batches) + 1
return int(batches)
@classmethod
def apply_mask(cls, image, mask):
return np.multiply(image, mask)
@classmethod
def anonymize_image(cls, image, mask):
image = Anonymizer.check_for_numpy(image)
image = image.reshape(image.shape[-2:])
image = np.float32(image)
mask = Anonymizer.check_for_numpy(mask)
mask = mask.reshape(mask.shape[-2:])
mask = np.uint8(mask)
im = Anonymizer.apply_mask(image, mask)
cv2.imwrite("sanity_mask.jpg", 255 * mask)
cv2.imwrite("sanity_image.jpg", image)
cv2.imwrite("sanity_join.jpg", im)
mask = Editor.invert_mask(mask)
mask = np.uint8(mask)
im = cv2.inpaint(im, mask, 10, cv2.INPAINT_TELEA)
cv2.imwrite("sanity_anon.jpg", im)
return im
def __init__(self, batch_size, image_paths, write_path, state_dict):
self.batch_size = batch_size
self.image_paths = glob.glob(image_paths)
self.batches = Anonymizer.get_number_of_batches(
self.image_paths, self.batch_size)
self.write_path = write_path
self.model = UNet()
self.state_dict = state_dict
def process_batch(self, batch):
# Grab a batch, shuffled according to the provided seed. Note that
# i-th image: samples[i][0], i-th mask: samples[i][1]
samples = Loader.get_batch(self.image_paths, self.batch_size, batch,
None)
samples.astype(float)
# Cast samples into torch.FloatTensor for interaction with U-Net
samples = torch.from_numpy(samples)
samples = samples.float()
# Cast into a CUDA tensor, if GPUs are available
if torch.cuda.is_available():
samples = samples.cuda()
# Isolate images and their masks
samples_images = samples[:, 0]
samples_masks = samples[:, 1]
# Reshape for interaction with U-Net
samples_images = samples_images.unsqueeze(1)
source = samples_images
samples_masks = samples_masks.unsqueeze(1)
# Run inputs through the model
output = self.model(samples_images)
# Clamp the target for proper interaction with BCELoss
target = torch.clamp(samples_masks, min=0, max=1)
del samples
return source, output, target
def anonymize(self):
if not os.path.isdir(self.write_path):
print("Making output directory")
os.mkdir(self.write_path)
count = 0
for batch in range(self.batches):
source, output, target = self.process_batch(batch)
source = Anonymizer.check_for_numpy(source)
source = source.reshape(source.shape[-2:])
source = np.float32(source)
binary_mask = Editor.make_binary_mask_from_torch(
output[0, :, :, :], 1.0)
inverted_binary_mask = Editor.invert_mask(
binary_mask) # Now a numpy array instead of torch tensor
anonymized_image = Anonymizer.anonymize_image(
source, inverted_binary_mask)
cv2.imwrite(self.write_path + "/orig_" + str(count) + ".jpg",
source)
cv2.imwrite(self.write_path + "/anon_" + str(count) + ".jpg",
anonymized_image)
count += 1
del batch, target, output, binary_mask, anonymized_image
def set_cuda(self):
if torch.cuda.is_available():
self.model = self.model.cuda()
def set_weights(self):
if torch.cuda.is_available():
buffered_state_dict = torch.load("weights/" + self.state_dict)
else:
buffered_state_dict = torch.load(
"weights/" + self.state_dict,
map_location=lambda storage, loc: storage)
self.model.load_state_dict(buffered_state_dict)
self.model.eval()
def main():
use_cuda = torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')
unet = UNet(input_dim, label_dim)
unet.load_state_dict(
torch.load('./checkpoints_1_19/checkpoint_30.pth', map_location='cpu'))
unet.eval()
if use_cuda:
unet.cuda()
criterion = nn.MSELoss()
"""
if sys.platform.startswith('win'):
num_workers = 0 # 0表示不用额外的进程来加速读取数据
else:
num_workers = 4
# 读取训练数据集
batch_size = 512
#准备数据
mnist_test_dataset_with_noise = MyMnistDataSet.MyMnistDataSet(root_dir='./mnist_dataset_noise',
label_root_dir='./mnist_dataset', type_name='test',
transform=transforms.ToTensor())
test_data_loader_with_noise = torch.utils.data.DataLoader(mnist_test_dataset_with_noise, batch_size, shuffle=False,
num_workers=num_workers)
#遍历数据已有模型进行reference
test_loss_sum, batch_count, start_time = 0.0, 0, time.time()
for X, y in test_data_loader_with_noise:
X = X.to(device)
y = y.to(device)
y_hat = unet(X)
l = criterion(y_hat, y)
test_loss_sum += l.cpu().item()
batch_count += 1
print('predict: batch_cout %d, test loss %.4f, time %.1f sec'
% (batch_count, test_loss_sum / batch_count, time.time() - start_time))
"""
transform = transforms.Compose([
transforms.CenterCrop(256),
transforms.ToTensor(),
])
dataset = SpectralDataSet(
root_dir=
'/mnt/liguanlin/DataSets/lowlight_hyperspectral_datasets/band_splited_dataset',
type_name='test',
transform=transform)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
count = 0
mean_psnr = 0
mean_ssim = 0
total_psnr = 0
total_ssim = 0
goundtruth_total_pnsr = 0
goundtruth_total_ssim = 0
for real, labels in tqdm(dataloader):
#print('real.shape', real.shape)
#print('labels.shape', labels.shape)
cur_batch_size = len(real)
# Flatten the image
real = real.to(device)
labels = labels.to(device)
pred = unet(real)
print(pred.shape)
print(real.shape)
save_image(pred, 'pred.png', nrow=2)
save_image(labels, 'labels.png', nrow=2)
pred_numpy = pred.detach().cpu().numpy()
label_numpy = labels.detach().cpu().numpy()
origin_numpy = real.detach().cpu().numpy()
print(pred_numpy.shape)
for i in range(pred_numpy.shape[0]):
numpy_img = pred_numpy[i].reshape(
(pred_numpy.shape[2], pred_numpy.shape[3]))
pred_numpy_img = numpy_img * 1023
pred_numpy_img = pred_numpy_img.astype(np.int16)
image = Image.fromarray(pred_numpy_img)
iamge_name = "./test_results/pred/" + str(i) + ".png"
image.save(iamge_name)
#pred_numpy_img_8bit = numpy_img * 255
label_numpy_img = label_numpy[i].reshape(
(label_numpy.shape[2], label_numpy.shape[3]))
label_numpy_img = label_numpy_img * 1023
label_numpy_img = label_numpy_img.astype(np.int16)
label_image = Image.fromarray(label_numpy_img)
label_iamge_name = "./test_results/label/" + str(i) + ".png"
label_image.save(label_iamge_name)
origin_numpy_img = origin_numpy[i].reshape(
(origin_numpy.shape[2], origin_numpy.shape[3]))
origin_numpy_img = origin_numpy_img * 1023
origin_numpy_img = origin_numpy_img.astype(np.int16)
origin_image = Image.fromarray(origin_numpy_img)
origin_iamge_name = "./test_results/original/" + str(i) + ".png"
origin_image.save(origin_iamge_name)
count += 1
total_psnr += caculate_psnr_16bit(pred_numpy_img, origin_numpy_img)
total_ssim += caculate_ssim_16bit(pred_numpy_img, label_numpy_img)
goundtruth_total_pnsr += caculate_psnr_16bit(
label_numpy_img, origin_numpy_img)
goundtruth_total_ssim += caculate_ssim_16bit(
label_numpy_img, pred_numpy_img)
if (count == 4):
break
mean_psnr = total_psnr / count
mean_ssim = total_ssim / count
print("count = ", count)
print("mean_psnr = ", mean_psnr)
print("mean_ssim = ", mean_ssim)
gound_truth_mean_psnr = goundtruth_total_pnsr / count
gound_truth_mean_ssim = goundtruth_total_ssim / count
print("gound_truth_mean_psnr = ", gound_truth_mean_psnr)
print("gound_truth_mean_ssim = ", gound_truth_mean_ssim)
class Runner(object):
def __init__(self,
hparams,
train_size: int,
class_weight: Optional[Tensor] = None):
# model, criterion, and prediction
self.model = UNet(ch_in=2, ch_out=1, **hparams.model)
self.sigmoid = torch.nn.Sigmoid()
self.criterion = torch.nn.BCEWithLogitsLoss(reduction='none')
self.class_weight = class_weight
# for prediction
self.frame2time = hparams.hop_size / hparams.sample_rate
self.T_6s = round(6 / self.frame2time) - 1
self.T_12s = round(12 / self.frame2time) - 1
self.metrics = ('precision', 'recall', 'F1')
# optimizer and scheduler
self.optimizer = AdamW(
self.model.parameters(),
lr=hparams.learning_rate,
weight_decay=hparams.weight_decay,
)
self.scheduler = CosineLRWithRestarts(self.optimizer,
batch_size=hparams.batch_size,
epoch_size=train_size,
**hparams.scheduler)
self.scheduler.step()
self.f1_last_restart = -1
# device
device_for_summary = self._init_device(hparams.device,
hparams.out_device)
# summary
self.writer = SummaryWriter(logdir=hparams.logdir)
path_summary = Path(self.writer.logdir, 'summary.txt')
if not path_summary.exists():
print_to_file(path_summary, summary,
(self.model,
(2, 128, 16 * hparams.model['stride'][1]**4)),
dict(device=device_for_summary))
# save hyperparameters
path_hparam = Path(self.writer.logdir, 'hparams.txt')
if not path_hparam.exists():
with path_hparam.open('w') as f:
for var in vars(hparams):
value = getattr(hparams, var)
print(f'{var}: {value}', file=f)
def _init_device(self, device, out_device) -> str:
if device == 'cpu':
self.device = torch.device('cpu')
self.out_device = torch.device('cpu')
self.str_device = 'cpu'
return 'cpu'
# device type
if type(device) == int:
device = [device]
elif type(device) == str:
device = [int(device[-1])]
else: # sequence of devices
if type(device[0]) == int:
device = device
else:
device = [int(d[-1]) for d in device]
# out_device type
if type(out_device) == int:
out_device = torch.device(f'cuda:{out_device}')
else:
out_device = torch.device(out_device)
self.device = torch.device(f'cuda:{device[0]}')
self.out_device = out_device
if len(device) > 1:
self.model = nn.DataParallel(self.model,
device_ids=device,
output_device=out_device)
self.str_device = ', '.join([f'cuda:{d}' for d in device])
else:
self.str_device = str(self.device)
self.model.cuda(device[0])
self.criterion.cuda(out_device)
if self.sigmoid:
self.sigmoid.cuda(device[0])
torch.cuda.set_device(device[0])
return 'cuda'
def calc_loss(self, y: Tensor, out: Tensor, Ts: Union[List[int],
int]) -> Tensor:
"""
:param y: (B, T) or (T,)
:param out: (B, T) or (T,)
:param Ts: length B list or int
:return:
"""
assert self.class_weight is not None
weight = (y > 0).float() * self.class_weight[1].item()
weight += (y == 0).float() * self.class_weight[0].item()
if y.dim() == 1: # if batch_size == 1
y = (y, )
out = (out, )
weight = (weight, )
Ts = (Ts, )
loss = torch.zeros(1, device=self.out_device)
for ii, T in enumerate(Ts):
loss_no_red = self.criterion(out[ii:ii + 1, ..., :T],
y[ii:ii + 1, :T])
loss += (loss_no_red * weight[ii:ii + 1, :T]).sum() / T
return loss
def predict(self, out_np: ndarray, Ts: Union[List[int], int]) \
-> Tuple[List[ndarray], List]:
""" peak-picking prediction
:param out_np: (B, T) or (T,)
:param Ts: length B list or int
:return: boundaries, thresholds
boundaries: length B list of boundary interval ndarrays
thresholds: length B list of threshold values
"""
if out_np.ndim == 1: # if batch_size == 1
out_np = (out_np, )
Ts = (Ts, )
boundaries = []
thresholds = []
for item, T in zip(out_np, Ts):
candid_idx = []
for idx in range(1, T - 1):
i_first = max(idx - self.T_6s, 0)
i_last = min(idx + self.T_6s + 1, T)
if item[idx] >= np.amax(item[i_first:i_last]):
candid_idx.append(idx)
boundary_idx = []
threshold = np.mean(item[candid_idx])
for idx in candid_idx:
if item[idx] > threshold:
boundary_idx.append(idx)
boundary_interval = np.array(
[[0] + boundary_idx, boundary_idx + [T]], dtype=np.float64).T
boundary_interval *= self.frame2time
boundaries.append(boundary_interval)
thresholds.append(threshold)
return boundaries, thresholds
@staticmethod
def evaluate(reference: Union[List[ndarray], ndarray],
prediction: Union[List[ndarray], ndarray]):
"""
:param reference: length B list of ndarray or just ndarray
:param prediction: length B list of ndarray or just ndarray
:return: (3,) ndarray
"""
if isinstance(reference, ndarray): # if batch_size == 1
reference = (reference, )
result = np.zeros(3)
for item_truth, item_pred in zip(reference, prediction):
mir_result = mir_eval.segment.detection(item_truth,
item_pred,
trim=True)
result += np.array(mir_result)
return result
# Running model for train, test and validation.
def run(self, dataloader, mode: str, epoch: int):
self.model.train() if mode == 'train' else self.model.eval()
if mode == 'test':
state_dict = torch.load(Path(self.writer.logdir, f'{epoch}.pt'))
if isinstance(self.model, nn.DataParallel):
self.model.module.load_state_dict(state_dict)
else:
self.model.load_state_dict(state_dict)
path_test_result = Path(self.writer.logdir, f'test_{epoch}')
os.makedirs(path_test_result, exist_ok=True)
else:
path_test_result = None
avg_loss = 0.
avg_eval = 0.
all_thresholds = dict()
print()
pbar = tqdm(dataloader,
desc=f'{mode} {epoch:3d}',
postfix='-',
dynamic_ncols=True)
for i_batch, (x, y, intervals, Ts, ids) in enumerate(pbar):
# data
n_batch = len(Ts) if hasattr(Ts, 'len') else 1
x = x.to(self.device) # B, C, F, T
x = dataloader.dataset.normalization.normalize_(x)
y = y.to(self.out_device) # B, T
# forward
out = self.model(x) # B, C, 1, T
out = out[..., 0, 0, :] # B, T
# loss
if mode != 'test':
if mode == 'valid':
with torch.autograd.detect_anomaly():
loss = self.calc_loss(y, out, Ts)
else:
loss = self.calc_loss(y, out, Ts)
else:
loss = 0
out_np = self.sigmoid(out).detach().cpu().numpy()
prediction, thresholds = self.predict(out_np, Ts)
eval_result = self.evaluate(intervals, prediction)
if mode == 'train':
# backward
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
self.scheduler.batch_step()
loss = loss.item()
elif mode == 'valid':
loss = loss.item()
if i_batch == 0: # save only the 0-th data
id_0, T_0 = ids[0], Ts[0]
out_np_0 = out_np[0, :T_0]
pred_0, truth_0 = prediction[0][1:, 0], intervals[0][1:, 0]
t_axis = np.arange(T_0) * self.frame2time
fig = draw_lineplot(t_axis, out_np_0, pred_0, truth_0,
id_0)
self.writer.add_figure(f'{mode}/out', fig, epoch)
np.save(Path(self.writer.logdir, f'{id_0}_{epoch}.npy'),
out_np_0)
np.save(
Path(self.writer.logdir, f'{id_0}_{epoch}_pred.npy'),
pred_0)
if epoch == 0:
np.save(Path(self.writer.logdir, f'{id_0}_truth.npy'),
truth_0)
else:
# save all test data
for id_, item_truth, item_pred, item_out, threshold, T \
in zip(ids, intervals, prediction, out_np, thresholds, Ts):
np.save(path_test_result / f'{id_}_truth.npy', item_truth)
np.save(path_test_result / f'{id_}.npy', item_out[:T])
np.save(path_test_result / f'{id_}_pred.npy', item_pred)
all_thresholds[str(id_)] = threshold
str_eval = np.array2string(eval_result / n_batch, precision=3)
pbar.set_postfix_str(f'{loss / n_batch:.3f}, {str_eval}')
avg_loss += loss
avg_eval += eval_result
avg_loss = avg_loss / len(dataloader.dataset)
avg_eval = avg_eval / len(dataloader.dataset)
if mode == 'test':
np.savez(path_test_result / f'thresholds.npz', **all_thresholds)
return avg_loss, avg_eval
def step(self, valid_f1: float, epoch: int):
"""
:param valid_f1:
:param epoch:
:return: test epoch or 0
"""
last_restart = self.scheduler.last_restart
self.scheduler.step() # scheduler.last_restart can be updated
if epoch == self.scheduler.last_restart:
if valid_f1 < self.f1_last_restart:
return last_restart
else:
self.f1_last_restart = valid_f1
torch.save(self.model.module.state_dict(),
Path(self.writer.logdir, f'{epoch}.pt'))
return 0
class Tester:
@classmethod
def partition_masks(cls, output, target):
# Partition the union of the output and target into a true positive mask,
# a false positive mask, and a false negative mask
true_positive_mask = torch.min(output, target)
false_positive_mask = output - true_positive_mask
false_negative_mask = target - true_positive_mask
return true_positive_mask, false_positive_mask, false_negative_mask
@classmethod
def get_partition_measures(cls, output, target):
true_positive_mask, false_positive_mask, false_negative_mask = Tester.partition_masks(output, target)
tp = torch.sum(true_positive_mask) / (torch.sum(true_positive_mask) + torch.sum(false_positive_mask))
fp = torch.sum(false_positive_mask) / (torch.sum(true_positive_mask) + torch.sum(false_positive_mask))
fn = torch.sum(false_negative_mask) / (torch.sum(true_positive_mask) + torch.sum(false_negative_mask))
return tp, fp, fn
@classmethod
def get_dice(cls, output, target):
tp, fp, fn = Tester.get_partition_measures(output, target)
if tp + fp + fn == 0:
return -1
dice = (2*tp)/(2*tp + fp + fn)
if math.isnan(dice):
return 0
return dice.item()
@classmethod
def get_intersection_over_union(cls, output, target):
tp, fp, fn = Tester.get_partition_measures(output, target)
if tp + fp + fn == 0:
return -1
iou = tp / (tp + fp + fn)
if math.isnan(iou):
return 0
return iou.item()
@classmethod
def get_accuracy(cls, output, target):
tp, fp, fn = Tester.get_partition_measures(output, target)
if tp + fp == 0:
return -1
accuracy = tp / (tp + fp)
if math.isnan(accuracy):
return 0
return accuracy.item()
@classmethod
def get_recall(cls, output, target):
tp, fp, fn = Tester.get_partition_measures(output, target)
if tp + fn == 0:
return -1
recall = tp / (tp + fn)
if math.isnan(recall):
return 0
return recall.item()
@classmethod
def get_number_of_batches(cls, image_paths, batch_size):
batches = len(image_paths) / batch_size
if not batches.is_integer():
batches = math.floor(batches) + 1
return int(batches)
@classmethod
def evaluate_loss(cls, criterion, output, target):
loss_1 = criterion(output, target)
loss_2 = 1 - Tester.get_intersection_over_union(output, target)
loss = loss_1 + 0.1 * loss_2
return loss
def __init__(self, side_length, batch_size, seed, image_paths, state_dict):
self.side_length = side_length
self.batch_size = batch_size
self.seed = seed
self.image_paths = glob.glob(image_paths)
self.batches = Tester.get_number_of_batches(self.image_paths, self.batch_size)
self.model = UNet()
self.loader = Loader(self.side_length)
self.state_dict = state_dict
def set_cuda(self):
if torch.cuda.is_available():
self.model = self.model.cuda()
def set_seed(self):
if self.seed is not None:
np.random.seed(self.seed)
def process_batch(self, batch):
# Grab a batch, shuffled according to the provided seed. Note that
# i-th image: samples[i][0], i-th mask: samples[i][1]
samples = Loader.get_batch(self.image_paths, self.batch_size, batch, self.seed)
samples.astype(float)
# Cast samples into torch.FloatTensor for interaction with U-Net
samples = torch.from_numpy(samples)
samples = samples.float()
# Cast into a CUDA tensor, if GPUs are available
if torch.cuda.is_available():
samples = samples.cuda()
# Isolate images and their masks
samples_images = samples[:, 0]
samples_masks = samples[:, 1]
# Reshape for interaction with U-Net
samples_images = samples_images.unsqueeze(1)
samples_masks = samples_masks.unsqueeze(1)
# Run inputs through the model
output = self.model(samples_images)
# Clamp the target for proper interaction with BCELoss
target = torch.clamp(samples_masks, min=0, max=1)
del samples
return output, target
def test_model(self):
if torch.cuda.is_available():
buffered_state_dict = torch.load("weights/" + self.state_dict)
else:
buffered_state_dict = torch.load("weights/" + self.state_dict, map_location=lambda storage, loc: storage)
self.model.load_state_dict(buffered_state_dict)
self.model.eval()
criterion = nn.BCELoss()
perfect_accuracy_count = 0
zero_accuracy_count = 0
image_count = 0
accuracy_list = []
recall_list = []
iou_list = []
dice_list = []
losses_list = []
for batch in range(self.batches):
output, target = self.process_batch(batch)
loss = Tester.evaluate_loss(criterion, output, target)
print("Batch:", batch)
print("~~~~~~~~~~~~~~~~~~~~~~~~~~")
print("~~~~~~~~~~~~~~~~~~~~~~~~~~")
for i in range(0, output.shape[0]):
image_count += 1
binary_mask = Editor.make_binary_mask_from_torch(output[i, :, :, :], 1.0)
# Metrics
accuracy = Tester.get_accuracy(binary_mask, target[i, :, :, :].cpu())
recall = Tester.get_recall(binary_mask, target[i, :, :, :].cpu())
iou = Tester.get_intersection_over_union(binary_mask, target[i, :, :, :].cpu())
dice = Tester.get_dice(binary_mask, target[i, :, :, :].cpu())
if accuracy == 1:
perfect_accuracy_count += 1
if accuracy == 0:
zero_accuracy_count += 1
accuracy_list.append(accuracy)
recall_list.append(recall)
iou_list.append(iou)
dice_list.append(dice)
print("Accuracy:", accuracy)
print("Recall:", recall)
print("IoU:", iou)
print("Dice:", dice,"\n")
print("Mean Accuracy:", mean(accuracy_list))
print("Mean Recall:", mean(recall_list))
print("Mean IoU:", mean(iou_list))
print("Mean Dice:", mean(dice_list))
print("~~~~~~~~~~~~~~~~~~~~~~~~~~")
loss_value = loss.item()
losses_list.append(loss_value)
print("Test loss:", loss_value)
print("~~~~~~~~~~~~~~~~~~~~~~~~~~")
del output
del target
mean_iou = mean(iou_list)
mean_accuracy = mean(accuracy_list)
mean_recall = mean(recall_list)
mean_dice = mean(dice_list)
print("Perfect Accuracy Percentage:", perfect_accuracy_count / image_count)
print("Zero Accuracy Percentage:", zero_accuracy_count / image_count)
print("Mean Accuracy:", mean_accuracy)
print("Mean Recall:", mean_recall)
print("Mean IoU:", mean_iou)
print("Mean Dice:", mean_dice)
dset_pred = UnetPred(PRED_INPUT,
keywords=[args.modality],
im_size=[args.size, args.size],
transform=tr.ToTensor())
pred_loader = DataLoader(dset_pred,
batch_size=args.test_batch_size,
shuffle=False,
num_workers=1)
print("Prediction Data : ", len(pred_loader.dataset))
# %% Loading in the model
model = UNet()
if args.cuda:
model.cuda()
if args.optimizer == 'SGD':
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.99)
if args.optimizer == 'ADAM':
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
betas=(args.beta1, args.beta2))
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
# Defining Loss Function
criterion = DICELossMultiClass()
def train(epoch, scheduler, loss_lsit):
class Trainer:
@classmethod
def intersection_over_union(cls, y, z):
iou = (torch.sum(torch.min(y, z))) / (torch.sum(torch.max(y, z)))
return iou
@classmethod
def get_number_of_batches(cls, image_paths, batch_size):
batches = len(image_paths) / batch_size
if not batches.is_integer():
batches = math.floor(batches) + 1
return int(batches)
@classmethod
def evaluate_loss(cls, criterion, output, target):
loss_1 = criterion(output, target)
loss_2 = 1 - Trainer.intersection_over_union(output, target)
loss = loss_1 + 0.1 * loss_2
return loss
def __init__(self, side_length, batch_size, epochs, learning_rate,
momentum_parameter, seed, image_paths, state_dict,
train_val_split):
self.side_length = side_length
self.batch_size = batch_size
self.epochs = epochs
self.learning_rate = learning_rate
self.momentum_parameter = momentum_parameter
self.seed = seed
self.image_paths = glob.glob(image_paths)
self.batches = Trainer.get_number_of_batches(self.image_paths,
self.batch_size)
self.model = UNet()
self.loader = Loader(self.side_length)
self.state_dict = state_dict
self.train_val_split = train_val_split
self.train_size = int(np.floor((self.train_val_split * self.batches)))
def set_cuda(self):
if torch.cuda.is_available():
self.model = self.model.cuda()
def set_seed(self):
if self.seed is not None:
np.random.seed(self.seed)
def process_batch(self, batch):
# Grab a batch, shuffled according to the provided seed. Note that
# i-th image: samples[i][0], i-th mask: samples[i][1]
samples = Loader.get_batch(self.image_paths, self.batch_size, batch,
self.seed)
samples.astype(float)
# Cast samples into torch.FloatTensor for interaction with U-Net
samples = torch.from_numpy(samples)
samples = samples.float()
# Cast into a CUDA tensor, if GPUs are available
if torch.cuda.is_available():
samples = samples.cuda()
# Isolate images and their masks
samples_images = samples[:, 0]
samples_masks = samples[:, 1]
# Reshape for interaction with U-Net
samples_images = samples_images.unsqueeze(1)
samples_masks = samples_masks.unsqueeze(1)
# Run inputs through the model
output = self.model(samples_images)
# Clamp the target for proper interaction with BCELoss
target = torch.clamp(samples_masks, min=0, max=1)
del samples
return output, target
def train_model(self):
self.model.train()
criterion = nn.BCELoss()
optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
iteration = 0
best_iteration = 0
best_loss = 10**10
losses_train = []
losses_val = []
iou_train = []
average_iou_train = []
iou_val = []
average_iou_val = []
print("BEGIN TRAINING")
print("TRAINING BATCHES:", self.train_size)
print("VALIDATION BATCHES:", self.batches - self.train_size)
print("BATCH SIZE:", self.batch_size)
print("EPOCHS:", self.epochs)
print("~~~~~~~~~~~~~~~~~~~~~~~~~~")
for k in range(0, self.epochs):
print("EPOCH:", k + 1)
print("~~~~~~~~~~~~~~~~~~~~~~~~~~")
# Train
for batch in range(0, self.train_size):
iteration = iteration + 1
output, target = self.process_batch(batch)
loss = Trainer.evaluate_loss(criterion, output, target)
print("EPOCH:", self.epochs)
print("Batch", batch, "of", self.train_size)
# Aggregate intersection over union scores for each element in the batch
for i in range(0, output.shape[0]):
binary_mask = Editor.make_binary_mask_from_torch(
output[i, :, :, :], 1.0)
iou = Trainer.intersection_over_union(
binary_mask, target[i, :, :, :].cpu())
iou_train.append(iou.item())
print("IoU:", iou.item())
# Clear data to prevent memory overload
del target
del output
# Clear gradients, back-propagate, and update weights
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Record the loss value
loss_value = loss.item()
if best_loss > loss_value:
best_loss = loss_value
best_iteration = iteration
losses_train.append(loss_value)
if batch == self.train_size - 1:
print("LOSS:", loss_value)
print("~~~~~~~~~~~~~~~~~~~~~~~~~~")
average_iou = sum(iou_train) / len(iou_train)
print("Average IoU:", average_iou)
average_iou_train.append(average_iou)
#Visualizer.save_loss_plot(average_iou_train, "average_iou_train.png")
# Validate
for batch in range(self.train_size, self.batches):
output, target = self.process_batch(batch)
loss = Trainer.evaluate_loss(criterion, output, target)
for i in range(0, output.shape[0]):
binary_mask = Editor.make_binary_mask_from_torch(
output[i, :, :, :], 1.0)
iou = Trainer.intersection_over_union(
binary_mask, target[i, :, :, :].cpu())
iou_val.append(iou.item())
print("IoU:", iou.item())
loss_value = loss.item()
losses_val.append(loss_value)
print("EPOCH:", self.epochs)
print("VALIDATION LOSS:", loss_value)
print("~~~~~~~~~~~~~~~~~~~~~~~~~~")
del output
del target
average_iou = sum(iou_val) / len(iou_val)
print("Average IoU:", average_iou)
average_iou_val.append(average_iou)
#Visualizer.save_loss_plot(average_iou_val, "average_iou_val.png")
print("Least loss", best_loss, "at iteration", best_iteration)
torch.save(self.model.state_dict(), "weights/" + self.state_dict)
def video_infer(args):
cap = cv2.VideoCapture(args.video)
_, frame = cap.read()
H, W = frame.shape[:2]
fourcc = cv2.VideoWriter_fourcc(*'DIVX')
out = cv2.VideoWriter(args.output, fourcc, 30, (W,H))
font = cv2.FONT_HERSHEY_SIMPLEX
# Background
if args.bg is not None:
BACKGROUND = cv2.imread(args.bg)[...,::-1]
BACKGROUND = cv2.resize(BACKGROUND, (W,H), interpolation=cv2.INTER_LINEAR)
KERNEL_SZ = 25
SIGMA = 0
# Alpha transperency
else:
COLOR1 = [90, 140, 154]
COLOR2 = [0, 0, 0]
if args.model=='unet':
model = UNet(backbone=args.net, num_classes=2, pretrained_backbone=None)
elif args.model=='deeplabv3_plus':
model = DeepLabV3Plus(backbone=args.net, num_classes=2, pretrained_backbone=None)
elif args.model=='hrnet':
model = HighResolutionNet(num_classes=2, pretrained_backbone=None)
if args.use_cuda:
model = model.cuda()
trained_dict = torch.load(args.checkpoint, map_location="cpu")['state_dict']
model.load_state_dict(trained_dict, strict=False)
model.eval()
while(cap.isOpened()):
start_time = time()
ret, frame = cap.read()
if ret:
image = frame[...,::-1]
h, w = image.shape[:2]
read_cam_time = time()
# Predict mask
X, pad_up, pad_left, h_new, w_new = utils.preprocessing(image, expected_size=args.input_sz, pad_value=0)
preproc_time = time()
with torch.no_grad():
if args.use_cuda:
mask = model(X.cuda())
if mask.shape[1] != h_new:
mask = F.interpolate(mask, size=(args.input_sz, args.input_sz), mode='bilinear', align_corners=True)
mask = mask[..., pad_up: pad_up+h_new, pad_left: pad_left+w_new]
mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=True)
mask = F.softmax(mask, dim=1)
mask = mask[0,1,...].cpu().numpy()
else:
mask = model(X)
mask = mask[..., pad_up: pad_up+h_new, pad_left: pad_left+w_new]
mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=True)
mask = F.softmax(mask, dim=1)
mask = mask[0,1,...].numpy()
predict_time = time()
# Draw result
if args.bg is None:
image_alpha = utils.draw_matting(image, mask)
#image_alpha = utils.draw_transperency(image, mask, COLOR1, COLOR2)
else:
image_alpha = utils.draw_fore_to_back(image, mask, BACKGROUND, kernel_sz=KERNEL_SZ, sigma=SIGMA)
draw_time = time()
# Print runtime
read = read_cam_time-start_time
preproc = preproc_time-read_cam_time
pred = predict_time-preproc_time
draw = draw_time-predict_time
total = read + preproc + pred + draw
fps = 1 / pred
print("read: %.3f [s]; preproc: %.3f [s]; pred: %.3f [s]; draw: %.3f [s]; total: %.3f [s]; fps: %.2f [Hz]" %
(read, preproc, pred, draw, total, fps))
# Wait for interupt
cv2.putText(image_alpha, "%.2f [fps]" % (fps), (10, 50), font, 1.5, (0, 255, 0), 2, cv2.LINE_AA)
out.write(image_alpha[..., ::-1])
if args.watch:
cv2.imshow('webcam', image_alpha[..., ::-1])
if cv2.waitKey(1) & 0xFF == ord('q'):
break
else:
break
def video_infer(args):
cap = cv2.VideoCapture(args.video)
_, frame = cap.read()
H, W = frame.shape[:2]
fps = cap.get(cv2.CAP_PROP_FPS)
out = cv2.VideoWriter(args.output,
cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps,
(W, H))
# Background
if args.bg is not None:
BACKGROUND = cv2.imread(args.bg)[..., ::-1]
BACKGROUND = cv2.resize(BACKGROUND, (W, H),
interpolation=cv2.INTER_LINEAR)
KERNEL_SZ = 25
SIGMA = 0
# Alpha transperency
else:
COLOR1 = [255, 0, 0]
COLOR2 = [0, 0, 255]
if args.model == 'unet':
model = UNet(backbone=args.net,
num_classes=2,
pretrained_backbone=None)
elif args.model == 'deeplabv3_plus':
model = DeepLabV3Plus(backbone=args.net,
num_classes=2,
pretrained_backbone=None)
if args.use_cuda:
model = model.cuda()
trained_dict = torch.load(args.checkpoint,
map_location="cpu")['state_dict']
model.load_state_dict(trained_dict, strict=False)
model.eval()
if W > H:
w_new = int(args.input_sz)
h_new = int(H * w_new / W)
else:
h_new = int(args.input_sz)
w_new = int(W * h_new / H)
disflow = cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST)
prev_gray = np.zeros((h_new, w_new), np.uint8)
prev_cfd = np.zeros((h_new, w_new), np.float32)
is_init = True
while (cap.isOpened()):
start_time = time()
ret, frame = cap.read()
if ret:
image = frame[..., ::-1]
h, w = image.shape[:2]
read_cam_time = time()
# Predict mask
X, pad_up, pad_left, h_new, w_new = utils.preprocessing(
image, expected_size=args.input_sz, pad_value=0)
preproc_time = time()
with torch.no_grad():
if args.use_cuda:
mask = model(X.cuda())
mask = mask[..., pad_up:pad_up + h_new,
pad_left:pad_left + w_new]
#mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=True)
mask = F.softmax(mask, dim=1)
mask = mask[0, 1, ...].cpu().numpy() #(213, 320)
else:
mask = model(X)
mask = mask[..., pad_up:pad_up + h_new,
pad_left:pad_left + w_new]
#mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=True)
mask = F.softmax(mask, dim=1)
mask = mask[0, 1, ...].numpy()
predict_time = time()
# optical tracking
cur_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
cur_gray = cv2.resize(cur_gray, (w_new, h_new))
scoremap = 255 * mask
optflow_map = postprocess(cur_gray, scoremap, prev_gray, prev_cfd,
disflow, is_init)
optical_flow_track_time = time()
prev_gray = cur_gray.copy()
prev_cfd = optflow_map.copy()
is_init = False
optflow_map = cv2.GaussianBlur(optflow_map, (3, 3), 0)
optflow_map = threshold_mask(optflow_map,
thresh_bg=0.2,
thresh_fg=0.8)
img_matting = np.repeat(optflow_map[:, :, np.newaxis], 3, axis=2)
bg_im = np.ones_like(img_matting) * 255
re_image = cv2.resize(image, (w_new, h_new))
comb = (img_matting * re_image + (1 - img_matting) * bg_im).astype(
np.uint8)
comb = cv2.resize(comb, (W, H))
comb = comb[..., ::-1]
# Print runtime
read = read_cam_time - start_time
preproc = preproc_time - read_cam_time
pred = predict_time - preproc_time
optical = optical_flow_track_time - predict_time
total = read + preproc + pred + optical
print(
"read: %.3f [s]; preproc: %.3f [s]; pred: %.3f [s]; optical: %.3f [s]; total: %.3f [s]; fps: %.2f [Hz]"
% (read, preproc, pred, optical, total, 1 / pred))
out.write(comb)
if args.watch:
cv2.imshow('webcam', comb[..., ::-1])
if cv2.waitKey(1) & 0xFF == ord('q'):
break
else:
break
cap.release()
out.release()
default='UNet',
type=str,
help='name of model for saving/loading weights')
parser.add_argument('--exp_name',
default='single_nuclei_lossless_augmentations',
type=str,
help='name of experiment for saving files')
args = parser.parse_args()
LR_SCHED = {0: args.lr, 100: args.lr * 0.1, 150: args.lr * 0.01}
torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
net = UNet()
net.cuda()
# set model filenames
MODEL_CKPT = '../models-pytorch/best_{}_{}.pth'.format(args.model_name,
args.exp_name)
train_loader, valid_loader, len_train, len_valid = get_cropimg_training_loaders(
imsize=args.img_size,
test_size=args.valid_size,
batch_size=args.batch_size,
augment_prob=args.aug_prob,
include_rnd=True)
msk_crit = nn.BCEWithLogitsLoss().cuda()
l1_crit = nn.L1Loss().cuda()
optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=args.l2)