-
Notifications
You must be signed in to change notification settings - Fork 4
/
standard_maskrcnn.py
executable file
·164 lines (121 loc) · 6.29 KB
/
standard_maskrcnn.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
import datetime
import torch
import os
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.engine.defaults import DefaultPredictor
from detectron2.evaluation.evaluator import inference_on_dataset
from detectron2.data.datasets import load_coco_json
from detectron2.data.datasets.coco import convert_to_coco_json
from detectron2.utils import logger
from mask2former.config import add_maskformer2_config
from utils.utils import *
from utils.custom_trainers import *
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
#########################
### PROGRAM VARIABLES ###
#########################
DO_TRAIN = True # Whether to train the model
DO_TEST = True # Whether to test the model
RESUME = False # If True, resumes from last_checkpoint
DATASET_FILENAME = 'data/prescaled/coco_annotation.json' # COCO style annotation file
IMAGE_DIR = './data/prescaled' # Folder where dataset was downloaded
SPLIT_FRAC = 0.8 # Fraction of data used in train
CONFIG_FILE = "configs/config_standard_maskrcnn.yaml" # Detectron2 style config file
TEST_WEIGHTS = "" # Checkpoint to use for testing
INITIAL_WEIGHTS = None # Path to a previous checkpoint to finetune or None
SHOW_ANNOTATIONS = False # Display images and labels before training
SHOW_AUGMENTATIONS = False # Display augmented images and labels before training
#########################
def init_maskrcnn(config_file, train_dicts, test_dicts, output_dir, fold_num=None, initial_weights=None):
# Configure Model - See all parameters here : https://detectron2.readthedocs.io/en/latest/modules/config.html#yaml-config-references
cfg = get_cfg()
add_maskformer2_config(cfg)
cfg.merge_from_file(config_file)
cfg.OUTPUT_DIR = output_dir
# Initialize network weights (to finetune)
if initial_weights is not None:
cfg.MODEL.WEIGHTS = initial_weights
# Rename dataset if using k-fold cross-validation
if fold_num is not None:
cfg.DATASETS.TRAIN = [cfg.DATASETS.TRAIN[0] + f"_fold{fold_num}"]
cfg.DATASETS.TEST = [cfg.DATASETS.TEST[0] + f"_fold{fold_num}"]
# Initialize Dataloaders
DatasetCatalog.register(cfg.DATASETS.TRAIN[0], lambda d=cfg.DATASETS.TRAIN[0]: train_dicts)
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).set(thing_classes=["log"])
DatasetCatalog.register(cfg.DATASETS.TEST[0], lambda d=cfg.DATASETS.TEST[0]: test_dicts)
MetadataCatalog.get(cfg.DATASETS.TEST[0]).set(thing_classes=["log"])
# Save Model config in output folder
os.makedirs(output_dir, exist_ok=True)
with open(f"{output_dir}/config.yaml", "w") as f:
f.write(cfg.dump())
# Copy annotation files to output folder (for backups)
convert_to_coco_json(cfg.DATASETS.TRAIN[0], os.path.join(output_dir, "annos_train.json"))
convert_to_coco_json(cfg.DATASETS.TEST[0], os.path.join(output_dir, "annos_test.json"))
return cfg
def train_maskrcnn(cfg, resume=False, show_annos=False, show_augs=False):
########################
#### TRAINING LOGIC ####
########################
# Visualize annotations
if show_annos:
dicts = DatasetCatalog.get(cfg.DATASETS.TRAIN[0])
visualize_annotations(dicts, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]))
# Visualize data augmentation
if show_augs:
train_loader = StandardMaskRCNNTrainer.build_train_loader(cfg)
visualize_data_augmentation(train_loader, MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), image_format=cfg.INPUT.FORMAT)
# Initialize trainer
trainer = StandardMaskRCNNTrainer(cfg)
trainer.resume_or_load(resume)
trainer.train()
# Save model (use half the space)
checkpointer = DetectionCheckpointer(trainer.model, save_dir=cfg.OUTPUT_DIR)
checkpointer.save("model_final") # save to output folder
def test_maskrcnn(cfg, test_weights):
#######################
#### TESTING LOGIC ####
#######################
if os.path.isfile(test_weights):
cfg.MODEL.WEIGHTS = test_weights
# Create Predictor
predictor = DefaultPredictor(cfg)
# Evaluate Predictions on Test Dataset
evaluator = StandardMaskRCNNTrainer.build_evaluator(cfg, cfg.DATASETS.TEST[0])
val_loader = StandardMaskRCNNTrainer.build_test_loader(cfg, cfg.DATASETS.TEST[0])
print(inference_on_dataset(predictor.model, val_loader, evaluator))
# Visualize Predictions
logs_metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
dicts_test = DatasetCatalog.get(cfg.DATASETS.TEST[0])
visualize_predictions(predictor, dicts_test, logs_metadata)
if __name__ == "__main__":
print('GPU available :', torch.cuda.is_available())
print('Torch version :', torch.__version__, '\n')
logger.setup_logger(name=__name__)
torch.autograd.set_detect_anomaly(False)
torch.autograd.profiler.profile(False)
torch.autograd.profiler.emit_nvtx(False)
# Init output folder
if (RESUME):
# Take latest run in outputs folder
output_dir = np.sort([x for x in os.listdir("./outputs/") if os.path.isdir("./outputs/"+x)])[-1]
output_dir = os.path.join("./outputs", output_dir)
config_file = os.path.join(output_dir, "config.yaml")
else:
output_dir = f'./outputs/{datetime.datetime.now().strftime("%Y-%m-%d_%H-%M")}'
random_state = np.random.RandomState(42)
dataset_dicts = np.array(load_coco_json(DATASET_FILENAME, IMAGE_DIR))
# Split train/test
dataset_idx = list(range(len(dataset_dicts)))
train_idx = random_state.choice(dataset_idx, int(SPLIT_FRAC * len(dataset_idx)), replace=False)
test_idx = list(set(dataset_idx) - set(train_idx))
train_dicts = dataset_dicts[train_idx]
test_dicts = dataset_dicts[test_idx]
cfg = init_maskrcnn(CONFIG_FILE, train_dicts, test_dicts, output_dir, initial_weights=INITIAL_WEIGHTS)
if DO_TRAIN:
train_maskrcnn(cfg, RESUME, SHOW_ANNOTATIONS, SHOW_AUGMENTATIONS)
TEST_WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") # Ensures trained network will be used for testing
if DO_TEST:
test_maskrcnn(cfg, TEST_WEIGHTS)