forked from perlfloccri/NuclearSegmentationPipeline
-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathtrain.py
178 lines (151 loc) · 6.22 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
175
176
177
178
import os
import sys
sys.path.append(r'E:\NuclearSegmentationPipeline\Config')
from Datahandling.Datasets import TisquantDatasetNew,ArtificialNucleiDataset,ArtificialNucleiDatasetNotConverted,MergedDataset, TisquantDataset
from Config.Config import UNETSettings
import random
import math
import re
import time
import numpy as np
import cv2
import matplotlib
import matplotlib.pyplot as plt
from scipy.io import loadmat
import skimage.transform as ski_transform
from skimage.measure import label
from utils import compute_final_mask
from config import Config
import utils
import model as modellib
import visualize
from model import log
#from Datasets import TisquantDataset
input_shape = [3, 256, 256]
# Root directory of the project
ROOT_DIR = r"E:\NuclearSegmentationPipeline\DeepLearningArchitectures\MaskRCNN"#os.getcwd()
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# Local path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
# Download COCO trained weights from Releases if needed
if not os.path.exists(COCO_MODEL_PATH):
utils.download_trained_weights(COCO_MODEL_PATH)
class NucleiConfig(Config):
"""Configuration for training on the toy shapes dataset.
Derives from the base Config class and overrides values specific
to the toy shapes dataset.
"""
# Give the configuration a recognizable name
#NAME = "Nuclei"
NAME = UNETSettings().network_info["net_description"]
# Train on 1 GPU and 8 images per GPU. We can put multiple images on each
# GPU because the images are small. Batch size is 8 (GPUs * images/GPU).
GPU_COUNT = 1
IMAGES_PER_GPU = 8
# Number of classes (including background)
NUM_CLASSES = 1 + 1 # background + Nucleus
# Use small images for faster training. Set the limits of the small side
# the large side, and that determines the image shape.
IMAGE_MIN_DIM = 256
IMAGE_MAX_DIM = 256
# Use smaller anchors because our image and objects are small
RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128) # anchor side in pixels
#RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128, 192)
BACKBONE_STRIDES = [4, 8, 16, 32, 64]
#BACKBONE_STRIDES = [4, 8, 16, 32, 64, 128]
# Reduce training ROIs per image because the images are small and have
# few objects. Aim to allow ROI sampling to pick 33% positive ROIs.
TRAIN_ROIS_PER_IMAGE = 32
# Use a small epoch since the data is simple
STEPS_PER_EPOCH = 100
# use small validation steps since the epoch is small
VALIDATION_STEPS = 5
config = NucleiConfig()
config.display()
def get_ax(rows=1, cols=1, size=8):
"""Return a Matplotlib Axes array to be used in
all visualizations in the notebook. Provide a
central point to control graph sizes.
Change the default size attribute to control the size
of rendered images
"""
_, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
return ax
"""
# Training dataset
dataset_train = TisquantDataset()
val_idx = dataset_train.load_data(width=config.IMAGE_SHAPE[0], height=config.IMAGE_SHAPE[1])
dataset_train.prepare()
# Validation dataset
dataset_val = TisquantDataset()
dataset_val.load_data(width=config.IMAGE_SHAPE[0], height=config.IMAGE_SHAPE[1],ids=val_idx,mode=2)
dataset_val.prepare()
"""
# Load Dataset
print ("Load dataset ...")
if UNETSettings().network_info["dataset"] == 'tisquant': #args.dataset
dataset= TisquantDatasetNew()
elif UNETSettings().network_info["dataset"] == 'artificialNuclei':
dataset = ArtificialNucleiDataset()
elif UNETSettings().network_info["dataset"] == 'artificialNucleiNotConverted':
dataset = ArtificialNucleiDatasetNotConverted()
elif UNETSettings().network_info["dataset"] == 'mergeTisquantArtificial':
datasets = []
dataset1 = TisquantDatasetNew()
dataset1.load_data(mode=1)
dataset2 = ArtificialNucleiDataset()
dataset2.load_data(mode=1)
datasets.append(dataset1)
datasets.append(dataset2)
dataset = MergedDataset(datasets)
else:
print('Dataset not valid')
sys.exit("Error")
# Load Dataset
dataset.load_data(mode=1)
dataset.prepare()
dataset_train, dataset_val = dataset.split_train_test()
images, labels = dataset_train.to_numpy()
test_images, test_labels = dataset_val.to_numpy()
"""
# Load and display random samples
image_ids = np.random.choice(dataset_train.image_ids, 4)
for image_id in image_ids:
image = dataset_train.load_image(image_id)
mask, class_ids = dataset_train.load_mask(image_id)
visualize.display_top_masks(image, mask, class_ids, dataset_train.class_names)
"""
# 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
model.load_weights(model.find_last()[1], by_name=True)
# Train the head branches
# Passing layers="heads" freezes all layers except the head
# layers. You can also pass a regular expression to select
# which layers to train by name pattern.
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE,
epochs=UNETSettings().network_info["max_epochs"],
layers='heads')
# Fine tune all layers
# Passing layers="all" trains all layers. You can also
# pass a regular expression to select which layers to
# train by name pattern.
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE / 10,
epochs=UNETSettings().network_info["max_epochs"],
layers="all")