Mask R-CNN上手指南:对象检测和分割实现无人机的检测

 !pip install mrcnnfrom mrcnn.config import Config
 from mrcnn import utils
 import mrcnn.model as modellib
 from mrcnn import visualize
 from mrcnn.model import log

 %tensorflow_version 1.x
 import tensorflow as tf

 !pip install keras==2.2.5

 dir = "Database1/"# filter out image that cant be read
 prob_list = ['video17_295','video19_1900'] # cant read format
 txt_list = [f for f in os.listdir(dir) if f.endswith(".txt") and f[:-4] not in prob_list]
 file_list = set([re.match("\w+(?=.)",f)[0] for f in txt_list])# create data list as tuple of (jpeg,txt)
 data_list = []
 for f in file_list:
     data_list.append((f+".JPEG",f+".txt"))

 img_box = []
 DIMENSION = 128 # set low resolution to decrease training timefor i in range(len(data_list)):
     # get bounding box and check if label exist
     with open(dir+data_list[i][1],"rb") as f:
     box = f.read().split()
     if len(box) != 5:
         continue # skip data if does not contain labelbox = [float(s) for s in box[1:]]# read imageimg = cv2.imread(dir+data_list[i][0])
     img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)# resize img to 128 x 128
     img = cv2.resize(img, (DIMENSION,DIMENSION), interpolation= cv2.INTER_LINEAR)# draw bounding box (for visualization purposes)
     resize1, resize2 = img.shape[0]/DIMENSION, img.shape[1]/DIMENSION
     p1,p2,p3,p4 = int(box[0]*img.shape[1]*resize2), int(box[1]*img.shape[0]*resize1) ,int(box[2]*img.shape[1]*resize2) ,int(box[3]*img.shape[0]*resize1)ymin, ymax, xmin, xmax = p2-p4//2, p2+p4//2, p1-p3//2, p1+p3//2draw = cv2.rectangle(img.copy(),(xmax,ymax),(xmin,ymin),color=(255,255,0),thickness =1)# store data if range of y is at least 20 pixels (remove data with small drones)
     if ymax - ymin >=20:
         X.append(img)
         y.append([ymin, ymax, xmin, xmax])
         img_box.append(draw)# convert to numpy arraysX = np.array(X).astype(np.uint8)
 y = np.array(y)
 img_box = np.array(img_box)

# define drones dataset using mrcnn utils classclass DronesDataset(utils.Dataset):
     def __init__(self,X,y): # init with numpy X,y
         self.X = X
         self.y = y
         super().__init__()def load_dataset(self):
         self.add_class("dataset",1,"drones") # only 1 class, drones
         for i in range(len(self.X)):
             self.add_image("dataset",i,path=None)def load_image(self,image_id):
         image = self.X[image_id] # where image_id is index of X
         return imagedef load_mask(self,image_id):
     # get details of image
     info = self.image_info[image_id]
     #create one array for all masks, each on a different channel
     masks = np.zeros([128, 128, len(self.X)], dtype='uint8')class_ids = []
     for i in range(len(self.y)):
         box = self.y[info["id"]]
         row_s, row_e = box[0], box[1]
         col_s, col_e = box[2], box[3]
         masks[row_s:row_e, col_s:col_e, i] = 1 # create mask with similar boundaries as bounding box
         class_ids.append(1)return masks, np.array(class_ids).astype(np.uint8)

 # train test split 80:20np.random.seed(42) # for reproducibility
 p = np.random.permutation(len(X))
 X = X[p].copy()
 y = y[p].copy()split = int(0.8 * len(X))X_train = X[:split]
 y_train = y[:split]X_val = X[split:]
 y_val = y[split:]

 # load dataset into mrcnn dataset classtrain_dataset = DronesDataset(X_train,y_train)
 train_dataset.load_dataset()
 train_dataset.prepare()val_dataset = DronesDataset(X_val,y_val)
 val_dataset.load_dataset()
 val_dataset.prepare()

   class DronesConfig(Config):
     # Give the configuration a recognizable name
     NAME = "drones"# Train on 1 GPU and 2 images per GPU.
     GPU_COUNT = 1
     IMAGES_PER_GPU = 2# Number of classes (including background)
     NUM_CLASSES = 1+1  # background + drones# Use small images for faster training.
     IMAGE_MIN_DIM = 128
     IMAGE_MAX_DIM = 128# Reduce training ROIs per image because the images are small and have few objects.
     TRAIN_ROIS_PER_IMAGE = 20# Use smaller anchors because our image and objects are small
     RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)  # anchor side in pixels# set appropriate step per epoch and validation step
     STEPS_PER_EPOCH = len(X_train)//(GPU_COUNT*IMAGES_PER_GPU)
     VALIDATION_STEPS = len(X_val)//(GPU_COUNT*IMAGES_PER_GPU)# Skip detections with < 70% confidence
     DETECTION_MIN_CONFIDENCE = 0.7config = DronesConfig()
 config.display()

# 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)

# Create model in training mode using gpuwith tf.device("/gpu:0"):
     model = modellib.MaskRCNN(mode="training", config=config,model_dir=MODEL_DIR)# Which weights to start with?
 init_with = "imagenet"  # imagenet, cocoif 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"])

model.train(train_dataset, val_dataset,learning_rate=config.LEARNING_RATE,epochs=5,layers='heads') # unfreeze head and just train on last layer

 model.train(train_dataset, val_dataset,
             learning_rate=config.LEARNING_RATE / 10,
             epochs=2,
             layers="all")

 # save weights
 model_path = os.path.join(MODEL_DIR, "mask_rcnn_drones.h5")
 model.keras_model.save_weights(model_path)

 # make inferenceclass InferenceConfig(DronesConfig):
     GPU_COUNT = 1
     IMAGES_PER_GPU = 1inference_config = InferenceConfig()# Recreate the model in inference mode
 model = modellib.MaskRCNN(mode="inference",config=inference_config, model_dir=MODEL_DIR)# Load trained weightsmodel_path = os.path.join(MODEL_DIR, "mask_rcnn_drones.h5")
 model.load_weights(model_path, by_name=True)

 def get_ax(rows=1, cols=1, size=8):
     _, ax = plt.subplots(rows, cols, figsize=(size*cols, size*rows))return ax# Test on a random image
 image_id = random.choice(val_dataset.image_ids)
 original_image, image_meta, gt_class_id, gt_bbox, gt_mask =\
 modellib.load_image_gt(val_dataset, inference_config,image_id, use_mini_mask=False)results = model.detect([original_image], verbose=1)
 r = results[0]visualize.display_instances(original_image, r['rois'], r['masks'], r['class_ids'],val_dataset.class_names, r['scores'], ax=get_ax())