/
train.py
107 lines (85 loc) · 3.14 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
# coding: utf-8
# hacked up from the shapes example file.
import sys
sys.path.append("Mask_RCNN")
import os
import random
import math
import time
import numpy as np
import random
import glob
import skimage
import osmmodelconfig
from config import Config
import imagestoosm.config as osmcfg
import utils
import model as modellib
from model import log
# Root directory of the project
#ROOT_DIR = os.getcwd()
ROOT_DIR = os.path.dirname(os.path.realpath(sys.argv[0]))
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# Path to COCO trained weights
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
config = osmmodelconfig.OsmModelConfig()
config.ROOT_DIR = ROOT_DIR
config.display()
fullTrainingDir = os.path.join( ROOT_DIR, osmcfg.trainDir,"*")
fullImageList = []
for imageDir in glob.glob(fullTrainingDir):
if ( os.path.isdir( os.path.join( fullTrainingDir, imageDir) )):
id = os.path.split(imageDir)[1]
fullImageList.append( id)
random.shuffle(fullImageList)
cutoffIndex = int(len(fullImageList)*.75)
trainingImages = fullImageList[0:cutoffIndex ]
validationImages = fullImageList[cutoffIndex:-1 ]
# Training dataset
dataset_train = osmmodelconfig.OsmImagesDataset(ROOT_DIR)
dataset_train.load(trainingImages, config.IMAGE_SHAPE[0], config.IMAGE_SHAPE[1])
dataset_train.prepare()
# Validation dataset
dataset_val = osmmodelconfig.OsmImagesDataset(ROOT_DIR)
dataset_val.load(validationImages, config.IMAGE_SHAPE[0], config.IMAGE_SHAPE[1])
dataset_val.prepare()
# Create model in training mode
model = modellib.MaskRCNN(mode="training", config=config,
model_dir=MODEL_DIR)
# Which weights to start with?
init_with = "coco" # imagenet, coco, or last
if init_with == "imagenet":
model.load_weights(model.get_imagenet_weights(), by_name=True)
elif init_with == "coco":
# Load weights trained on MS COCO, but skip layers that
# are different due to the different number of classes
# See README for instructions to download the COCO weights
model.load_weights(COCO_MODEL_PATH, by_name=True,
exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
"mrcnn_bbox", "mrcnn_mask"])
elif init_with == "last":
# Load the last model you trained and continue training
print(model.find_last()[1])
model.load_weights(model.find_last()[1], by_name=True)
if ( init_with != "last") :
# Training - Stage 1
# Adjust epochs and layers as needed
print("Training network heads")
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE,
epochs=10,
layers='heads')
# Training - Stage 2
# Finetune layers from ResNet stage 4 and up
print("Training Resnet layer 3+")
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE/10,
epochs=100,
layers='3+')
# Finetune layers from ResNet stage 3 and up
print("Training all")
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE / 100,
epochs=1000,
layers='all')