-
Notifications
You must be signed in to change notification settings - Fork 2
/
train.py
173 lines (153 loc) · 8.98 KB
/
train.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
import argparse
import albumentations as albu
from torch.cuda import device_count
from torch.utils.data import DataLoader
import segmentation_models_pytorch as smp
from tools.models import SegmentationModel
from tools.datasets import SegmentationDataset
from tools.supervisely_tools import read_supervisely_project
from tools.data_processing import split_data, get_logging_labels
from tools.utils import DynamicWeighting, StaticWeighting
def main(args):
img_paths, ann_paths, dataset_names = read_supervisely_project(sly_project_dir=args.dataset_dir,
included_datasets=args.included_datasets,
excluded_datasets=args.excluded_datasets)
subsets = split_data(img_paths=img_paths,
ann_paths=ann_paths,
dataset_names=dataset_names,
class_name=args.class_name,
seed=11,
ratio=args.ratio,
normal_datasets=['rsna_normal', 'chest_xray_normal'])
preprocessing_params = smp.encoders.get_preprocessing_params(encoder_name=args.encoder_name,
pretrained=args.encoder_weights)
augmentation_params = albu.Compose([
albu.CLAHE(p=0.2),
albu.RandomSizedCrop(min_max_height=(int(0.7 * args.input_size[0]), int(0.9 * args.input_size[0])),
height=args.input_size[0],
width=args.input_size[1],
w2h_ratio=1.0,
p=0.2),
albu.Rotate(limit=15, p=0.5),
albu.HorizontalFlip(p=0.5),
albu.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.2)
])
datasets = {}
for subset_name in subsets:
_augmentation_params = augmentation_params if subset_name == 'train' else None
dataset = SegmentationDataset(img_paths=subsets[subset_name][0],
ann_paths=subsets[subset_name][1],
input_size=args.input_size,
class_name=args.class_name,
augmentation_params=_augmentation_params,
transform_params=preprocessing_params)
datasets[subset_name] = dataset
# Used only for augmentation debugging
# import cv2
# import torch
# import numpy as np
# idx = 0
# img_tensor, mask_tensor = datasets['train'][idx]
# mean = np.array(preprocessing_params['mean'])
# std = np.array(preprocessing_params['std'])
# _img = ((img_tensor.permute(1, 2, 0).cpu().detach().numpy() * std) + mean) * 255
# img = _img.astype(np.uint8)
# _mask = (torch.squeeze(mask_tensor).cpu().detach().numpy()) * 255
# mask = _mask.astype(np.uint8)
# cv2.imwrite('img.png', img)
# cv2.imwrite('mask.png', mask)
# If debug is frozen, use num_workers = 0
num_workers = 0 * device_count()
train_loader = DataLoader(datasets['train'], batch_size=args.batch_size, num_workers=num_workers, shuffle=True)
val_loader = DataLoader(datasets['val'], batch_size=args.batch_size, num_workers=num_workers)
test_loader = DataLoader(datasets['test'], batch_size=args.batch_size, num_workers=num_workers)
# Use all images from the logging folder without exclusion
img_paths_logging, ann_paths_logging, dataset_names_logging = read_supervisely_project(sly_project_dir=args.logging_dir,
included_datasets=None,
excluded_datasets=None)
logging_dataset = SegmentationDataset(img_paths=img_paths_logging,
ann_paths=ann_paths_logging,
input_size=args.input_size,
class_name=args.class_name,
augmentation_params=None,
transform_params=preprocessing_params)
logging_loader = DataLoader(logging_dataset, batch_size=1, num_workers=num_workers)
aux_params = None
if args.use_cls_head:
aux_params = dict(pooling='avg',
dropout=0.20,
activation='sigmoid',
classes=1)
if not args.use_cls_head:
args.loss_cls = None
# weights_strategy = StaticWeighting(w1=1.0, w2=1.0)
weights_strategy = DynamicWeighting(alpha=0.05)
model = SegmentationModel(model_name=args.model_name,
encoder_name=args.encoder_name,
encoder_weights=args.encoder_weights,
aux_params=aux_params,
batch_size=args.batch_size,
epochs=args.epochs,
class_name=args.class_name,
loss_seg=args.loss_seg,
loss_cls=args.loss_cls,
weights_strategy=weights_strategy,
optimizer=args.optimizer,
lr=args.lr,
es_patience=args.es_patience,
es_min_delta=args.es_min_delta,
monitor_metric=args.monitor_metric,
input_size=args.input_size,
save_dir=args.save_dir,
logging_labels=get_logging_labels([args.class_name]),
wandb_project_name=args.wandb_project_name,
wandb_api_key=args.wandb_api_key)
model.train(train_loader, val_loader, test_loader, logging_loader)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Segmentation pipeline')
parser.add_argument('--dataset_dir', default='dataset/covid_segmentation_single_crop', type=str,
help='dataset/covid_segmentation, '
'dataset/covid_segmentation_single_crop, '
'dataset/covid_segmentation_double_crop,'
'dataset/lungs_segmentation')
parser.add_argument('--included_datasets', default=None, type=str)
parser.add_argument('--excluded_datasets', default=None, type=str)
parser.add_argument('--ratio', nargs='+', default=(0.8, 0.1, 0.1), type=float, help='train, val, and test sizes')
parser.add_argument('--model_name', default='Unet', type=str, help='Unet, Unet++, DeepLabV3, DeepLabV3+, FPN, Linknet, PSPNet, PAN and MAnet')
parser.add_argument('--input_size', nargs='+', default=(512, 512), type=int)
parser.add_argument('--encoder_name', default='resnet18', type=str)
parser.add_argument('--encoder_weights', default='imagenet', type=str, help='imagenet, ssl or swsl')
parser.add_argument('--batch_size', default=8, type=int)
parser.add_argument('--loss_seg', default='Dice', type=str, help='Dice, Jaccard, BCE, BCEL, Lovasz or Focal')
parser.add_argument('--loss_cls', default='SL1', type=str, help='BCE, SL1 or L1')
parser.add_argument('--optimizer', default='Adam', type=str, help='SGD, Adam, AdamW, RMSprop, Adam_amsgrad or AdamW_amsgrad')
parser.add_argument('--lr', default=0.001, type=float)
parser.add_argument('--es_patience', default=10, type=int)
parser.add_argument('--es_min_delta', default=0.01, type=float)
parser.add_argument('--monitor_metric', default='f1_seg', type=str)
parser.add_argument('--epochs', default=30, type=int)
parser.add_argument('--use_cls_head', action='store_true')
parser.add_argument('--save_dir', default='models', type=str)
parser.add_argument('--wandb_project_name', default=None, type=str)
parser.add_argument('--wandb_api_key', default='b45cbe889f5dc79d1e9a0c54013e6ab8e8afb871', type=str)
args = parser.parse_args()
# Used only for debugging
# args.excluded_datasets = [
# 'covid-chestxray-dataset',
# 'COVID-19-Radiography-Database',
# 'Figure1-COVID-chestxray-dataset',
# 'rsna_normal',
# 'chest_xray_normal'
# ]
if 'covid' in args.dataset_dir:
args.class_name = 'COVID-19'
args.wandb_project_name = 'covid_segmentation' if not isinstance(args.wandb_project_name, str) else args.wandb_project_name
args.logging_dir = args.dataset_dir + '_logging'
elif 'lungs' in args.dataset_dir:
args.class_name = 'Lungs'
args.wandb_project_name = 'lungs_segmentation' if not isinstance(args.wandb_project_name, str) else args.wandb_project_name
args.logging_dir = args.dataset_dir + '_logging'
else:
raise ValueError('There is no class name for dataset {:s}'.format(args.dataset_dir))
print('\nW&B project name: {:s}'.format(args.wandb_project_name))
main(args)