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))
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, 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 i, res in enumerate(results):
masks = res['masks']
for i in range(masks.shape[-1]):
mask = masks[:, :, i]
segmentation = cocomask.encode(
np.asfortranarray(mask, dtype=np.uint8))
score = 1
if len(res['scores'] > i):
score = res['scores'][i]
coco_id = res['coco_id']
point_sets.append((segmentation, score, coco_id))
print("Contours extracted")
return point_sets
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)
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 __init__(self):
"""Initalise a new instance of the MaskInterface class.
This is the constructor for the MaskInterface class which also
initalises a new instance of the Mask R-CNN network.
Args:
None
"""
# Load Mask_RCNN
config = InferenceConfig()
coco_path = os.path.join(os.getcwd(), "Mask_RCNN", "mask_rcnn_coco.h5")
self.model = modellib.MaskRCNN(mode="inference",
config=config,
model_dir="./logs")
self.model.load_weights(coco_path, by_name=True)
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 save_masked_images(self):
# Import COCO config
ROOT_DIR = os.path.abspath("./")
sys.path.append(ROOT_DIR) # To find local version of the librar
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# Local path to trained weights file
COCO_MODEL_PATH = ROOT_DIR + '/mask_rcnn_coco.h5'
IMAGE_DIR = self.image_path
OUTPUT_DIR = self.masked_image_path
class InferenceConfig(coco.CocoConfig):
# Set batch size to 1 since we'll be running inference on
# one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
GPU_COUNT = 1
IMAGES_PER_GPU = 1
config = InferenceConfig()
config.display()
model = modellib.MaskRCNN(mode="inference",
model_dir=MODEL_DIR,
config=config)
model.load_weights(COCO_MODEL_PATH, by_name=True)
# 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'
]
foldernames = [
f for f in os.listdir(IMAGE_DIR)
if f.isnumeric() and not f.startswith('.')
]
foldernames.sort()
for foldername in foldernames:
print(foldername)
if 1:
CURRENT_IMAGE_DIR = IMAGE_DIR + foldername + '/'
CURRENT_OUTPUT_DIR = OUTPUT_DIR + foldername + '/'
if not os.path.isdir(CURRENT_OUTPUT_DIR):
os.mkdir(CURRENT_OUTPUT_DIR)
filenames = sorted(os.listdir(CURRENT_IMAGE_DIR))
for filename in filenames:
print(filename)
if not filename.startswith('.'):
img = skimage.io.imread(CURRENT_IMAGE_DIR + filename)
results = model.detect([img], verbose=1)
# Visualize results
r = results[0]
r['masks'], r['rois'], r['class_ids'], r[
'scores'] = self.filter_masks(
r['masks'], r['rois'], r['class_ids'],
r['scores'])
visualize.display_instances(
img,
r['rois'],
r['masks'],
r['class_ids'],
class_names,
r['scores'],
save_option=1,
save_path=CURRENT_OUTPUT_DIR + filename)
plt.close('all')
def __init__(self, fig, ax, args):
self.ax = ax
self.ax.set_yticklabels([])
self.ax.set_xticklabels([])
self.img_dir = args['image_dir']
self.index = 0
self.fig = fig
self.polys = []
self.zoom_scale, self.points, self.prev, self.submit_p, self.lines, self.circles = 1.2, [], None, None, [], []
self.zoom_id = fig.canvas.mpl_connect('scroll_event', self.zoom)
self.click_id = fig.canvas.mpl_connect('button_press_event',
self.onclick)
self.clickrel_id = fig.canvas.mpl_connect('button_release_event',
self.onclick_release)
self.keyboard_id = fig.canvas.mpl_connect('key_press_event',
self.onkeyboard)
self.axradio = plt.axes([0.0, 0.0, 0.2, 1])
self.axbringprev = plt.axes([0.3, 0.05, 0.17, 0.05])
self.axreset = plt.axes([0.48, 0.05, 0.1, 0.05])
self.axsubmit = plt.axes([0.59, 0.05, 0.1, 0.05])
self.axprev = plt.axes([0.7, 0.05, 0.1, 0.05])
self.axnext = plt.axes([0.81, 0.05, 0.1, 0.05])
self.b_bringprev = Button(self.axbringprev,
'Bring Previous Annotations')
self.b_bringprev.on_clicked(self.bring_prev)
self.b_reset = Button(self.axreset, 'Reset')
self.b_reset.on_clicked(self.reset)
self.b_submit = Button(self.axsubmit, 'Submit')
self.b_submit.on_clicked(self.submit)
self.b_next = Button(self.axnext, 'Next')
self.b_next.on_clicked(self.next)
self.b_prev = Button(self.axprev, 'Prev')
self.b_prev.on_clicked(self.previous)
self.button_axes = [
self.axbringprev, self.axreset, self.axsubmit, self.axprev,
self.axnext, self.axradio
]
self.existing_polys = []
self.existing_patches = []
self.selected_poly = False
self.objects = []
self.feedback = args['feedback']
self.right_click = False
self.text = ''
with open(args['class_file'], 'r') as f:
self.class_names = [
x.strip() for x in f.readlines() if x.strip() != ""
]
self.radio = RadioButtons(self.axradio, self.class_names)
self.class_names = ('BG', ) + tuple(self.class_names)
self.img_paths = sorted(glob.glob(os.path.join(self.img_dir, '*.jpg')))
if len(self.img_paths) == 0:
self.img_paths = sorted(
glob.glob(os.path.join(self.img_dir, '*.png')))
if os.path.exists(self.img_paths[self.index][:-3] + 'txt'):
self.index = len(glob.glob(os.path.join(self.img_dir, '*.txt')))
self.checkpoint = self.index
try:
im = Image.open(self.img_paths[self.index])
except IndexError:
print(
"Reached end of dataset! Delete some TXT files if you want to relabel some images in the folder"
)
exit()
width, height = im.size
im.close()
image = plt.imread(self.img_paths[self.index])
if args['feedback']:
from mask_rcnn import model as modellib
from mask_rcnn.get_json_config import get_demo_config
#from skimage.measure import find_contours
from .contours import find_contours
from mask_rcnn.visualize_cv2 import random_colors
config = get_demo_config(len(self.class_names) - 2, True)
if args['config_path'] is not None:
config.from_json(args['config_path'])
# Create model object in inference mode.
model = modellib.MaskRCNN(
mode="inference",
model_dir='/'.join(args['weights_path'].split('/')[:-2]),
config=config)
# Load weights trained on MS-COCO
model.load_weights(args['weights_path'], by_name=True)
r = model.detect([image], verbose=0)[0]
# Number of instances
N = r['rois'].shape[0]
masks = r['masks']
# Generate random colors
colors = random_colors(N)
# Show area outside image boundaries.
height, width = image.shape[:2]
class_ids, scores = r['class_ids'], r['scores']
for i in range(N):
color = colors[i]
# Label
class_id = class_ids[i]
score = scores[i] if scores is not None else None
label = self.class_names[class_id]
# Mask
mask = masks[:, :, i]
# Mask Polygon
# Pad to ensure proper polygons for masks that touch image edges.
padded_mask = np.zeros((mask.shape[0] + 2, mask.shape[1] + 2),
dtype=np.uint8)
padded_mask[1:-1, 1:-1] = mask
contours = find_contours(padded_mask, 0.5)
for verts in contours:
# Subtract the padding and flip (y, x) to (x, y)
verts = np.fliplr(verts) - 1
pat = PatchCollection([Polygon(verts, closed=True)],
facecolor='green',
linewidths=0,
alpha=0.6)
self.ax.add_collection(pat)
self.objects.append(label)
self.existing_patches.append(pat)
self.existing_polys.append(
Polygon(verts,
closed=True,
alpha=0.25,
facecolor='red'))
self.ax.imshow(image, aspect='auto')
self.text += str(self.index) + '\n' + os.path.abspath(self.img_paths[
self.index]) + '\n' + str(width) + ' ' + str(height) + '\n\n'
def __init__(self, fig, ax, img_dir, classes, model_path, json_file):
self.RS = RectangleSelector(ax, self.line_select_callback,
drawtype='box', useblit=True,
button=[1, 3], # don't use middle button
minspanx=5, minspany=5,
spancoords='pixels',
interactive=True)
ax.set_yticklabels([])
ax.set_xticklabels([])
#self.classes, self.img_paths, _ = read_JSON_file(json_file)
with open(classes, 'r') as f:
self.classes, img_paths = sorted([x.strip().split(',')[0] for x in f.readlines()]), glob.glob(os.path.abspath(os.path.join(img_dir, '*.jpg')))
plt.tight_layout()
self.ax = ax
self.fig = fig
self.axradio = plt.axes([0.0, 0.0, 0.1, 1])
self.radio = RadioButtons(self.axradio, self.classes)
self.zoom_scale = 1.2
self.zoom_id = self.fig.canvas.mpl_connect('scroll_event', self.zoom)
self.keyboard_id = self.fig.canvas.mpl_connect('key_press_event', self.onkeyboard)
self.selected_poly = False
self.axsave = plt.axes([0.81, 0.05, 0.1, 0.05])
self.b_save = Button(self.axsave, 'Save')
self.b_save.on_clicked(self.save)
self.objects, self.existing_patches, self.existing_rects = [], [], []
self.num_pred = 0
if json_file is None:
self.images, self.annotations = [], []
self.index = 0
self.ann_id = 0
else:
with open(json_file, 'r') as g:
d = json.loads(g.read())
self.images, self.annotations = d['images'], d['annotations']
self.index = len(self.images)
self.ann_id = len(self.annotations)
prev_files = [x['file_name'] for x in self.images]
for i, f in enumerate(img_paths):
im = Image.open(f)
width, height = im.size
dic = {'file_name': f, 'id': self.index+i, 'height': height, 'width': width}
if f not in prev_files:
self.images.append(dic)
else:
self.index+=1
image = plt.imread(self.images[self.index]['file_name'])
self.ax.imshow(image, aspect='auto')
if not args['no_feedback']:
from mask_rcnn.get_json_config import get_demo_config
from mask_rcnn import model as modellib
from mask_rcnn.visualize_cv2 import random_colors
self.config = get_demo_config(len(self.classes)-1, True)
if 'config_path' in args:
self.config.from_json(args['config_path'])
plt.connect('draw_event', self.persist)
# Create model object in inference mode.
self.model = modellib.MaskRCNN(mode="inference", model_dir='/'.join(args['weights_path'].split('/')[:-2]), config=self.config)
# Load weights trained on MS-COCO
self.model.load_weights(args['weights_path'], by_name=True)
r = self.model.detect([image], verbose=0)[0]
# Number of instances
N = r['rois'].shape[0]
masks = r['masks']
# Show area outside image boundaries.
height, width = image.shape[:2]
class_ids, scores, rois = r['class_ids'], r['scores'], r['rois'],
for i in range(N):
# Label
class_id = class_ids[i]
score = scores[i] if scores is not None else None
label = self.classes[class_id-1]
pat = patches.Rectangle((rois[i][1], rois[i][0]), rois[i][3]-rois[i][1], rois[i][2]-rois[i][0], linewidth=1, edgecolor='r',facecolor='r', alpha=0.4)
rect = self.ax.add_patch(pat)
self.objects.append(label)
self.existing_patches.append(pat.get_bbox().get_points())
self.existing_rects.append(pat)
self.num_pred = len(self.objects)
iaa.Affine(rotate=(-30, 30)),
iaa.CropAndPad(percent=(-0.25, 0.25))])
dataset_train = FishDataset()
dataset_train.load_data(args.folders, 'whole_body', name=args.project)
dataset_train.prepare()
# Validation dataset
dataset_val = FishDataset('val')
dataset_val.load_data(args.folders, 'whole_body', name=args.project)
dataset_val.prepare()
config.MEAN_PIXEL = dataset_train.mean
config.display()
model = modellib.MaskRCNN(mode="training", config=config, model_dir=args.logs)
if 'mask_rcnn' in args.model.lower() and 'logs' not in args.model.lower():
model.load_weights(args.model, by_name=True, \
exclude=["mrcnn_class_logits", "mrcnn_bbox_fc", "mrcnn_bbox", "mrcnn_mask"])
else:
model.load_weights(args.model, by_name=True)
schedfactor = 5
lr = 1e-3
schedule = lambda epoch: lr*np.power(0.15, epoch // schedfactor)
temps = 40
print("Training network heads on augmented dataset")
model.train(dataset_train, dataset_val,
