import os
import sys
import json
import datetime
import numpy as np
import skimage.draw
import tensorflow as tf
import matplotlib.pyplot as plt
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
ROOT_DIR = os.path.abspath("../")
print(ROOT_DIR)
sys.path.append(ROOT_DIR)
from mrcnn.config import Config
from mrcnn import model as modellib, utils
from pathlib import Path
COCO_WEIGHTS_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
print(COCO_WEIGHTS_PATH)
DEFAULT_LOGS_DIR = os.path.join(ROOT_DIR, "logs")
print(DEFAULT_LOGS_DIR)
class LandslideConfig(Config):
"""用于对 toy balloon 数据集进行训练的配置.
派生自基本的Config类, 并覆盖一些值.
"""
NAME = "landslide"
IMAGES_PER_GPU = 1
NUM_CLASSES = 1 + 1
STEPS_PER_EPOCH = 100
DETECTION_MIN_CONFIDENCE = 0.9
IMAGE_MIN_DIM = 128
IMAGE_MAX_DIM = 128
class LandslideDataset(utils.Dataset):
def load_landslide(self, dataset_dir, subset):
"""加载 Balloon 数据集的子集.
dataset_dir: 数据集的根目录.
要加载的子集: train or val
"""
self.add_class("landslide", 1, "landslide")
assert subset in ["train", "val"]
dataset_dir = os.path.join(dataset_dir, subset)
annotations = json.load(open(os.path.join(dataset_dir, "via_region_data.json")))
annotations = list(annotations.values())
annotations = [a for a in annotations if a['regions']]
for a in annotations:
if type(a['regions']) is dict:
polygons = [r['shape_attributes'] for r in a['regions'].values()]
else:
polygons = [r['shape_attributes'] for r in a['regions']]
image_path = os.path.join(dataset_dir, a['filename'])
image = skimage.io.imread(image_path)
height, width = image.shape[:2]
self.add_image(
"landslide",
image_id=a['filename'],
path=image_path,
width=width, height=height,
polygons=polygons)
def load_mask(self, image_id):
"""生成图像的实例 mask.
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_info = self.image_info[image_id]
if image_info["source"] != "landslide":
return super(self.__class__, self).load_mask(image_id)
info = self.image_info[image_id]
mask = np.zeros([info["height"], info["width"], len(info["polygons"])],
dtype=np.uint8)
for i, p in enumerate(info["polygons"]):
rr, cc = skimage.draw.polygon(p['all_points_y'], p['all_points_x'])
mask[rr, cc, i] = 1
return mask.astype(np.bool), np.ones([mask.shape[-1]], dtype=np.int32)
def image_reference(self, image_id):
"""返回图像的路径."""
info = self.image_info[image_id]
if info["source"] == "landslide":
return info["path"]
else:
super(self.__class__, self).image_reference(image_id)
def train(model):
"""训练模型."""
dataset_train = LandslideDataset()
dataset_train.load_landslide(args.dataset, "train")
dataset_train.prepare()
dataset_val = LandslideDataset()
dataset_val.load_landslide(args.dataset, "val")
dataset_val.prepare()
print("Training network heads")
model.train(dataset_train, dataset_val,
learning_rate=config.LEARNING_RATE,
epochs=30,
layers='heads')
def color_splash(image, mask):
"""应用 颜色飞溅 效果.
image: RGB image [height, width, 3]
mask: instance segmentation mask [height, width, instance count]
Returns 结果图像.
"""
gray = skimage.color.gray2rgb(skimage.color.rgb2gray(image)) * 255
if mask.shape[-1] > 0:
mask = (np.sum(mask, -1, keepdims=True) >= 1)
splash = np.where(mask, image, gray).astype(np.uint8)
else:
splash = gray.astype(np.uint8)
return splash
def detect_and_color_splash(model, image_path=None, video_path=None):
assert image_path or video_path
if image_path:
print("Running on {}".format(args.image))
image = skimage.io.imread(args.image)
r = model.detect([image], verbose=1)[0]
splash = color_splash(image, r['masks'])
file_name = "splash_{:%Y%m%dT%H%M%S}.png".format(datetime.datetime.now())
skimage.io.imsave(file_name, splash)
elif video_path:
import cv2
vcapture = cv2.VideoCapture(video_path)
width = int(vcapture.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vcapture.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = vcapture.get(cv2.CAP_PROP_FPS)
file_name = "splash_{:%Y%m%dT%H%M%S}.avi".format(datetime.datetime.now())
vwriter = cv2.VideoWriter(file_name,
cv2.VideoWriter_fourcc(*'MJPG'),
fps, (width, height))
count = 0
success = True
while success:
print("frame: ", count)
success, image = vcapture.read()
if success:
image = image[..., ::-1]
r = model.detect([image], verbose=0)[0]
splash = color_splash(image, r['masks'])
splash = splash[..., ::-1]
vwriter.write(splash)
count += 1
vwriter.release()
print("Saved to ", file_name)
def image_name(images_path,file_tpye):
images = []
for root,dirs,files in os.walk(images_path):
for file in files:
if os.path.splitext(file)[1] == file_tpye:
images.append(os.path.join(root,file))
return images
def get_ax(rows=1, cols=1, size=16):
_, ax = plt.subplots(rows, cols, figsize=(size * cols, size * rows))
return ax
def detect_and_show(model, image_path=None):
import visualize_cv2
dataset = LandslideDataset()
dataset.load_landslide(args.dataset, "val")
dataset.prepare()
images = image_name(Path(args.image), ".png")
print(images)
count = 1
for i in images:
print("Running on {}".format(i))
image = skimage.io.imread(i)
results = model.detect([image], verbose=1)
r = results[0]
file_name = "detected_{:%Y%m%dT%H%M%S}{count}.png".format(datetime.datetime.now(), count=count)
visualize_cv2.save_image(image, file_name, r['rois'], r['masks'],
r['class_ids'], r['scores'], dataset.class_names,
filter_classs_names=['landslide'], scores_thresh=0.7, mode=0)
print("Saved to ", file_name)
count = count + 1
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(
description='Train Mask R-CNN to detect landslide.')
parser.add_argument("command",
metavar="",
help="'train' or 'splash'")
parser.add_argument('--dataset', required=False,
metavar="/path/to/landslide/dataset/",
help='Directory of the landslide dataset')
parser.add_argument('--weights', required=True,
metavar="/path/to/mask_rcnn_coco.h5",
help="Path to weights .h5 file or 'coco'")
parser.add_argument('--logs', required=False,
default=DEFAULT_LOGS_DIR,
metavar="/path/to/logs/",
help='Logs and checkpoints directory (default=logs/)')
parser.add_argument('--image', required=False,
metavar="path or URL to image",
help='Image to apply the color splash effect on')
parser.add_argument('--video', required=False,
metavar="path or URL to video",
help='Video to apply the color splash effect on')
args = parser.parse_args()
if args.command == "train":
assert args.dataset, "Argument --dataset is required for training"
elif args.command == "splash":
assert args.image or args.video,\
"Provide --image or --video to apply color splash"
print("Weights: ", args.weights)
print("Dataset: ", args.dataset)
print("Logs: ", args.logs)
if args.command == "train":
config = LandslideConfig()
else:
class InferenceConfig(LandslideConfig):
GPU_COUNT = 1
IMAGES_PER_GPU = 1
config = InferenceConfig()
config.display()
if args.command == "train":
model = modellib.MaskRCNN(mode="training", config=config,
model_dir=args.logs)
else:
model = modellib.MaskRCNN(mode="inference", config=config,
model_dir=args.logs)
if args.weights.lower() == "coco":
weights_path = COCO_WEIGHTS_PATH
if not os.path.exists(weights_path):
utils.download_trained_weights(weights_path)
elif args.weights.lower() == "last":
weights_path = model.find_last()
elif args.weights.lower() == "imagenet":
weights_path = model.get_imagenet_weights()
else:
weights_path = args.weights
print("Loading weights ", weights_path)
if args.weights.lower() == "coco":
model.load_weights(weights_path, by_name=True, exclude=[
"mrcnn_class_logits", "mrcnn_bbox_fc",
"mrcnn_bbox", "mrcnn_mask"])
else:
model.load_weights(weights_path, by_name=True)
if args.command == "train":
train(model)
elif args.command == "splash":
detect_and_color_splash(model, image_path=args.image,
video_path=args.video)
elif args.command == "show":
detect_and_show(model, image_path=args.image)
else:
print("'{}' is not recognized. "
"Use 'train' or 'splash'".format(args.command))