def train():
config = MyMaskRcnnConfig()
config.display()
dataset_train, dataset_val = get_datasets()
# 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":
if not os.path.isfile(COCO_MODEL_PATH):
utils.download_trained_weights(COCO_MODEL_PATH)
# 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(train_dataset=dataset_train,
val_dataset=dataset_val,
learning_rate=config.LEARNING_RATE,
epochs=10,
layers='heads')
model.keras_model.save_weights(os.path.join(MODEL_DIR, "stage1.h5"),
overwrite=True)
# Training - Stage 2
# Finetune layers from ResNet stage 4 and up
print("Training Resnet layer 3+")
model.train(train_dataset=dataset_train,
val_dataset=dataset_val,
learning_rate=config.LEARNING_RATE / 10,
epochs=100,
layers='3+')
model.keras_model.save_weights(os.path.join(MODEL_DIR, "stage2.h5"),
overwrite=True)
class Predictor:
config = MyMaskRcnnConfig()
class InferenceConfig(config.__class__):
# Run detection on one image at a time
GPU_COUNT = 1
IMAGES_PER_GPU = 1
IMAGE_MIN_DIM = 256
IMAGE_MAX_DIM = 256
def __init__(self, weights_path: str):
if not os.path.isfile(weights_path):
raise RuntimeError(
"Weights cannot be found at: {}".format(weights_path))
self.weights_path = weights_path
self._model = None
def predict_array(self, img_data: np.ndarray, extent=None, do_rectangularization=True, tile=None) \
-> List[List[Tuple[int, int]]]:
if not self._model:
print("Loading model")
inference_config = self.InferenceConfig()
# Create model in training mode
model = modellib.MaskRCNN(mode="inference",
config=inference_config,
model_dir="log")
model.load_weights(self.weights_path, by_name=True)
self._model = model
model = self._model
print("Predicting...")
res = model.detect([img_data], verbose=1)
print("Prediction done")
print("Extracting contours...")
point_sets = self._get_contours(masks=res[0]['masks'])
print("Contours extracted")
point_sets_mapped = []
col, row = tile
for points in point_sets:
pp = list(
map(lambda p: (p[0] + col * 256, p[1] + row * 256), points))
if pp:
point_sets_mapped.append(pp)
point_sets = point_sets_mapped
rectangularized_outlines = []
if do_rectangularization:
point_sets = list(
map(lambda point_set: rectangularize(point_set), point_sets))
if not extent:
rectangularized_outlines = point_sets
else:
for o in point_sets:
georeffed = georeference(o, extent)
if georeffed:
rectangularized_outlines.append(georeffed)
return rectangularized_outlines
def predict_path(self,
img_path: str,
extent=None) -> List[List[Tuple[int, int]]]:
img = Image.open(img_path)
data = np.asarray(img, dtype="uint8")
return self.predict_array(img_data=data, extent=extent)
@staticmethod
def _get_contours(masks: np.ndarray) -> List[List[Tuple[int, int]]]:
contours: List[List[Tuple[int, int]]] = []
for i in range(masks.shape[-1]):
mask = masks[:, :, i]
if mask.any():
conts = find_contours(mask, 0.5)
for c in conts:
rr, cc = polygon_perimeter(c[:, 0],
c[:, 1],
shape=mask.shape,
clip=False)
points = tuple(zip(cc, rr))
contours.append(points)
return contours
class Predictor:
config = MyMaskRcnnConfig()
class InferenceConfig(config.__class__):
# Run detection on one image at a time
GPU_COUNT = 1
# IMAGES_PER_GPU = 30
IMAGES_PER_GPU = 1
IMAGE_MIN_DIM = 256
IMAGE_MAX_DIM = 256
def __init__(self, weights_path: str):
if not os.path.isfile(weights_path):
raise RuntimeError("Weights cannot be found at: {}".format(weights_path))
self.weights_path = weights_path
self._model = None
def predict_arrays(self, images: List[Tuple[np.ndarray, str]], verbose=1) -> List:
batch_size = 30
if not self._model:
print("Loading model")
inference_config = self.InferenceConfig()
inference_config.BATCH_SIZE = batch_size
inference_config.IMAGES_PER_GPU = batch_size
print("Predicting {} images".format(len(images)))
# Create model in training mode
model = modellib.MaskRCNN(mode="inference", config=inference_config, model_dir="log")
model.load_weights(self.weights_path, by_name=True)
self._model = model
model = self._model
batches = int(math.ceil(len(images) / batch_size))
all_point_sets = []
for i in range(batches):
start = i * batch_size
end = start + batch_size
img_with_id_batch = images[start:end]
if len(img_with_id_batch) < batch_size:
inference_config = self.InferenceConfig()
inference_config.BATCH_SIZE = len(img_with_id_batch)
inference_config.IMAGES_PER_GPU = len(img_with_id_batch)
model = modellib.MaskRCNN(mode="inference", config=inference_config, model_dir="log")
model.load_weights(self.weights_path, by_name=True)
print("Predicting batch {}/{}".format(i+1, batches))
img_batch = list(map(lambda i: i[0], img_with_id_batch))
id_batch = list(map(lambda i: i[1], img_with_id_batch))
results = model.detect(img_batch, image_ids=id_batch, verbose=verbose)
print("Extracting contours...")
for res in results:
point_sets = get_contours(masks=res['masks'], classes=res['class_ids'])
image_id = res['coco_id']
for points, class_name in point_sets:
all_point_sets.append((points, image_id, class_name))
print("Contours extracted")
K.clear_session()
return all_point_sets
def predict_path(self, img_path: str, verbose=1) -> List[List[Tuple[int, int]]]:
return self.predict_paths([img_path], verbose=verbose)
def predict_paths(self, all_paths: List[str], verbose=1) -> List[List[Tuple[int, int]]]:
all_images = []
for p in all_paths:
data = cv2.imread(p)
coco_img_id = int(os.path.basename(p).replace(".jpg", ""))
all_images.append((data, coco_img_id))
return self.predict_arrays(images=all_images, verbose=verbose)
class Predictor:
config = MyMaskRcnnConfig()
class InferenceConfig(config.__class__):
# Run detection on one image at a time
GPU_COUNT = 1
IMAGES_PER_GPU = 1
IMAGE_MIN_DIM = IMAGE_WIDTH
IMAGE_MAX_DIM = IMAGE_WIDTH
def __init__(self, weights_path: str):
if not os.path.isfile(weights_path):
raise RuntimeError(
"Weights cannot be found at: {}".format(weights_path))
self.weights_path = weights_path
self._model = None
def predict_array(self, img_data: np.ndarray, extent=None, do_rectangularization=True, tile=None) \
-> List[List[Tuple[int, int]]]:
if not tile:
tile = (0, 0)
if not self._model:
print("Loading model")
inference_config = self.InferenceConfig()
# Create model in training mode
model = modellib.MaskRCNN(mode="inference",
config=inference_config,
model_dir="log")
model.load_weights(self.weights_path, by_name=True)
self._model = model
model = self._model
print("Predicting...")
res = model.detect([img_data], verbose=1)
print("Prediction done")
print("Extracting contours...")
point_sets = get_contours(masks=res[0]['masks'])
point_sets = list(map(lambda point_set: list(point_set), point_sets))
print("Contours extracted")
rectangularized_outlines = []
if do_rectangularization:
point_sets = list(
map(lambda point_set: rectangularize(point_set), point_sets))
point_sets_mapped = []
col, row = tile
for points in point_sets:
pp = list(
map(lambda p: (p[0] + col * 256, p[1] + row * 256), points))
if pp:
point_sets_mapped.append(pp)
point_sets = point_sets_mapped
if not extent:
rectangularized_outlines = point_sets
else:
for o in point_sets:
georeffed = georeference(o, extent)
if georeffed:
rectangularized_outlines.append(georeffed)
return rectangularized_outlines
def predict_path(self,
img_path: str,
extent=None) -> List[List[Tuple[int, int]]]:
img = Image.open(img_path)
data = np.asarray(img, dtype="uint8")
return self.predict_array(img_data=data, extent=extent)