目標檢測是一種計算機視覺技術，用於識別和定位圖像中的目標。有很多檢測算法存在，這裡有一個很好的總結。
Mask R-CNN是目標檢測的擴展，它為圖像中檢測到的每個目標生成邊界框和分割掩模。這篇文章是關於使用Mask R-CNN訓練自定義數據集的指南，希望它能幫助你們中的一些人簡化這個過程。
庫和包
算法的主要包是mrcnn。下載庫並將其導入到環境中。
<code>!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/<code>
mrcnn還不能與TensorFlow 2.0兼容，所以請確保您恢復到TensorFlow 1.x。因為我是在Colab上開發的，所以我將使用magic函數來恢復到TensorFlow 1.x。這也是TF被詬病的地方，兼容基本靠改。
<code>%tensorflow_version 1.x
import tensorflow as tf/<code>
在TensorFlow 2.0中，tf.random_shuffle被重命名為tf.random.shuffle，從而導致不兼容問題。通過更改mrcnn代碼中的shuffle函數，可以使用TensorFlow 2.0。
使用Colab最好把Keras到以前的版本，如果遇到錯誤的話，這樣才做，沒有的話就忽略吧。
<code>!pip install keras==2.2.5/<code>
預處理
mrcnn包在接收的數據格式方面相當靈活。我們這裡直接處理成NumPy數組。
在此之前，cv2無法正確讀取video17295和 video191900。因此，我過濾掉了這些圖像並創建了一個文件名列表。
<code>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"))/<code>
接下來要做的事情
· 檢查標籤是否存在（某些圖像不包含無人機）
· 讀取和處理圖像
· 讀取和處理邊界框的座標
· 可視化目的繪製邊界框
<code>X,y = [], []
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)/<code>
在轉換為NumPy數組之前，我獲取數據集的一個子集，作為測試可以減少訓練時間。
如果你有計算能力的話，可以省略。
MRCNN處理
現在來看看mrcnn本身，我們需要在訓練過程之前定義一個mrcnn數據集類。這個數據集類提供圖像的信息，比如它所屬的類和對象在其中的位置。mrcnn.utils包含這個類
這裡的事情變得有點棘手，需要閱讀一些源代碼。這些是你需要修改的功能：
· add_class，用於確定模型的類數
· 添加映像，在其中定義映像id和映像路徑（如果適用）
· 加載圖像，其中加載圖像數據
· 加載掩碼，獲取有關圖像的掩碼/邊框的信息
<code># 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)/<code>
我們已經將圖像格式化為NumPy數組，因此可以簡單地用數組初始化Dataset類，並通過索引到數組中來加載圖像和邊界框。
接下來分割訓練和測試集。
<code># 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:]/<code>
現在將數據加載到數據集類中。
<code># 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()/<code>
prepare（）函數使用圖像ID和類ID信息為mrcnn模型準備數據，下面是我們從mrcnn導入的config類的修改。Config類確定訓練中使用的變量，應該根據數據集進行調整。
下面的這些變量並非詳盡無遺，您可以參考文檔中的完整列表。
<code>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()/<code>
根據您的計算能力，您可能需要相應地調整這些變量。否則，您將面臨卡在"Epoch 1"的問題，並且不會給出錯誤消息。甚至還有針對這個問題提出的GitHub問題，並提出了許多解決方案。如果你遇到這種情況，一定要檢查一下，並測試一下這些建議中的一些。
https://github.com/matterport/Mask_RCNN/issues/287
MRCNN 訓練
mrcnn通過COCO和ImageNet數據集進行了訓練。所以這裡只要使用這些預先訓練的權重進行遷移學習，我們需要將其下載到環境中（記住首先定義根目錄）
<code># 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)/<code>
創建模型並使用預先訓練的權重。
<code># 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"])/<code>現在，我們可以開始進行實際訓練。
<code>model.train(train_dataset, val_dataset,learning_rate=config.LEARNING_RATE,epochs=5,layers='heads') # unfreeze head and just train on last layer/<code>
我只訓練最後一層來檢測數據集中的無人機。如果時間允許，您還應該通過訓練前面的所有層來微調模型。
<code>model.train(train_dataset, val_dataset, learning_rate=config.LEARNING_RATE / 10, epochs=2, layers="all")/<code>
完成了mrcnn模型的訓練後。可以用這兩行代碼保存模型的權重。
<code># save weights
model_path = os.path.join(MODEL_DIR, "mask_rcnn_drones.h5")
model.keras_model.save_weights(model_path)/<code>
MRCNN推斷
要對其他圖片進行推理，需要創建一個具有自定義配置的新推理模型。
<code># make inference
class 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)/<code>
可視化
<code>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())/<code>
好了，我們已經訓練了一個帶有自定義數據集的mrcnn模型。
分享到:
閱讀更多 deephub 的文章
