def __init__(self) :
super(GANCascadedUNet,self).__init__()
self.layer_1 = UNet(1,1, segmentation = False)
self.layer_2 = UNet(2,4, segmentation = False)
self.layer_3 = UNet(6,1, segmentation = False)
self.layer_4 = UNet(2,4, segmentation = False)
self.layer_5 = UNet(6,1, segmentation = False)
def start():
model = UNet(num_channels=1, num_classes=2)
criterion = DICELossMultiClass()
dir_path = './Data/images/'
sizes = []
file_names = []
for f in os.listdir(dir_path):
im = Image.open(os.path.join(dir_path, f))
print(im.size)
sizes.append(im.size)
file_names.append(f)
print(sizes)
print(file_names)
test_dataset = TestDataset('./Data/',
im_size=[256, 256],
transform=tr.ToTensor())
test_loader = DataLoader(test_dataset,
batch_size=4,
shuffle=False,
num_workers=1)
print("Test Data : ", len(test_loader.dataset))
model.load_state_dict(
torch.load('unet-model-16-100-0.001', map_location='cpu'))
test_only(model, test_loader, criterion, sizes, file_names)
def load_net_cifar(model_loc):
""" Make a model
Network must be saved in the form model_name-depth, where this is a unique identifier
"""
model_file = Path(model_loc).name.rsplit('_')[0]
model_name = model_file.split('-')[0]
print('Loading model_file', model_file)
if (model_name == 'lenet'):
model = LeNet(10, 32)
elif (model_name == 'vggnet'):
model = VGG(int(model_file.split('-')[1]), 10, 32)
elif (model_name == 'resnet'):
model = ResNet(int(model_file.split('-')[1]), 10, 32)
elif (model_name == 'wide'):
model = Wide_ResNet(
model_file.split('-')[2][0:2],
model_file.split('-')[2][2:4], 0, 10, 32)
elif (model_name == 'unet'):
# this is a terrible way to do it but should be saving class, den parameters to load them individually
loss = DenoiseLoss(n=1, hard_mining=0, norm=False)
classifier = ResNet(50, num_classes=10, IM_SIZE=32)
if (Path(model_loc).name.rsplit('_')[1] == 'model'):
denoiser = UNet(3, 3, stochastic=False)
elif (Path(model_loc).name.rsplit('_')[1] == 'stochastic'):
denoiser = UNet(3, 3, stochastic=True)
model = DenoiseNet(classifer=classifier, denoiser=denoiser, loss=loss)
else:
print(
'Error : Network should be either [LeNet / VGGNet / ResNet / Wide_ResNet / unet'
)
sys.exit(0)
model.load_state_dict(torch.load(model_loc)['state_dict'])
return model
def main(opt):
writer = SummaryWriter()
log_dir = writer.get_logdir()
os.makedirs(os.path.join(log_dir, "images"), exist_ok=True)
os.makedirs(os.path.join(log_dir, "test"), exist_ok=True)
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
# Initialize generator and discriminator
generator = UNet(opt.sample_num, opt.channels, opt.batch_size, opt.alpha)
discriminator = Discriminator(opt.batch_size, opt.alpha)
generator.to(device=device)
discriminator.to(device=device)
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(),
lr=opt.lr_g,
betas=(opt.b1, opt.b2))
optimizer_D = torch.optim.Adam(discriminator.parameters(),
lr=opt.lr_d,
betas=(opt.b1, opt.b2))
if opt.mode == 'train':
generator = train(writer, log_dir, device, generator, discriminator,
optimizer_G, optimizer_D, opt)
test(opt, log_dir, generator=generator)
if opt.mode == 'test':
test(opt, log_dir)
test_moving(opt, log_dir)
def main():
# Define the net
net = UNet(in_channels=Dataset.in_channels,
out_channels=Dataset.out_channels) #.cuda()
criterion = DiceLoss()
optimizer = optim.Adam(net.parameters(), lr=1e-4)
# Parpered the dataset
dataloader = torch.utils.data.DataLoader(
Dataset('/home/mooziisp/GitRepos/unet/data/membrane/train/image'),
batch_size=1,
shuffle=True,
num_workers=2,
pin_memory=True)
# TODO validation
# TODO acc and loss record
# TODO intergated with tensorboard
for epoch in range(100):
for i, (images, targets) in enumerate(dataloader):
#images, targets = images.cuda(), targets.cuda()
preds = net(images)
loss = criterion(preds, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f'{epoch*30+i}it, loss: {loss.item():.3f}')
return net
def __init__(self, proxy_map, opts,startup_check=False):
super(SpecificWorker, self).__init__(proxy_map)
self.Period = 50
if startup_check:
self.startup_check()
else:
self.timer.timeout.connect(self.compute)
self.timer.start(self.Period)
self.opts = opts # Copy user options
self.wait_frames = 5 # Wait for these many frames to get stored before calling gaitset (see gaitSet/pretreatment.py)
# Dictonary to store mapping between tracking id and segmentation mask
self.id2mask = {}
# Segmentation module
segmentation_model = UNet(backbone="resnet18", num_classes=2)
segmentation_model.load_state_dict(torch.load("./src/PretrainedModels/UNet_ResNet18.pth", map_location="cpu")['state_dict'], strict=False)
self.segmentation_inference = ImageListInference(segmentation_model,opts)
# Gait feature extraction
self.gait_model = SetNet(256)
self.gait_model.load_state_dict(torch.load("./src/PretrainedModels/GaitSet_CASIA-B_73_False_256_0.2_128_full_30-80000-encoder.ptm",map_location="cpu"),strict=False)
if self.opts.gpu >= 0:
self.gait_model.cuda()
self.gait_model.eval()
# Variables to store data for computation
self.lock = False # Lock from taking any more input until features are calculated
self.input_image_list = [] # Store images in this list before performing segmentation
self.input_tracking_id = [] # Store respective tracking id here
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 __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 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 __init__(self, hparams):
super(TimeTransfer, self).__init__()
self.hparams = hparams
# networks
self.unet = UNet(3, hparams.hidden_dim)
self.data_dir = os.path.expanduser(hparams.data_dir)
self.split_indices = prefix_sum(hparams.data_split)
self.device = torch.device('cuda' if hparams.gpus > 0 else 'cpu')
self.criteria = PerceptualLoss().to(self.device)
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 test(opt, log_dir, generator=None):
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
if generator == None:
generator = UNet(opt.sample_num, opt.channels, opt.batch_size,
opt.alpha)
checkpoint = torch.load(opt.load_model, map_location=device)
generator.load_state_dict(checkpoint['g_state_dict'])
del checkpoint
torch.cuda.empty_cache()
generator.to(device)
generator.eval()
dataloader = torch.utils.data.DataLoader(MyDataset_test(opt),
opt.batch_size,
shuffle=True,
num_workers=0)
for i, (imgs, filename) in enumerate(dataloader):
with torch.no_grad():
test_img = generator(imgs.to(device))
filename = filename[0].split('/')[-1]
filename = "test/" + filename + '.png'
test_img = convert_im(test_img,
os.path.join(log_dir, filename),
nrow=5,
normalize=True,
save_im=True)
def test_moving(opt, log_dir, generator=None):
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
if generator == None:
generator = UNet(opt.sample_num, opt.channels, opt.batch_size,
opt.alpha)
checkpoint = torch.load(opt.load_model, map_location=device)
generator.load_state_dict(checkpoint['g_state_dict'])
del checkpoint
torch.cuda.empty_cache()
generator.to(device)
generator.eval()
dataloader = torch.utils.data.DataLoader(
MyDataset_test_moving(opt),
opt.batch_size,
shuffle=True,
num_workers=opt.num_workers_dataloader)
for i, (imgs, filename) in enumerate(dataloader):
with torch.no_grad():
filename = filename[0].split('/')[-1]
for k in range(len(imgs)):
test_img = generator(imgs[k].to(device))
folder_path = os.path.join(log_dir, "test/%s" % filename)
os.makedirs(folder_path, exist_ok=True)
filename_ = filename + '_' + str(k) + '.png'
test_img = convert_im(test_img,
os.path.join(folder_path, filename_),
nrow=5,
normalize=True,
save_im=True)
def train(device, model_path, dataset_path):
"""
Trains the network according on the dataset_path
"""
network = UNet(1, 3).to(device)
optimizer = torch.optim.Adam(network.parameters())
criteria = torch.nn.MSELoss()
dataset = GrayColorDataset(dataset_path, transform=train_transform)
loader = torch.utils.data.DataLoader(dataset,
batch_size=16,
shuffle=True,
num_workers=cpu_count())
if os.path.exists(model_path):
network.load_state_dict(torch.load(model_path))
for _ in tqdm.trange(10, desc="Epoch"):
network.train()
for gray, color in tqdm.tqdm(loader, desc="Training", leave=False):
gray, color = gray.to(device), color.to(device)
optimizer.zero_grad()
pred_color = network(gray)
loss = criteria(pred_color, color)
loss.backward()
optimizer.step()
torch.save(network.state_dict(), model_path)
def denoise_3d(mat, model=None, filepath='/home/jupyter/notebooks/checkpoints/3d_denoise.pt', return_detached=True,
batch_size=5000, device=torch.device('cuda')):
if model is None:
model = UNet(in_channels=1, num_classes=1, out_channels=[4, 8, 16], num_conv=2, n_dim=3,
kernel_size=[3, 3, 3], same_shape=True).to(device)
model.load_state_dict(torch.load(filepath))
with torch.no_grad():
num_batches = (mat.size(0) + batch_size - 1)//batch_size
mat = torch.cat([model(mat[batch_size*i:batch_size*(i+1)]) for i in range(num_batches)], dim=0)
if return_detached:
mat = mat.detach()
for k in [k for k in locals().keys() if k!='mat']:
del locals()[k]
torch.cuda.empty_cache()
return mat
def get_model(m_config, d_config, args):
if m_config.unet:
kwargs = {
'n_channels':
3,
'dropout':
m_config.dropout,
'deeper':
not (d_config.dataset == 'CIFAR10'
or d_config.dataset == 'CELEBA'),
'ngf':
m_config.ngf
}
return UNet(**kwargs).to(args.device)
if d_config.dataset == "FFHQ":
model = NCSNv2Deepest
elif d_config.dataset == 'LSUN':
model = NCSNv2Deeper
else:
model = NCSNv2
sigmoid = args.target == 'dae'
model_args = [m_config, d_config, sigmoid, args.no_dilation, args.std]
return model(*model_args).to(args.device)
def get_model(hparams):
if hparams['model'] == 'discrete':
return C4UNet(hparams['in_channels'], hparams['out_channels'],
group_type=hparams['group'],
N=hparams['N'],
n_features=hparams['n_features'],
loss_type=hparams['loss_func'],
lr=hparams['lr'])
elif hparams['model'] == 'standard':
return UNet(hparams['in_channels'], hparams['out_channels'], n_features=hparams['n_features'], loss_type=hparams['loss_func'], lr=hparams['lr'])
elif hparams['model'] == 'harmonic':
return HarmonicUNet(hparams['in_channels'],
hparams['out_channels'],
n_features=hparams['n_features'],
group_type=hparams['group'],
max_freq=hparams['max_freq'],
loss_type=hparams['loss_func'],
lr=hparams['lr'])
elif hparams['model'] == 'steerable':
gspace = gspaces.Rot2dOnR2(-1, maximum_frequency=hparams['max_freq'])
return SteerableCNN(gspace,
hparams['in_channels'],
hparams['out_channels'],
n_blocks=hparams['n_blocks'],
n_features=hparams['n_features'],
irrep_type=hparams['irrep_type'],
loss_type=hparams['loss_func'],
lr=hparams['lr'])
else:
raise ValueError(f'Unsupported model type: {hparams["model"]}')
def prior_update_gan_noise(particles, inputs, model_no):
netG = net.G(5, 3, 64)
netG.load_state_dict(
torch.load(
'%s/netG_epoch_%d.pth' %
('/home/gtx1080/Abduallah/pix2pix.pytorch/imglog/fluid_noise',
model_no)))
netG.eval()
netG.cuda()
mu = np.load(
'/home/gtx1080/Sync/Kun/30_min_ele/non_ele_whole/train/mu.npy')
var = np.load(
'/home/gtx1080/Sync/Kun/30_min_ele/non_ele_whole/train/var.npy')
particles = np.concatenate((particles, inputs), axis=1)
particles, label = normalization_test(particles, particles[:, 0:3, :, :],
mu, var)
particles = torch.tensor(particles, dtype=torch.float).cuda()
with torch.no_grad():
particles = netG(particles)
particles = np.array(particles.cpu())
particles[:,
0, :, :] = particles[:, 0, :, :] * np.sqrt(var[5, 0]) + mu[5, 0]
particles[:,
1, :, :] = particles[:, 1, :, :] * np.sqrt(var[6, 0]) + mu[6, 0]
particles[:,
2, :, :] = particles[:, 2, :, :] * np.sqrt(var[7, 0]) + mu[7, 0]
particles[:, 0, :, :] = np.maximum(particles[:, 0, :, :], 0)
return particles
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 __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 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 == 'MDeNet':
print('Mine MDeNet..................')
model = MDeNet(num_classes=num_classes)
elif model_type == 'EncDec':
print('Mine EncDec..................')
model = EncDec(num_classes=num_classes)
elif model_type == 'hourglass':
model = hourglass(num_classes=num_classes)
elif model_type == 'MDeNetplus':
print('load MDeNetplus..................')
model = MDeNetplus(num_classes=num_classes)
elif model_type == 'UNet':
model = UNet(num_classes=num_classes)
else:
print('I am here')
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 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 denoise_trace(trace, model=None, filepath='/home/jupyter/notebooks/checkpoints/denoise_trace.pt', return_detached=True,
device=torch.device('cuda')):
if model is None:
model = UNet(in_channels=1, num_classes=1, out_channels=[8, 16, 32], num_conv=2,
n_dim=1, kernel_size=3).to(device)
model.load_state_dict(torch.load(filepath))
with torch.no_grad():
mean = trace.mean()
std = trace.std()
pred = model((trace-mean)/std)
pred = model(pred)
pred = pred * std + mean
if return_detached:
pred = pred.detach()
for k in [k for k in locals().keys() if k!='pred']:
del locals()[k]
torch.cuda.empty_cache()
return pred
def attention_map(mat, model=None, filepath='/home/jupyter/notebooks/checkpoints/segmentation_count_hardmask.pt',
batch_size=5000, return_detached=True, device=torch.device('cuda')):
if model is None:
model = UNet(in_channels=1, num_classes=1, out_channels=[4, 8, 16], num_conv=2, n_dim=3,
kernel_size=[3, 3, 3], same_shape=True).to(device)
model.load_state_dict(torch.load(filepath))
nrow, ncol = mat.shape[1:]
if batch_size*nrow*ncol > 1e7:
batch_size = int(1e7 / (nrow*ncol))
with torch.no_grad():
num_batches = (mat.size(0) + batch_size - 1)//batch_size
mat = torch.cat([model(mat[batch_size*i:batch_size*(i+1)]) for i in range(num_batches)], dim=0).mean(0)
if return_detached:
mat = mat.detach()
for k in [k for k in locals().keys() if k!='mat']:
del locals()[k]
torch.cuda.empty_cache()
return mat
def __init__(self, train_dl, val_dl):
self.device = ('cuda:0' if torch.cuda.is_available() else 'cpu')
self.train_dl = train_dl
self.val_dl = val_dl
self.loss = Loss()
self.net = UNet(1).to(self.device)
self.net.apply(Model._init_weights)
self.criterion = self.loss.BCEDiceLoss
self.optim = None
self.scheduler = None
self._init_optim(LR, BETAS)
self.cycles = 0
self.hist = {'train': [], 'val': [], 'loss': []}
utils.create_dir('./pt')
utils.log_data_to_txt('train_log', f'\nUsing device {self.device}')
def run(conf, data):
score = ConfusionMatrix(num_classes)
for file in os.listdir(data['splits_dir']):
conf['log_key'] = file
split = du.load_split_json(data['splits_dir'] + sep + file)
model = UNet(conf['input_channels'], conf['num_classes'], reduce_by=2)
optimizer = optim.Adam(model.parameters(), lr=conf['learning_rate'])
if conf['distribute']:
model = torch.nn.DataParallel(model)
model.float()
optimizer = optim.Adam(model.module.parameters(), lr=conf['learning_rate'])
try:
trainer = KernelTrainer(run_conf=conf, model=model, optimizer=optimizer)
if conf.get('mode') == 'train':
train_loader = KernelDataset.get_loader(shuffle=True, mode='train', transforms=transforms,
images=split['train'], data_conf=data,
run_conf=conf)
validation_loader = KernelDataset.get_loader(shuffle=True, mode='validation',
transforms=transforms,
images=split['validation'], data_conf=data,
run_conf=conf)
print('### Train Val Batch size:', len(train_loader), len(validation_loader))
# trainer.resume_from_checkpoint(parallel_trained=conf.get('parallel_trained'), key='latest')
trainer.train(train_loader=train_loader, validation_loader=validation_loader)
test_loader = KernelDataset.get_loader(shuffle=False, mode='test', transforms=transforms,
images=split['test'], data_conf=data, run_conf=conf)
trainer.resume_from_checkpoint(checkpoint_file= conf.get("checkpoint_file"), parallel_trained=conf.get('parallel_trained'), key='best')
trainer.test(data_loader=test_loader, global_score=score, logger=trainer.test_logger)
except Exception as e:
traceback.print_exc()
with open(conf.get('log_dir', 'net_logs') + os.sep + 'global_score.txt', 'w') as lg:
lg.write(f'{score.precision()},{score.recall()},{score.f1()},{score.accuracy()}')
lg.flush()
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 main():
device = torch.device("cuda")
inputPath = "D:/VolumeSuperResolution-InputData/sparse-rendering/cleveland60/sample00004_sparse.npy"
input = torch.from_numpy(np.load(inputPath)).to(device=device,
dtype=torch.float32)
print("input:", input.shape)
B, C, H, W = input.shape
mask = torch.abs(input[:, 0:1, :, :])
print("Filled pixels:", end='')
for b in range(B):
print(" %d" % int(torch.sum(mask[b])), end='')
print(" of %d pixels" % (H * W))
print(
"Now test different u-net configuations and report the filled pixels for the first batch"
)
wf = 2
with torch.no_grad():
for depth in range(1, 10):
unet = UNet(in_channels=C,
depth=depth,
wf=wf,
padding='partial',
return_masks=True)
unet.to(device)
output, masks = unet(input, mask)
print("Depth:", depth)
print(" output shape:", output.shape)
print(" Layer: shape + filled pixels")
for i, m in enumerate(masks):
print(
" ", m.shape, " -> %d of %d pixels" %
(int(torch.sum(m[0])), m.shape[2] * m.shape[3]))
print("Save mask")
img = mask[0, 0].cpu().numpy()
plt.imsave("unet-test-mask.png", img, cmap='Greys')
def build_model(model_type):
if model_type == 'unet':
model = UNet(num_classes=2)
if model_type == 'dcan':
model = UNet_DCAN(num_classes=2)
if model_type == 'dmtn':
model = UNet_DMTN(num_classes=2)
if model_type == 'psinet':
model = PsiNet(num_classes=2)
if model_type == 'convmcd':
model = UNet_ConvMCD(num_classes=2)
return model
