def getBBoxes(PILImg):
"""
get boxes given a PIL image.
:param imagePath: PIL Image
:return: ndarray of bboxes
"""
config = InferenceConfig()
config.display()
# Create model object in inference mode.
model = modellib.MaskRCNN(mode="inference", model_dir='mask_rcnn_coco.hy', config=config)
# Load weights trained on MS-COCO
#model.load_weights('mask_rcnn_balloon.h5', by_name=True)
model.load_weights('mask_rcnn_coco.h5', by_name=True)#, exclude=[ "mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"])
# COCO Class names
class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
'kite', 'baseball bat', 'baseball glove', 'skateboard',
'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
'teddy bear', 'hair drier', 'toothbrush']
# Load a random image from the images folder
#image = skimage.io.imread('coco_2014_test/COCO_val2014_000000060623.jpg')
#image = skimage.io.imread(imagePath)
image = np.asarray(PILImg)
# original image
#plt.figure(figsize=(12,10))
#skimage.io.imshow(image)
# Run detection
results = model.detect([image], verbose=1)
# Visualize results
r = results[0]
#print(type(r['rois']))
#print(r['rois'])
return r['rois']
#visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores'])
def train(self):
MODEL_DIR = "weights/mrnn/"
# Local path to trained weights file
COCO_MODEL_PATH = "../mask_rcnn_coco.h5"
config = RSNAConfig()
config.display()
assert config.IMAGE_SHAPE[0] == config.IMAGE_SHAPE[1]
dataset_train = RSNADataset()
dataset_train.load_shapes(self.images_path, self.train_csv_file, config.IMAGE_SHAPE[0])
dataset_train.prepare()
# Validation dataset
dataset_val = RSNADataset()
dataset_val.load_shapes(self.images_path, "data/validation.csv", config.IMAGE_SHAPE[0])
dataset_val.prepare()
# 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)
model = modellib.MaskRCNN(mode="training", config=config,
model_dir=MODEL_DIR)
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(), by_name=True)
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE,
epochs=50,
layers='heads')
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE / 10,
epochs=50,
layers="all")
def test(self):
MODEL_DIR = "weights/mrnn/"
inference_config = InferenceConfig()
# Recreate the model in inference mode
model = modellib.MaskRCNN(mode="inference",
config=inference_config,
model_dir=MODEL_DIR)
# Get path to saved weights
# Either set a specific path or find last trained weights
# model_path = os.path.join(ROOT_DIR, ".h5 file name here")
model_path = model.find_last()
# Load trained weights
print("Loading weights from ", model_path)
model.load_weights(model_path, by_name=True)
config = RSNAConfig()
# dataset_train = RSNADataset()
# dataset_train.load_shapes(self.images_path, self.train_csv_file, config.IMAGE_SHAPE[0])
# dataset_train.prepare()
dataset_val = RSNADataset()
dataset_val.load_shapes(self.test_images_path, "data/test.csv", config.IMAGE_SHAPE[0])
dataset_val.prepare()
for image_id in dataset_val.image_ids:
info = dataset_val.image_info[image_id]
image_pid = (info['path'].replace("\\", "/").split("/")[-1])[0:-4]
original_image, image_meta, gt_class_id, gt_bbox, gt_mask = \
modellib.load_image_gt(dataset_val, inference_config,
image_id, use_mini_mask=False)
# visualize.display_instances(original_image, gt_bbox, gt_mask, gt_class_id,
# dataset_train.class_names, figsize=(8, 8))
results = model.detect([original_image], verbose=1)
r = results[0]
r['rois'] = r['rois'][r['class_ids'] == 1]
r['scores'] = r['scores'][r['class_ids'] == 1]
r['masks'] = r['masks'][:,:,r['class_ids'] == 1]
r['class_ids'] = r['class_ids'][r['class_ids'] == 1]
# visualize.display_instances(original_image, r['rois'], r['masks'], r['class_ids'],
# dataset_val.class_names, r['scores'], figsize=(8, 8))
# plt.show()
if r['masks'].shape[2] == 0:
r['masks'] = np.zeros((256, 256))
else:
r['masks'] = r['masks'][:,:] * r['scores']
r['masks'] = np.sum(r['masks'], axis=-1)
r['masks'][r['masks'] > 1] = 1
r['masks'] = r['masks'] * 255
annot = Image.fromarray(np.uint8(r['masks']))
annot.save("preds/" + image_pid + ".png")
h = 0
rois = detections['rois']
masks = detections['masks']
class_ids = detections['class_ids']
scores = detections['class_ids']
mask = np.zeros_like(frame_bgr)
mask_p = np.zeros(frame_bgr.shape[:-1], dtype=np.bool)
mask_n = np.zeros(frame_bgr.shape[:-1], dtype=np.bool)
for i, roi in enumerate(rois):
if class_ids[i] == class_names.index('airplane'):
mask_p = np.bitwise_or(masks[:, :, i], mask_p)
else:
mask_n = np.bitwise_or(masks[:, :, i], mask_n)
mask[:, :, 1] = mask_p.astype(np.uint8) * 255
mask[:, :, 2] = mask_n.astype(np.uint8) * 255
overlap = cv2.addWeighted(frame_bgr, 0.5, mask, 0.5, 20)
font = cv2.FONT_HERSHEY_SIMPLEX
for i, roi in enumerate(rois):
y1, x1, y2, x2 = roi
cv2.rectangle(overlap, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(overlap, class_names[class_ids[i]], (x1, y1 + 25), font, 1,
(255, 255, 255), 2, cv2.LINE_AA)
return overlap
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
model.load_weights(COCO_MODEL_PATH, by_name=True)
model.keras_model._make_predict_function()
predict(np.zeros((512, 512, 3), dtype=np.uint8))
print('> Done')