def main():
args = _get_args()
image_path = args.image_path
class_names = {1: 'mask', 2: 'no_mask', 3: 'mask_worn_incorrectly'}
inf_config = InferenceConfig()
model = MaskRCNN(mode='inference', config=inf_config, model_dir='./')
model_path = model.find_last()
model.load_weights(model_path, by_name=True)
for filename in os.listdir(image_path):
if filename.split('.')[-1] not in ['jpg', 'png', 'jpeg']:
continue
img = skimage.io.imread(image_path + filename)
img_arr = np.array(img)
img_arr = img_arr[:, :, :3]
results = model.detect([img_arr], verbose=1)
r = results[0]
display_instances(img,
r['rois'],
r['masks'],
r['class_ids'],
class_names,
r['scores'],
figsize=(10, 10))
def load_mask(self, image_idx, visualization = False):
"""Generate instance masks for an image.
Returns:
masks: A bool array of shape [height, width, instance count] with
one mask per instance.
class_ids: a 1D array of class IDs of the instance masks.
"""
image_id = self.image_info[image_idx]['id']
masks = None
class_ids = []
class_names = []
boxes = []
# Construct masks from bounding boxes
for img in self.images:
# If found the right Image in the JSON file...
if image_id == self.get_image_id_from_pathname_entry_in_json(img['image']['pathname']):
# For each bounding box in 'objects'....
masks = np.zeros((img['image']['shape']['r'], img['image']['shape']['c'], len(img['objects'])))
for o_idx, o in enumerate(img['objects']):
label = o['category']
min_y, min_x = o['bounding_box']['minimum']['r'], o['bounding_box']['minimum']['c']
max_y, max_x = o['bounding_box']['maximum']['r'], o['bounding_box']['maximum']['c']
masks[min_y:max_y, min_x:max_x, o_idx] = 1
class_names.append(label)
boxes.append((min_y, min_x, max_y, max_x))
class_ids.append(self.CLASSES[label])
break
if visualization:
fig, ax = plt.subplots(1, figsize = (10,10))
visualize.display_instances(self.load_image(image_idx), np.array(boxes), masks, np.array(class_ids), class_names, ax = ax)
fig.savefig('Mask for Image: ' + image_id + '.png', bbox_inches='tight')
return masks, np.array(class_ids)
def inference(model, dataset_dir):
"""Run inference on images in the given directory."""
print("Running on {}".format(dataset_dir))
# Create directory
if not os.path.exists(RESULTS_DIR):
os.makedirs(RESULTS_DIR)
submit_dir = "submit_{:%Y%m%dT%H%M%S}".format(datetime.datetime.now())
submit_dir = os.path.join(RESULTS_DIR, submit_dir)
os.makedirs(submit_dir)
# Read dataset
dataset = MalariaDataset()
dataset.load_dataset(dataset_dir, is_test=True)
dataset.prepare()
# Load over images
submission = []
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 10))
for image_id in dataset.image_ids:
# Load image and run inference
image = dataset.load_image(image_id)
raw_image_id = dataset.image_info[image_id]["id"]
# Detect objects
r = model.detect([image], verbose=0)[0]
# Save raw image
boxes, masks, class_ids, class_names = dataset.load_masks_and_boxes(
image_id)
visualize.display_instances(image,
boxes,
masks,
class_ids,
class_names,
show_bbox=True,
show_mask=False,
title="Ground Truth",
ax=ax1)
# Save prediction
visualize.display_instances(image,
r['rois'],
r['masks'],
r['class_ids'],
dataset.class_names,
r['scores'],
show_bbox=True,
show_mask=False,
title="Predictions",
ax=ax2)
fig.suptitle("Image: " + str(raw_image_id))
fig.savefig("{}/{}.png".format(submit_dir, raw_image_id),
bbox_inches='tight')
ax1.clear()
ax2.clear()
plt.cla()
print("---- DONE DETECTING ----")
def inference(inference_model, dataset, limit, tag=''):
"""Run detection on images in the given directory."""
# Create directory
if not os.path.exists(RESULTS_DIR):
os.makedirs(RESULTS_DIR)
time_dir = "{:%Y%m%dT%H%M%S}_{}".format(datetime.datetime.now(), tag)
time_dir = os.path.join(RESULTS_DIR, time_dir)
os.makedirs(time_dir)
if limit < 0:
limit = dataset.num_images
print("inference all: {} images".format(limit))
output_file = open('./test_prediction_{}.txt'.format(tag), 'w')
# Load over images
progress_bar = tqdm(dataset.image_ids[:limit])
for image_id_ in progress_bar:
# Load image and run detection
image_ = dataset.load_image(image_id_)
# Detect objects
r_ = inference_model.detect([image_], verbose=0)[0]
# Encode image to RLE. Returns a string of multiple lines
source_id = dataset.image_info[image_id_]["id"]
# Save image with masks
if len(r_['class_ids']) > 0:
progress_bar.set_description(
'[*] {}th image has {} instance(s).'.format(
image_id_, len(r_['class_ids'])))
visualize.display_instances(image_,
r_['rois'],
r_['masks'],
r_['class_ids'],
dataset.class_names,
r_['scores'],
show_bbox=True,
show_mask=False,
title="Predictions")
plt.savefig("{}/{}".format(time_dir, source_id))
plt.close()
image_name = os.path.basename(
dataset.image_info[image_id_]["path"])
output_result_format(image_name, r_['rois'], r_['class_ids'],
r_['scores'], output_file)
else:
plt.imshow(image_)
plt.savefig("{}/noinstance_{}".format(time_dir, source_id))
progress_bar.set_description(
'[*] {}th image have no instance.'.format(image_id_))
plt.close()
output_file.close()
def detect(model, dataset_dir, subset):
"""Run detection on images in the given directory."""
print("Running on {}".format(dataset_dir))
# Create directory
if not os.path.exists(RESULTS_DIR):
os.makedirs(RESULTS_DIR)
submit_dir = "submit_{:%Y%m%dT%H%M%S}".format(datetime.datetime.now())
submit_dir = os.path.join(RESULTS_DIR, submit_dir)
os.makedirs(submit_dir)
# Read dataset
dataset = NucleusDataset()
dataset.load_nucleus(dataset_dir, subset)
dataset.prepare()
# Load over images
submission = []
for image_id in dataset.image_ids:
# Load image and run detection
image = dataset.load_image(image_id)
# Detect objects
r = model.detect([image], verbose=0)[0]
# Encode image to RLE. Returns a string of multiple lines
source_id = dataset.image_info[image_id]["id"]
rle = mask_to_rle(source_id, r["masks"], r["scores"])
submission.append(rle)
# Save image with masks
visualize.display_instances(image,
r['rois'],
r['masks'],
r['class_ids'],
dataset.class_names,
r['scores'],
show_bbox=False,
show_mask=False,
title="Predictions")
plt.savefig("{}/{}.png".format(submit_dir,
dataset.image_info[image_id]["id"]))
# Save to csv file
submission = "ImageId,EncodedPixels\n" + "\n".join(submission)
file_path = os.path.join(submit_dir, "submit.csv")
with open(file_path, "w") as f:
f.write(submission)
print("Saved to ", submit_dir)
def detect(self, images, save=False):
"""
Detecting images
:param images: images loaded to memory
:param save: if should save or collect to self.frames
"""
results = self.model.detect(images, verbose=0)
# Visualize results
for img in images:
r = results.pop(0)
path = visualize.display_instances(img,
r['rois'],
r['masks'],
r['class_ids'],
class_names,
r['scores'],
colors=color_labels)
if not save:
self.frames.append(cv2.imread(path))
os.remove(path)
def func(filepath,filename):
# Root directory of the project
ROOT_DIR = os.path.abspath("../")
# Import Mask RCNN
sys.path.append(ROOT_DIR) # To find local version of the library
from Mask_RCNN.mrcnn import utils
import Mask_RCNN.mrcnn.model as modellib
from Mask_RCNN.mrcnn import visualize
# Import COCO config
sys.path.append(os.path.join(ROOT_DIR, "samples/coco/")) # To find local version
from Mask_RCNN.samples.coco.coco import CocoConfig
# 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)
# Directory of images to run detection on
IMAGE_DIR = os.path.join(ROOT_DIR, "images")
class InferenceConfig(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()
# Create model object in inference mode.
model = modellib.Mask_RCNN(mode="inference", model_dir=MODEL_DIR, config=config)
# Load weights trained on MS-COCO
model.load_weights(COCO_MODEL_PATH, by_name=True)
# COCO Class names
# Index of the class in the list is its ID. For example, to get ID of
# the teddy bear class, use: class_names.index('teddy bear')
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
# file_names = next(os.walk(IMAGE_DIR))[2]
image = skimage.io.imread(filepath)#os.path.join(IMAGE_DIR, random.choice(file_names)))
# Run detection
results = model.detect([image], verbose=1)
# Visualize results
r = results[0]
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
class_names, r['scores'],filename=filename)
# func()
def run_mask_rcnn(save_name,
target_dir,
image_dir,
model=None,
one=False,
display=False,
save_images=False,
detection_confidence=0.5,
detection_nms_threshold=0.3):
'''
:param save_name: The name of the h5py data file to create.
:param target_dir: The directory to save all the generated data
:param image_dir: The directory where the images are
:param model: use a mask-rcnn model with external configs
:param one: only run for one test image
:param display: Show the generated images with masks
:param save_images: save the generated images with masks
:return: N/A
'''
#TODO: Allow for batch processing of images
file_names = next(os.walk(image_dir))[2]
if (one):
file_name = random.choice(file_names)
while (file_name == ".DS_store"):
file_name = random.choice(file_names)
file_names = [file_name]
#Build the model
if (model == None):
model = create_model(detection_confidence=detection_confidence,
detection_nms_threshold=0.3)
if (not one):
save_path = os.path.join(target_dir, save_name)
f = h5py.File(save_path, "w")
f.create_dataset("class_names",
data=[np.string_(j) for j in class_names])
file_names.sort()
print(file_names)
i = 0
for file_name in file_names:
if (file_name == ".DS_Store"):
continue
image = skimage.io.imread(os.path.join(image_dir, file_name))
# Run detection
results = model.detect([image], verbose=0)
r = results[0]
if (not one):
f.create_dataset("frame{}/image".format(i), data=image)
for k, v in r.items():
f.create_dataset("frame{}/{}".format(i, k), data=v)
if (save_images):
save_dir_images = os.path.join(target_dir, "images/")
if not os.path.exists(save_dir_images):
os.makedirs(save_dir_images)
save_path_images = os.path.join(target_dir,
"images/frame{}.jpg".format(i))
visualize.save_instances(image, save_path_images, r['rois'],
r['masks'], r['class_ids'], class_names,
r['scores'])
if (display):
visualize.display_instances(image, r['rois'], r['masks'],
r['class_ids'], class_names,
r['scores'])
i += 1
number_of_frames = len(f.keys()) - 1
f.create_dataset("frame_number", data=[number_of_frames])
if (not one):
f.close()
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
# Load weights trained on MS-COCO
model.load_weights(COCO_MODEL_PATH, by_name=True)
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
file_names = next(os.walk(IMAGE_DIR))[2]
image = skimage.io.imread(os.path.join(IMAGE_DIR, random.choice(file_names)))
# Run detection
results = model.detect([image], verbose=1)
# Visualize results
r = results[0]
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
class_names, r['scores'])