def load_model(model_path, flow_subdir = 'validation', data_root_dir = None, nms_args = None, flow_type = None, data_type = None):
model_path = Path(model_path)
bn = model_path.parent.name
data_type_bn, _, remain = bn.partition('+F')
if data_type is None:
data_type = data_type_bn
flow_type_bn, _, remain = remain.partition('+roi')
if flow_type is None:
flow_type = flow_type_bn
model_name, _, remain = remain.partition('unet-')[-1].partition('_')
model_name = 'unet-' + model_name
remain = remain.split('_')
loss_type = remain[0]
if nms_args is None:
if 'reg' in loss_type:
nms_args = dict(nms_threshold_abs = 0.4, nms_threshold_rel = None)
else:
nms_args = dict(nms_threshold_abs = 0.0, nms_threshold_rel = 0.2)
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_args = model_types[model_name]
model_args['unet_pad_mode'] = 'reflect'
model = CellDetectorWithClassifier(**model_args,
unet_n_inputs = n_ch_in,
unet_n_ouputs = n_ch_out,
loss_type = loss_type,
return_belive_maps = True,
**nms_args
)
model.load_state_dict(state['state_dict'])
model.eval()
if data_root_dir is None:
data_root_dir = data_args['root_data_dir']
flow_args = flow_types[flow_type]
data_flow = CoordFlow(data_root_dir / flow_subdir,
**flow_args,
is_preloaded = False
)
return model, data_flow, state['epoch']
def main(
data_type='limphocytes-20x',
flow_type='limphocytes',
root_model_dir=Path.home() /
'workspace/localization/results/locmax_detection/limphocytes/20x/limphocytes-20x',
checkpoint_name='model_best.pth.tar',
max_dist=10,
thresh2check=np.arange(0.05, 1, 0.05),
cuda_id=0,
):
device = get_device(cuda_id)
data_args = data_types[data_type]
root_data_dir = data_args['root_data_dir']
flow_args = flow_types[flow_type]
root_model_dir = Path(root_model_dir)
train_flow = CoordFlow(root_data_dir / 'train',
**flow_args,
is_preloaded=True)
val_flow = CoordFlow(root_data_dir / 'validation',
**flow_args,
is_preloaded=True)
model_paths = root_model_dir.rglob(checkpoint_name)
model_paths = list(model_paths)
#model_paths = [x for x in model_paths if 'maxlikelihood' in x.parent.name]
results = {}
for model_path in tqdm.tqdm(model_paths):
res = get_scores_with_best_threshold(model_path, device, thresh2check,
max_dist, train_flow, val_flow)
bn = model_path.parent.name
results[bn] = res
save_name = root_model_dir / 'scores.p'
with open(save_name, 'wb') as fid:
pickle.dump(results, fid)
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']
root_dir = '/Users/avelinojaver/OneDrive - Nexus365/bladder_cancer_tils/eosinophils/training/20x/'
num_workers = 4
batch_size = 4 #512
gauss_sigma = 1.5
device = 'cpu'
flow_args = dict(roi_size=16,
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]
def train(data_type='woundhealing-v2-mix',
flow_type=None,
model_name='unet-simple',
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,
num_folds=None,
val_fold_id=None,
val_dist=5):
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)
if root_data_dir.is_dir():
assert val_fold_id is None
train_dir = root_data_dir / 'train'
test_dir = root_data_dir / 'validation'
print(root_data_dir)
train_flow = CoordFlow(train_dir,
samples_per_epoch=train_samples_per_epoch,
roi_size=roi_size,
**flow_args,
is_preloaded=is_preloaded)
val_flow = CoordFlow(test_dir,
roi_size=roi_size,
**flow_args,
is_preloaded=is_preloaded)
else:
data_type += f'fold-{val_fold_id}-{num_folds}'
train_fold_ids = [
x + 1 for x in range(num_folds) if x + 1 != val_fold_id
]
train_flow = CoordFlow(root_data_dir,
samples_per_epoch=train_samples_per_epoch,
roi_size=roi_size,
**flow_args,
folds2include=train_fold_ids,
num_folds=num_folds,
is_preloaded=True)
val_flow = CoordFlow(root_data_dir,
roi_size=roi_size,
**flow_args,
folds2include=val_fold_id,
num_folds=num_folds,
is_preloaded=True)
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,
val_dist=val_dist)
def train(
data_type = 'heba',
model_name = 'retinanet-resnet34',
loss_type = 'focal',
cuda_id = 0,
log_dir = None,
batch_size = 16,
num_workers = 1,
root_data_dir = None,
lr = 1e-5,
**argkws
):
if log_dir is None:
log_dir = log_dir_root_dflt / 'bbox_detection' / data_type
dflts = data_types_dflts[data_type]
root_data_dir = dflts['root_data_dir']
roi_size = dflts['roi_size']
min_radius = dflts['min_radius']
scale_int = dflts['scale_int']
num_classes = dflts['num_classes']
if root_data_dir is None:
root_data_dir = root_data_dir
train_dir = root_data_dir / 'train'
test_dir = root_data_dir / 'validation'
bbox_encoder = BoxEncoder(img_size = (roi_size, roi_size),
pyramid_levels = [1, 2, 3, 4, 5],
aspect_ratios = [(1.,1.)]
)
print(root_data_dir)
train_flow = CoordFlow(train_dir,
samples_per_epoch = 20480,
bbox_encoder = bbox_encoder,
min_radius = min_radius,
roi_size = roi_size,
scale_int = scale_int
)
val_flow = CoordFlow(test_dir,
samples_per_epoch = 640,
bbox_encoder = bbox_encoder,
min_radius = min_radius,
roi_size = roi_size,
scale_int = scale_int
)
if model_name.startswith('retinanet'):
backbone = model_name.split('-')[1]
model = RetinaNet(backbone = backbone,
is_fpn_small = True,
num_classes = num_classes,
num_anchors = 3)
device = get_device(cuda_id)
loss_type = f'{loss_type}-{num_classes}'
criterion = get_loss(loss_type)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
date_str = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
save_prefix = f'{data_type}_{model_name}_{loss_type}_{date_str}'
save_prefix = f'{save_prefix}_adam_lr{lr}_batch{batch_size}'
train_bboxes(save_prefix,
model,
device,
train_flow,
val_flow,
criterion,
optimizer,
log_dir = log_dir,
batch_size = batch_size,
num_workers = num_workers,
**argkws
)