def load_model(model_path, **argkws):
model_path = Path(model_path)
bn = model_path.parent.name
parts = bn.split('_')
model_type = parts[1]
loss_type = parts[2]
state = torch.load(model_path, map_location='cpu')
model = get_model(model_type, 1, 1, loss_type, **argkws)
model.load_state_dict(state['state_dict'])
model.eval()
return model, state['epoch']
def load_model(model_path,
flow_subdir = 'validation',
data_root_dir = None,
flow_type = None,
data_type = None,
**argkws
):
model_path = Path(model_path)
bn = model_path.parent.name
data_type_bn, _, remain = bn.partition('+F')
flow_type_bn, _, remain = remain.partition('+roi')
model_name, _, remain = remain.partition('_')[-1].partition('_')
remain = remain.split('_')
loss_type = remain[0]
if data_type is None:
data_type = data_type_bn
if flow_type is None:
flow_type = flow_type_bn
flow_args = flow_types[flow_type]
data_args = data_types[data_type]
if data_root_dir is None:
data_root_dir = data_args['root_data_dir']
state = torch.load(model_path, map_location = 'cpu')
n_ch_in = data_args['n_ch_in']
n_ch_out = data_args['n_ch_out']
model = get_model(model_name, n_ch_in, n_ch_out, loss_type, **argkws)
model.load_state_dict(state['state_dict'])
model.eval()
data_flow = CoordFlow(data_root_dir / flow_subdir,
**flow_args,
is_preloaded = False
)
return model, data_flow, state['epoch']
def load_model(model_path, **argkws):
model_path = Path(model_path)
bn = model_path.parent.name
data_type, _, remain = bn.partition('+F')
flow_type, _, remain = remain.partition('+roi')
model_name, _, remain = remain.partition('_')[-1].partition('_')
remain = remain.split('_')
loss_type = remain[0]
state = torch.load(model_path, map_location='cpu')
data_args = data_types[data_type]
n_ch_in = data_args['n_ch_in']
n_ch_out = data_args['n_ch_out']
model = get_model(model_name, n_ch_in, n_ch_out, loss_type)
model.load_state_dict(state['state_dict'])
model.eval()
return model, state['epoch']
def train(data_type='woundhealing-v2-mix',
flow_type=None,
model_name='unet-simple',
use_classifier=False,
loss_type='l1smooth-G1.5',
cuda_id=0,
log_dir=None,
batch_size=256,
n_epochs=2000,
save_frequency=200,
num_workers=0,
root_data_dir=None,
optimizer_name='adam',
lr_scheduler_name='',
lr=1e-5,
weight_decay=0.0,
momentum=0.9,
roi_size=64,
is_preloaded=False,
hard_mining_freq=None,
model_path_init=None,
train_samples_per_epoch=40960):
data_args = data_types[data_type]
dflt_root_data_dir = data_args['root_data_dir']
n_ch_in = data_args['n_ch_in']
n_ch_out = data_args['n_ch_out']
if flow_type is None:
flow_type = data_args['dflt_flow_type']
flow_args = flow_types[flow_type]
if log_dir is None:
if 'log_prefix' in data_args:
log_dir = LOG_DIR_DFLT / data_args['log_prefix'] / data_type
else:
log_dir = LOG_DIR_DFLT / data_type
if root_data_dir is None:
root_data_dir = dflt_root_data_dir
root_data_dir = Path(root_data_dir)
train_dir = root_data_dir / 'train'
test_dir = root_data_dir / 'validation'
if '-merged' in data_type:
flow_func = CoordFlowMerged
else:
flow_func = CoordFlow
print(flow_func)
print(root_data_dir)
train_flow = flow_func(train_dir,
samples_per_epoch=train_samples_per_epoch,
roi_size=roi_size,
**flow_args,
is_preloaded=is_preloaded)
val_flow = flow_func(test_dir,
roi_size=roi_size,
**flow_args,
is_preloaded=is_preloaded)
model = get_model(model_name, n_ch_in, n_ch_out, loss_type)
if model_path_init is not None:
model_name += '-pretrained'
state = torch.load(model_path_init, map_location='cpu')
model.load_state_dict(state['state_dict'])
device = get_device(cuda_id)
optimizer = get_optimizer(optimizer_name,
model,
lr=lr,
momentum=momentum,
weight_decay=weight_decay)
lr_scheduler = get_scheduler(lr_scheduler_name, optimizer)
date_str = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
hard_mining_str = '' if hard_mining_freq is None else f'+hard-neg-{hard_mining_freq}'
lr_scheduler_name = '+' + lr_scheduler_name if lr_scheduler_name else ''
save_prefix = f'{data_type}+F{flow_type}+roi{roi_size}{hard_mining_str}_{model_name}_{loss_type}_{date_str}'
save_prefix = f'{save_prefix}_{optimizer_name}{lr_scheduler_name}_lr{lr}_wd{weight_decay}_batch{batch_size}'
train_locmax(save_prefix,
model,
device,
train_flow,
val_flow,
optimizer,
lr_scheduler=lr_scheduler,
log_dir=log_dir,
batch_size=batch_size,
num_workers=num_workers,
hard_mining_freq=hard_mining_freq,
n_epochs=n_epochs,
save_frequency=save_frequency)
scale_int=(0, 4095),
prob_unseeded_patch=0.5,
zoom_range=(0.97, 1.03),
int_aug_offset=(-0.2, 0.2),
int_aug_expansion=(0.7, 1.3),
samples_per_epoch=batch_size * 100)
gen = CoordFlow(root_dir, **flow_args)
loader = DataLoader(gen,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
collate_fn=collate_simple)
model = get_model('ind+clf+unet-simple', 3, 2, 'maxlikelihood')
for images, targets in tqdm.tqdm(loader):
images = torch.from_numpy(np.stack(images)).to(device)
targets = [{
k: torch.from_numpy(v).to(device)
for k, v in target.items()
} for target in targets]
#%%
model.train()
losses = model(images, targets)
loss = sum([v for v in losses.values()])
loss.backward()
#%%
# model_type = 'unet-simple'
assert model_path.exists()
#img_src_dir = Path.home() / 'workspace/localization/data/woundhealing/raw/'
img_src_dir = Path.home(
) / 'workspace/localization/data/woundhealing/HitDataJan2020/'
save_dir = Path.home(
) / 'workspace/localization/predictions/woundhealing_experiments' / bn
n_ch_in = 1
n_ch_out = 1
device = get_device(cuda_id)
model = get_model(model_type,
n_ch_in,
n_ch_out,
loss_type,
nms_threshold_abs=nms_threshold_abs,
nms_threshold_rel=nms_threshold_rel)
state = torch.load(model_path, map_location='cpu')
model.load_state_dict(state['state_dict'])
model.eval()
model = model.to(device)
#%%
img_paths = [
x for x in img_src_dir.rglob('*.tif') if not x.name.startswith('.')
]
for img_path in tqdm.tqdm(img_paths):
img = cv2.imread(str(img_path), -1)
