def test():
device = torch.device(conf.cuda if torch.cuda.is_available() else "cpu")
test_dataset = Testinging_Dataset(conf.data_path_test,
conf.test_noise_param,
conf.crop_img_size)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
print('Loading model from: {}'.format(conf.model_path_test))
model = UNet(in_channels=conf.img_channel, out_channels=conf.img_channel)
print('loading model')
model.load_state_dict(torch.load(conf.model_path_test))
model.eval()
model.to(device)
result_dir = conf.denoised_dir
if not os.path.exists(result_dir):
os.mkdir(result_dir)
for batch_idx, (source, img_cropped) in enumerate(test_loader):
source_img = tvF.to_pil_image(source.squeeze(0))
img_truth = img_cropped.squeeze(0).numpy().astype(np.uint8)
source = source.to(device)
denoised_img = model(source).detach().cpu()
img_name = test_loader.dataset.image_list[batch_idx]
denoised_result = tvF.to_pil_image(
torch.clamp(denoised_img.squeeze(0), 0, 1))
fname = os.path.splitext(img_name)[0]
source_img.save(os.path.join(result_dir, f'{fname}-noisy.png'))
denoised_result.save(os.path.join(result_dir, f'{fname}-denoised.png'))
io.imsave(os.path.join(result_dir, f'{fname}-ground_truth.png'),
img_truth)
def get_unet(self, use_distributed_data_parallel=True):
""" Creates a new network and returns. If machine has multiple GPUs, uses them.
"""
net = UNet(n_channels=self.channel_count, n_classes=1, bilinear=True,
running_on_gpu=(self.gpu_number is not None))
net.to(self.device)
if not use_distributed_data_parallel:
return net
net = nn.parallel.DistributedDataParallel(net, device_ids=[self.gpu_number])
return net
def train():
device = torch.device(conf.cuda if torch.cuda.is_available() else "cpu")
dataset = Training_Dataset(conf.data_path_train, conf.gaussian_noise_param,
conf.crop_img_size)
dataset_length = len(dataset)
train_loader = DataLoader(dataset,
batch_size=4,
shuffle=True,
num_workers=4)
model = UNet(in_channels=conf.img_channel, out_channels=conf.img_channel)
criterion = nn.MSELoss()
model = model.to(device)
optim = Adam(model.parameters(),
lr=conf.learning_rate,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=0,
amsgrad=True)
scheduler = lr_scheduler.StepLR(optim, step_size=100, gamma=0.5)
model.train()
print(model)
print("Starting Training Loop...")
since = time.time()
for epoch in range(conf.max_epoch):
print('Epoch {}/{}'.format(epoch, conf.max_epoch - 1))
print('-' * 10)
running_loss = 0.0
scheduler.step()
for batch_idx, (source, target) in enumerate(train_loader):
source = source.to(device)
target = target.to(device)
optim.zero_grad()
denoised_source = model(source)
loss = criterion(denoised_source, target)
loss.backward()
optim.step()
running_loss += loss.item() * source.size(0)
print('Current loss {} and current batch idx {}'.format(
loss.item(), batch_idx))
epoch_loss = running_loss / dataset_length
print('{} Loss: {:.4f}'.format('current ' + str(epoch), epoch_loss))
if (epoch + 1) % conf.save_per_epoch == 0:
save_model(model, epoch + 1)
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
def submit_mnms(model_path, input_data_directory, output_data_directory,
device):
data_paths = load_path(input_data_directory)
net = UNet(n_channels=1, n_classes=4, bilinear=True)
net.load_state_dict(torch.load(model_path, map_location=device))
net.to(device)
for path in data_paths:
ED_np, ES_np = load_phase(path) # HxWxF
ED_masks = []
ES_masks = []
for i in range(ED_np.shape[2]):
img_np = ED_np[:, :, i]
img_tensor = pre_transform(img_np)
img_tensor = img_tensor.to(device)
mask = predict_img(net, img_tensor)
mask = post_transform(img_np, mask[0:3, :, :])
ED_masks.append(mask)
for i in range(ES_np.shape[2]):
img_np = ES_np[:, :, i]
img_tensor = pre_transform(img_np)
img_tensor = img_tensor.to(device)
mask = predict_img(net, img_tensor)
mask = post_transform(img_np, mask[0:3, :, :])
ES_masks.append(mask)
ED_masks = np.concatenate(ED_masks, axis=2)
ES_masks = np.concatenate(ES_masks, axis=2)
save_phase(ED_masks, ES_masks, output_data_directory, path)
def select_model(model_name, init_msg):
logger = get_logger()
logger.info(init_msg)
if model_name == "SETR-PUP":
_, model = get_SETR_PUP()
elif model_name == "SETR-MLA":
_, model = get_SETR_MLA()
elif model_name == "TransUNet-Base":
model = get_TransUNet_base()
elif model_name == "TransUNet-Large":
model = get_TransUNet_large()
elif model_name == "UNet":
model = UNet(CLASS_NUM)
model = model.to(device)
return logger, model
discriminator_g = GlobalDiscriminator()
discriminator_l = LocalDiscriminator()
resume=False
if(len(sys.argv)>1 and sys.argv[1]=='resume'):
resume=True
# Load model if available
if(resume==True):
print('Resuming training....')
generator.load_state_dict(torch.load(os.path.join(model_path,'model_gen_latest')))
discriminator_g.load_state_dict(torch.load(os.path.join(model_path,'model_gdis_latest')))
discriminator_l.load_state_dict(torch.load(os.path.join(model_path,'model_ldis_latest')))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
generator = generator.to(device)
discriminator_g = discriminator_g.to(device)
discriminator_l = discriminator_l.to(device)
optimizer_g = optim.Adam(discriminator_g.parameters(), lr=0.00005)
optimizer_l = optim.Adam(discriminator_l.parameters(), lr=0.00005)
gen_optimizer = optim.Adam(generator.parameters(), lr=0.0002)
lossdis = nn.BCELoss()
lossgen = FocalLoss()
lamda = 75
data_loader = load_images(data_path)
num_epochs = 2000
class Train(object):
def __init__(self, configs):
self.batch_size = configs.get("batch_size", "16")
self.epochs = configs.get("epochs", "100")
self.lr = configs.get("lr", "0.0001")
device_args = configs.get("device", "cuda")
self.device = torch.device(
"cpu" if not torch.cuda.is_available() else device_args)
self.workers = configs.get("workers", "4")
self.vis_images = configs.get("vis_images", "200")
self.vis_freq = configs.get("vis_freq", "10")
self.weights = configs.get("weights", "./weights")
if not os.path.exists(self.weights):
os.mkdir(self.weights)
self.logs = configs.get("logs", "./logs")
if not os.path.exists(self.weights):
os.mkdir(self.weights)
self.images_path = configs.get("images_path", "./data")
self.is_resize = config.get("is_resize", False)
self.image_short_side = config.get("image_short_side", 256)
self.is_padding = config.get("is_padding", False)
is_multi_gpu = config.get("DateParallel", False)
pre_train = config.get("pre_train", False)
model_path = config.get("model_path", './weights/unet_idcard_adam.pth')
# self.image_size = configs.get("image_size", "256")
# self.aug_scale = configs.get("aug_scale", "0.05")
# self.aug_angle = configs.get("aug_angle", "15")
self.step = 0
self.dsc_loss = DiceLoss()
self.model = UNet(in_channels=Dataset.in_channels,
out_channels=Dataset.out_channels)
if pre_train:
self.model.load_state_dict(torch.load(model_path,
map_location=self.device),
strict=False)
if is_multi_gpu:
self.model = nn.DataParallel(self.model)
self.model.to(self.device)
self.best_validation_dsc = 0.0
self.loader_train, self.loader_valid = self.data_loaders()
self.params = [p for p in self.model.parameters() if p.requires_grad]
self.optimizer = optim.Adam(self.params,
lr=self.lr,
weight_decay=0.0005)
# self.optimizer = torch.optim.SGD(self.params, lr=self.lr, momentum=0.9, weight_decay=0.0005)
self.scheduler = lr_scheduler.LR_Scheduler_Head(
'poly', self.lr, self.epochs, len(self.loader_train))
def datasets(self):
train_datasets = Dataset(
images_dir=self.images_path,
# image_size=self.image_size,
subset="train", # train
transform=get_transforms(train=True),
is_resize=self.is_resize,
image_short_side=self.image_short_side,
is_padding=self.is_padding)
# valid_datasets = train_datasets
valid_datasets = Dataset(
images_dir=self.images_path,
# image_size=self.image_size,
subset="validation", # validation
transform=get_transforms(train=False),
is_resize=self.is_resize,
image_short_side=self.image_short_side,
is_padding=False)
return train_datasets, valid_datasets
def data_loaders(self):
dataset_train, dataset_valid = self.datasets()
loader_train = DataLoader(
dataset_train,
batch_size=self.batch_size,
shuffle=True,
drop_last=True,
num_workers=self.workers,
)
loader_valid = DataLoader(
dataset_valid,
batch_size=1,
drop_last=False,
num_workers=self.workers,
)
return loader_train, loader_valid
@staticmethod
def dsc_per_volume(validation_pred, validation_true):
assert len(validation_pred) == len(validation_true)
dsc_list = []
for p in range(len(validation_pred)):
y_pred = np.array([validation_pred[p]])
y_true = np.array([validation_true[p]])
dsc_list.append(dsc(y_pred, y_true))
return dsc_list
@staticmethod
def get_logger(filename, verbosity=1, name=None):
level_dict = {0: logging.DEBUG, 1: logging.INFO, 2: logging.WARNING}
formatter = logging.Formatter(
"[%(asctime)s][%(filename)s][line:%(lineno)d][%(levelname)s] %(message)s"
)
logger = logging.getLogger(name)
logger.setLevel(level_dict[verbosity])
fh = logging.FileHandler(filename, "w")
fh.setFormatter(formatter)
logger.addHandler(fh)
sh = logging.StreamHandler()
sh.setFormatter(formatter)
logger.addHandler(sh)
return logger
def train_one_epoch(self, epoch):
self.model.train()
loss_train = []
for i, data in enumerate(self.loader_train):
self.scheduler(self.optimizer, i, epoch, self.best_validation_dsc)
x, y_true = data
x, y_true = x.to(self.device), y_true.to(self.device)
y_pred = self.model(x)
# print('1111', y_pred.size())
# print('2222', y_true.size())
loss = self.dsc_loss(y_pred, y_true)
loss_train.append(loss.item())
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# lr_scheduler.step()
if self.step % 200 == 0:
print('Epoch:[{}/{}]\t iter:[{}]\t loss={:.5f}\t '.format(
epoch, self.epochs, i, loss))
self.step += 1
def eval_model(self, patience):
self.model.eval()
loss_valid = []
validation_pred = []
validation_true = []
# early_stopping = EarlyStopping(patience=patience, verbose=True)
for i, data in enumerate(self.loader_valid):
x, y_true = data
x, y_true = x.to(self.device), y_true.to(self.device)
# print(x.size())
# print(333,x[0][2])
with torch.no_grad():
y_pred = self.model(x)
loss = self.dsc_loss(y_pred, y_true)
# print(y_pred.shape)
mask = y_pred > 0.5
mask = mask * 255
mask = mask.cpu().numpy()[0][0]
# print(mask)
# print(mask.shape())
cv2.imwrite('result.png', mask)
loss_valid.append(loss.item())
y_pred_np = y_pred.detach().cpu().numpy()
validation_pred.extend(
[y_pred_np[s] for s in range(y_pred_np.shape[0])])
y_true_np = y_true.detach().cpu().numpy()
validation_true.extend(
[y_true_np[s] for s in range(y_true_np.shape[0])])
# early_stopping(loss_valid, self.model)
# if early_stopping.early_stop:
# print('Early stopping')
# import sys
# sys.exit(1)
mean_dsc = np.mean(
self.dsc_per_volume(
validation_pred,
validation_true,
))
# print('mean_dsc:', mean_dsc)
if mean_dsc > self.best_validation_dsc:
self.best_validation_dsc = mean_dsc
torch.save(self.model.state_dict(),
os.path.join(self.weights, "unet_xia_adam.pth"))
print("Best validation mean DSC: {:4f}".format(
self.best_validation_dsc))
def main(self):
# print('train is begin.....')
# print('load data end.....')
# loaders = {"train": loader_train, "valid": loader_valid}
for epoch in tqdm(range(self.epochs), total=self.epochs):
self.train_one_epoch(epoch)
self.eval_model(patience=10)
torch.save(self.model.state_dict(),
os.path.join(self.weights, "unet_final.pth"))
def train(cont=False):
# for tensorboard tracking
logger = get_logger()
logger.info("(1) Initiating Training ... ")
logger.info("Training on device: {}".format(device))
writer = SummaryWriter()
# init model
aux_layers = None
if net == "SETR-PUP":
aux_layers, model = get_SETR_PUP()
elif net == "SETR-MLA":
aux_layers, model = get_SETR_MLA()
elif net == "TransUNet-Base":
model = get_TransUNet_base()
elif net == "TransUNet-Large":
model = get_TransUNet_large()
elif net == "UNet":
model = UNet(CLASS_NUM)
# prepare dataset
cluster_model = get_clustering_model(logger)
train_dataset = CityscapeDataset(img_dir=data_dir,
img_dim=IMG_DIM,
mode="train",
cluster_model=cluster_model)
valid_dataset = CityscapeDataset(img_dir=data_dir,
img_dim=IMG_DIM,
mode="val",
cluster_model=cluster_model)
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True)
valid_loader = DataLoader(valid_dataset,
batch_size=batch_size,
shuffle=False)
logger.info("(2) Dataset Initiated. ")
# optimizer
epochs = epoch_num if epoch_num > 0 else iteration_num // len(
train_loader) + 1
optim = SGD(model.parameters(),
lr=lrate,
momentum=momentum,
weight_decay=wdecay)
# optim = Adam(model.parameters(), lr=lrate)
scheduler = lr_scheduler.MultiStepLR(
optim, milestones=[int(epochs * fine_tune_ratio)], gamma=0.1)
cur_epoch = 0
best_loss = float('inf')
epochs_since_improvement = 0
# for continue training
if cont:
model, optim, cur_epoch, best_loss = load_ckpt_continue_training(
best_ckpt_src, model, optim, logger)
logger.info("Current best loss: {0}".format(best_loss))
with warnings.catch_warnings():
warnings.simplefilter("ignore")
for i in range(cur_epoch):
scheduler.step()
else:
model = nn.DataParallel(model)
model = model.to(device)
logger.info("(3) Model Initiated ... ")
logger.info("Training model: {}".format(net) + ". Training Started.")
# loss
ce_loss = CrossEntropyLoss()
if use_dice_loss:
dice_loss = DiceLoss(CLASS_NUM)
# loop over epochs
iter_count = 0
epoch_bar = tqdm.tqdm(total=epochs,
desc="Epoch",
position=cur_epoch,
leave=True)
logger.info("Total epochs: {0}. Starting from epoch {1}.".format(
epochs, cur_epoch + 1))
for e in range(epochs - cur_epoch):
epoch = e + cur_epoch
# Training.
model.train()
trainLossMeter = LossMeter()
train_batch_bar = tqdm.tqdm(total=len(train_loader),
desc="TrainBatch",
position=0,
leave=True)
for batch_num, (orig_img, mask_img) in enumerate(train_loader):
orig_img, mask_img = orig_img.float().to(
device), mask_img.float().to(device)
if net == "TransUNet-Base" or net == "TransUNet-Large":
pred = model(orig_img)
elif net == "SETR-PUP" or net == "SETR-MLA":
if aux_layers is not None:
pred, _ = model(orig_img)
else:
pred = model(orig_img)
elif net == "UNet":
pred = model(orig_img)
loss_ce = ce_loss(pred, mask_img[:].long())
if use_dice_loss:
loss_dice = dice_loss(pred, mask_img, softmax=True)
loss = 0.5 * (loss_ce + loss_dice)
else:
loss = loss_ce
# Backward Propagation, Update weight and metrics
optim.zero_grad()
loss.backward()
optim.step()
# update learning rate
for param_group in optim.param_groups:
orig_lr = param_group['lr']
param_group['lr'] = orig_lr * (1.0 -
iter_count / iteration_num)**0.9
iter_count += 1
# Update loss
trainLossMeter.update(loss.item())
# print status
if (batch_num + 1) % print_freq == 0:
status = 'Epoch: [{0}][{1}/{2}]\t' \
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(epoch+1, batch_num+1, len(train_loader), loss=trainLossMeter)
logger.info(status)
# log loss to tensorboard
if (batch_num + 1) % tensorboard_freq == 0:
writer.add_scalar(
'Train_Loss_{0}'.format(tensorboard_freq),
trainLossMeter.avg,
epoch * (len(train_loader) / tensorboard_freq) +
(batch_num + 1) / tensorboard_freq)
train_batch_bar.update(1)
writer.add_scalar('Train_Loss_epoch', trainLossMeter.avg, epoch)
# Validation.
model.eval()
validLossMeter = LossMeter()
valid_batch_bar = tqdm.tqdm(total=len(valid_loader),
desc="ValidBatch",
position=0,
leave=True)
with torch.no_grad():
for batch_num, (orig_img, mask_img) in enumerate(valid_loader):
orig_img, mask_img = orig_img.float().to(
device), mask_img.float().to(device)
if net == "TransUNet-Base" or net == "TransUNet-Large":
pred = model(orig_img)
elif net == "SETR-PUP" or net == "SETR-MLA":
if aux_layers is not None:
pred, _ = model(orig_img)
else:
pred = model(orig_img)
elif net == "UNet":
pred = model(orig_img)
loss_ce = ce_loss(pred, mask_img[:].long())
if use_dice_loss:
loss_dice = dice_loss(pred, mask_img, softmax=True)
loss = 0.5 * (loss_ce + loss_dice)
else:
loss = loss_ce
# Update loss
validLossMeter.update(loss.item())
# print status
if (batch_num + 1) % print_freq == 0:
status = 'Validation: [{0}][{1}/{2}]\t' \
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(epoch+1, batch_num+1, len(valid_loader), loss=validLossMeter)
logger.info(status)
# log loss to tensorboard
if (batch_num + 1) % tensorboard_freq == 0:
writer.add_scalar(
'Valid_Loss_{0}'.format(tensorboard_freq),
validLossMeter.avg,
epoch * (len(valid_loader) / tensorboard_freq) +
(batch_num + 1) / tensorboard_freq)
valid_batch_bar.update(1)
valid_loss = validLossMeter.avg
writer.add_scalar('Valid_Loss_epoch', valid_loss, epoch)
logger.info("Validation Loss of epoch [{0}/{1}]: {2}\n".format(
epoch + 1, epochs, valid_loss))
# update optim scheduler
scheduler.step()
# save checkpoint
is_best = valid_loss < best_loss
best_loss_tmp = min(valid_loss, best_loss)
if not is_best:
epochs_since_improvement += 1
logger.info("Epochs since last improvement: %d\n" %
(epochs_since_improvement, ))
if epochs_since_improvement == early_stop_tolerance:
break # early stopping.
else:
epochs_since_improvement = 0
state = {
'epoch': epoch,
'loss': best_loss_tmp,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optim.state_dict(),
}
torch.save(state, ckpt_src)
logger.info("Checkpoint updated.")
best_loss = best_loss_tmp
epoch_bar.update(1)
writer.close()
import glob
import numpy as np
import torch
import os
import cv2
from torchvision import transforms
from unet_model import UNet
if __name__ == "__main__":
# 选择设备，有cuda用cuda，没有就用cpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 加载网络，图片单通道，分类为1。
net = UNet(n_channels=1, n_classes=1)
# 将网络拷贝到deivce中
net.to(device=device)
# 加载模型参数
# net.load_state_dict(torch.load('best_model_100X.pth', map_location=device))
net.load_state_dict(torch.load('best_model.pth', map_location=device))
# 测试模式
net.eval()
# 读取所有图片路径
tests_path = glob.glob('dataS/test/*.png')
# tests_path = glob.glob('data100X/test/*.png')
# 遍历素有图片
for test_path in tests_path:
# 保存结果地址
save_res_path = test_path.split('.')[0] + '_res.png'
# 读取图片
img = cv2.imread(test_path)
# 转为灰度图
# load best model weights
model.load_state_dict(best_model_wts)
return model
if __name__ == '__main__':
lr = 0.001
model = UNet(n_channels=1)
#num_ftrs = model.fc.in_features
# Here the size of each output sample is set to 2
# Alternatively it can be generalized to nn.Linear(num_ftrs, len(class_names))
#model.fc = nn.Linear(num_ftrs, 2)
model = model.to(device)
criterion = nn.CrossEntropyLoss()
# Observe that all parameters are being optimized
optimizer_ft = optim.Adam(model.parameters(), lr=lr) #, momentum=0.9)
# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft,
step_size=7,
gamma=0.1)
# Finetune training the convnet and evaluation
model = train_model(model,
criterion,
optimizer_ft,
import numpy as np
import torch
import matplotlib.pyplot as plt
import tqdm
import cv2
from unet_model import UNet
#device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
device = 'cpu'
best_model_name = 'best_model.pt'
best_model = torch.load(best_model_name)
model = UNet()
model.load_state_dict(best_model['state_dict'])
model.eval()
model.to(device)
test_dir = '../data/poster/images/'
out_dir = '../data/poster/model/'
test_images = [os.path.join(test_dir, x) for x in os.listdir(test_dir)]
counter = 0
i = 0
for i in range(len(test_images)):
test_image_one = test_images[i]
#if 'post' not in test_image_one:
# i += 1
# continue
#counter += 1
#i += 1
print(i, test_image_one)
class OnePredict(object):
def __init__(self, params):
self.params = params
self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
self.model_path = params['model_path']
self.model = UNet(in_channels=3, out_channels=1)
self.threshold = 0.5
self.resume()
# self.model.eval()
self.transform = get_transforms_3()
self.is_resize = True
self.image_short_side = 1024
self.init_torch_tensor()
self.model.eval()
def init_torch_tensor(self):
torch.set_default_tensor_type('torch.FloatTensor')
if torch.cuda.is_available():
self.device = torch.device('cuda')
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
self.device = torch.device('cpu')
# self.model.to(self.device)
def resume(self):
self.model.load_state_dict(torch.load(self.model_path, map_location=self.device), strict=False)
self.model.to(self.device)
def resize_img(self, img):
'''输入PIL格式的图片'''
width, height = img.size
# print('111', img.size)
if self.is_resize:
if height < width:
new_height = self.image_short_side
new_width = int(math.ceil(new_height / height * width / 32) * 32)
else:
new_width = self.image_short_side
new_height = int(math.ceil(new_width / width * height / 32) * 32)
else:
if height < width:
scale = int(height / 32)
new_image_short_side = scale * 32
new_height = new_image_short_side
new_width = int(math.ceil(new_height / height * width / 32) * 32)
else:
scale = int(width / 32)
new_image_short_side = scale * 32
new_width = new_image_short_side
new_height = int(math.ceil(new_width / width * height / 32) * 32)
# print('test1:', np.array(img))
# print('new:', (new_width, new_height))
resized_img = img.resize((new_width, new_height), Image.ANTIALIAS)
# print(new_height, new_width)
# print('test2:', np.array(resized_img))
return resized_img
def format_output(self):
pass
@staticmethod
def pre_process(img):
return img
@staticmethod
def pad_sample(img):
a = img.size[0]
b = img.size[1]
if a == b:
return img
diff = (max(a, b) - min(a, b)) / 2.0
if a > b:
padding = (0, int(np.floor(diff)), 0, int(np.ceil(diff)))
else:
padding = (int(np.floor(diff)), 0, int(np.ceil(diff)), 0)
img = ImageOps.expand(img, border=padding, fill=0) ##left,top,right,bottom
assert img.size[0] == img.size[1]
return img
def post_process(self, preds, img):
mask = preds > self.threshold
mask = mask * 255
# print(mask.size())
mask = mask.cpu().numpy()[0][0]
# print(mask)
# print(mask.shape())
cv2.imwrite('mask.png', mask)
mask = np.array(mask, np.uint8)
contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
# print(contours)
# img = img.cpu()
img = np.array(img, np.uint8)
cv2.drawContours(img, contours, -1, (0, 0, 255), 1)
cv2.imwrite('result2.png', img)
boxes = []
return boxes
@staticmethod
def demo_visualize():
pass
def inference(self, img_path, is_visualize=True, is_format_output=False):
img = cv2.imread(img_path, cv2.COLOR_BGR2RGB)
img = Image.fromarray(img).convert("RGB")
# img = Image.open(img_path).convert("RGB")
# print('222', np.array(img))
# img = self.pad_sample(img)
img = self.resize_img(img)
# print('333', img.size)
# print('-----', np.array(img))
ori_img = img
img.save('img.png')
# img = [img]
print('111', np.array(img))
img = self.transform(img)
print('222', np.array(img))
img = img.unsqueeze(0)
img = img.to(self.device)
# print('1111', img.size())
# print(img)
# print(img)
with torch.no_grad():
s1 = time.time()
preds = self.model(img)
print(preds)
s2 = time.time()
print(s2 - s1)
# boxes, scores = SegDetectorRepresenter().represent(pred=preds, height=h, width=w, is_output_polygon=False)
boxes = self.post_process(preds, ori_img)
class PG(object):
def __init__(self, configs, env):
self.configs = configs
self.env = env
self.action_size = (64, 1024)
# n_channels=3 for RGB images
# n_classes is the number of probabilities you want to get per pixel
# - For 1 class and background, use n_classes=1
# - For 2 classes, use n_classes=1
# - For N > 2 classes, use n_classes=N
# TODO now I assume input<->output size are equal, which might not be true, so we need some modifications onto Unet if necesary
self.actor = UNet(
n_channels=3, n_classes=1, bilinear=True
) # [B,C, H_in=372, W_in=1242] -> [B, C, H_out=64, W_out=1024]
self.optimizer = Adam(self.actor.parameters(), lr=configs['lr'])
self.actor.to(device)
def get_action(self, state, deterministic=False):
"""Given the state, produces an action, the probability of the action, the log probability of the action, and
the argmax action"""
action_probabilities = self.actor(
state) # output size should be [B*H*W]
action_probabilities = torch.sigmoid(
action_probabilities) # make sure the probs are in range [0,1]
# B, _, _, _ = action_probabilities.shape
action_probabilities = action_probabilities[:, :, :self.
action_size[0], :self.
action_size[1]]
action_probabilities = torch.squeeze(action_probabilities, 1)
# assert action_probabilities.size()[1, 2] == self.action_size, "Actor output the wrong size"
# action_probabilities_flat = action_probabilities.contiguous().view(B, -1)
# TODO leave this to future process; seems it will get the index
max_probability_action = torch.argmax(action_probabilities, dim=-1)
if deterministic:
# using deteministic policy during test time
action = action_probabilities(action_probabilities > 0.5).cpu()
else:
# using stochastic policy during traning time
action_distribution = Bernoulli(
action_probabilities
) # this creates a distribution to sample from
action = action_distribution.sample().cpu(
) # sample the discrete action and copy it to cpu
# Have to deal with situation of 0.0 probabilities because we can't do log 0
z = action_probabilities == 0.0
z = z.float() * 1e-8
log_action_probabilities = torch.log(action_probabilities + z)
return action, action_probabilities, log_action_probabilities, max_probability_action
def compute_loss(self, obs, act, rew):
"""make loss function whose gradient, for the right data, is policy gradient"""
# TODO we may do not need to calculate it for the second time.
act_baseline, _, logp, _ = self.get_action(obs, deterministic=True)
# advantage
_, rew_baseline, _, _ = self.env.step(act_baseline, obs=obs)
advantage = rew.to(device).float() - rew_baseline.to(device).float()
loss = logp * Variable(advantage).expand_as(act)
loss = loss.mean()
return loss
def update(self, batch_obs, batch_acts, batch_rews):
"""take a single policy gradient update step for a batch"""
self.optimizer.zero_grad()
batch_loss = self.compute_loss(
obs=torch.as_tensor(batch_obs, dtype=torch.float32),
act=torch.as_tensor(batch_acts, dtype=torch.int32),
rew=torch.as_tensor(batch_rews, dtype=torch.int32),
)
batch_loss.backward()
self.optimizer.step()
return batch_loss
output_dir = dir_names.valout_dir + '/' + experiment_name + '/'
model_file = model_dir + 'model_20.pth'
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if not load_model:
c.force_create(model_dir)
c.force_create(tfboard_dir)
num_class = 3
# load unet model
if not load_model:
net = UNet(num_class=num_class)
net = net.to(device)
else:
net = UNet(num_class=num_class)
net.load_state_dict(
torch.load(model_file, map_location=torch.device(device)))
net = net.to(device)
net.eval()
#Initialize class to convert labels to color images
color_transform = Colorize()
#Set up data
#Define image dataset (reads in full images and segmentations)
image_dataset = p.ImageDataset(csv_file=c.train_val_csv)
# Split dataset into train and validation
