forked from facebookresearch/detectron2
-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
174 lines (154 loc) · 6.9 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
174
from detectron2.engine import DefaultTrainer
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
import os
import sys
from detectron2.data.datasets import bottle_loader
import detectron2
from detectron2.utils.logger import setup_logger
import csv
setup_logger()
# import some common libraries
import numpy as np
import random
import shutil
# import some common detectron2 utilities
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2
def gen_cfg_train(model, weights, dataset):
cfg = get_cfg()
cfg.merge_from_file("./configs/COCO-Detection/" + model)
cfg.DATASETS.TRAIN = (dataset + '_train',)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = "detectron2://COCO-Detection/" + os.path.splitext(model)[0] + '/' + weights # initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 3000 # 300 iterations seems good enough, but you can certainly train longer
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4 # only has one class (ballon)
cfg.OUTPUT_DIR = 'output_' + dataset
return cfg
def gen_cfg_test(dataset, model, dataset_name):
#cfg = gen_cfg_train(model, weights, dataset)
cfg = get_cfg()
#cfg.merge_from_file("./configs/COCO-Detection/" + model)
cfg.OUTPUT_DIR = 'output_' + dataset
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9 # set the testing threshold for this model
cfg.DATASETS.TEST = (dataset_name + '_test', )
cfg.TEST.DETECTIONS_PER_IMAGE = 1
return cfg
def train_model(path, model, weights, dataset, action_type='train', mode="full"):
bottle_loader.register_dataset(path, dataset, action_type, mode)
cfg = gen_cfg_train(model, weights, dataset)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
def test_model(path, model, weights, dataset, action_type='test', mode="full", visualize=False):
dataset_name = os.path.basename(path)
test = bottle_loader.register_dataset(path, dataset_name, action_type, mode)
bottle_loader.register_dataset(path, dataset, 'train', mode)
cfg_test = gen_cfg_test(dataset, model, dataset_name)
cfg = gen_cfg_train(model, weights, dataset)
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
evaluator = COCOEvaluator("%s_%s" % (dataset_name, action_type), cfg_test, False, output_dir="./output_%s/" % (dataset))
val_loader = build_detection_test_loader(cfg_test, "%s_%s" % (dataset, 'train'))
result = inference_on_dataset(trainer.model, val_loader, evaluator)
#Visualize the test
if visualize:
visualize_images_dict(dataset_name, test, MetadataCatalog.get('%s_%s' % (dataset, 'train')), cfg, dataset_name)
return result
def visualize_cfg(cfg, dataset):
cfg.DATASETS.TEST = (dataset + '_test', )
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.9 # set the testing threshold for this model
predictor = DefaultPredictor(cfg)
return predictor
def visualize_images_dict(folder, dict_data, bottle_metadata, cfg, dataset_name):
path = os.path.join(cfg.OUTPUT_DIR, folder)
if os.path.isdir(path):
shutil.rmtree(path)
os.mkdir(path)
dataset_dicts = dict_data
predictor = visualize_cfg(cfg, dataset_name)
for d in dataset_dicts:
im = cv2.imread(d["file_name"])
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
metadata=bottle_metadata,
scale=1.0 # remove the colors of unsegmented pixels
)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
image = v.get_image()[:, :, ::-1]
v_gt = Visualizer(image[:,:,::-1],
metadata=bottle_metadata,
scale=1.5)
v_gt = v_gt.draw_dataset_dict(d)
image = v_gt.get_image()[:,:,::-1]
cv2.imwrite(os.path.join(path, os.path.basename(d['file_name'])), image)
v_gt = Visualizer(im[:,:,::-1],
metadata=bottle_metadata,
scale=1.5)
v_gt = v_gt.draw_dataset_dict(d)
image = v_gt.get_image()[:,:,::-1]
cv2.imwrite(os.path.join(path, 'gt_' + os.path.basename(d['file_name'])), image)
# def visualize_images_dict(folder, dict_data, bottle_metadata, cfg):
# path = os.path.join(cfg.OUTPUT_DIR, folder)
# if os.path.isdir(path):
# shutil.rmtree(path)
# os.mkdir(path)
# dataset_dicts = dict_data
# predictor = visualize_cfg(cfg)
# for d in dataset_dicts:
# im = cv2.imread(d["file_name"])
# outputs = predictor(im)
# v = Visualizer(im[:, :, ::-1],
# metadata=bottle_metadata,
# scale=0.8 # remove the colors of unsegmented pixels
# )
# v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
# image = v.get_image()[:, :, ::-1]
# cv2.imwrite(os.path.join(path, 'instance_' + os.path.basename(d['file_name'])), image)
# v = Visualizer(image[:, :, ::-1],
# metadata=bottle_metadata,
# scale=1.0)
# v = v.draw_dataset_dict(d)
# image = v.get_image()[:, :, ::-1]
# #Draw the ground truth as well:
# cv2.imwrite(os.path.join(path, os.path.basename(d['file_name'])), image)
def main(args):
run = args[1]
path = args[2]
if run == 'train':
train_model(args[2], args[3], args[4], args[5], mode=args[6])
elif run == 'test':
csv_file = None
if len(args) > 7:
csv_file = args[7]
epoch = 0
if len(args) > 8:
epoch = args[8]
result = test_model(args[2], args[3], args[4], args[5], mode=args[6])
if csv_file is not None:
#TODO: Print to row of csv
values = result['bbox']
AP50 = values['AP50']
AP75 = values['AP75']
AP = values['AP']
print([AP50, AP75, AP])
with open(os.path.join(os.getcwd(), csv_file), 'a', newline='') as file:
writer = csv.writer(file)
writer.writerows([[epoch, AP50, AP75, AP]])
if __name__ == "__main__":
main(sys.argv)
#train_model('faster_rcnn_R_50_C4_3x.yaml', '137849393/model_final_f97cb7.pkl', 'bottle')