無人機檢測：Mask R-CNN的分步解析（附代碼）

無人機檢測：Mask R-CNN的分步解析

對象檢測是一種計算機視覺技術，用於識別和定位圖像中的對象。

MaskR-CNN是對象檢測的擴展，因為它會為圖像中檢測到的每個對象生成邊界框和分割蒙版。

本文會介紹如何使用Mask R-CNN訓練自定義數據集，並希望可以幫助大家簡化該過程。

庫和包

該算法的主要軟件包是mrcnn。首先下載並將庫導入到環境中。

<code>！
pip
 

從
mrcnn
.config
導入
安裝
mrcnn
 

從
mrcnn
導入配置
從
utils
 

導入
mrcnn
.model
作為
從
mrcnn
導入的

modellib
 從
mrcnn
.model
導入可視化/<code>

到現在，我們只知道這些就是需要的import語句。

至於TensorFlow，mrcnn尚未與TensorFlow 2.0兼容，因此請確保恢復到TensorFlow 1.x

如果我沒看錯，則在TensorFlow 2.0中將tf.random_shuffle重命名為tf.random.shuffle，從而導致不兼容問題。通過更改mrcnn代碼中的shuffle函數，可以使用TensorFlow 2.0。

<code>
！pip
 
install
 
keras
 
==
 
2.2
.5
/<code>

預處理

mrcnn包在其接受的數據格式方面相當靈活。因此，由於其簡單的操作，將其處理成NumPy數組。

在此之前，我意識到cv2無法正確讀取video17_295和video19_1900。因此過濾掉了這些圖像並創建了文件名列表。

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

接下來要做的事很少；

1. 檢查標籤是否存在（某些圖像不包含無人機）
2. 讀取並處理圖像
3. 讀取和處理邊界框的座標
4. 繪製邊界框以進行可視化

<code>  = 
[], []
 
img_box
 = 
[]
DIMENSION
 = 
128 # set low resolution to decrease training time
for
 
i in range(len(data_list)):
 
    
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 label
box
 = 
[float(s) for s in box[1:]]
 
img
 = 
cv2.imread(dir+data_list[i][0])
    
img
 = 
cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
 
    
img
 = 
cv2.resize(img, (DIMENSION,DIMENSION), interpolation= cv2.INTER_LINEAR)
 
      
resize2 = img.shape[0]/DIMENSION, img.shape[1]/DIMENSION
      = 
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)
  
ymax, xmin, xmax = p2-p4//2, p2+p4//2, p1-p3//2, p1+p3//2
 
draw
 = 
cv2.rectangle(img.copy(),(xmax,ymax),(xmin,ymin),color=(255,255,0),thickness =1)
 
    
if
 
ymax - ymin >=20:
        
X.append(img)
          
ymax, xmin, xmax])
        
img_box.append(draw)
 
X
 = 
np.array(X).astype(np.uint8)
y
 = 
np.array(y)
img_box
 = 
np.array(img_box)
/<code>

在轉換為NumPy數組之前，我獲取了數據集的一個子填充，以減少訓練時間。如果具有計算能力，可以忽略它。這是一些示例圖像。

MRCNN-處理

現在，要正確看待mrcnn，我們需要在訓練過程之前定義一個mrcnn Dataset類。該Dataset類提供圖像的信息，例如圖像所屬的類以及對象在其中的位置。我們之前導入的mrcnn.utils包含此Dataset類。

這是有些棘手的問題，需要對源代碼進行一些閱讀。

這些是需要修改的功能；

1. add_class，它確定模型的類數
2. add_image，在其中定義image_id和圖像的路徑（如果適用）
3. load_image，圖像數據加載到的位置
4. load_mask，它獲取有關圖像的蒙版/邊框的信息

<code> 
class
 
DronesDataset
(
utils
.
Dataset
):

    
def
 
__init__ 
(
self
,X,y)
:  
        
self
.X = X
        
self
.y = y
        
super
().__init_
_
()
def
 
load_dataset
(
self
)
:
        
self
.add_class(
"dataset"
,
1
,
"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]  
        
return
 image
def
 
load_mask
(
self
,image_id)
:
     
    info = 
self
.image_info[image_id]
     
    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
  
        class_ids.append( 
1
)
return
 masks, np.array(class_ids).astype(np.uint8)/<code>

由於我們努力將圖像格式化為NumPy數組，因此我們可以簡單地使用數組初始化Dataset類，並通過索引數組來加載圖像和邊框。

接下來要做一次火車測試，以拆分傳統方法，

<code> 
np.random.seed(42)  
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>

現在將數據加載到Dataset類中。

<code>＃將數據集加載到mrcnn數據集類中
train_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（）函數使用image_ids和class_ids信息為mrcnn模型準備數據，

接下來是對我們從mrcnn導入的config類的修改。Config類確定了訓練中使用的變量，應根據數據集進行調整。以下這些變量並不詳盡，可以參考文檔以獲取完整列表。

<code>
class
 
DronesConfig(Config):
 
    
NAME
 = 
"drones"
 
    
GPU_COUNT
 = 
1
     
IMAGES_PER_GPU
 = 
2
 
    
NUM_CLASSES
 = 
1+1  # background + drones
 
    
IMAGE_MIN_DIM
 = 
128
    
IMAGE_MAX_DIM
 = 
128
 
    
TRAIN_ROIS_PER_IMAGE
 = 
20
 
    
RPN_ANCHOR_SCALES
 = 
(8, 16, 32, 64, 128)  # anchor side in pixels
 
    
STEPS_PER_EPOCH
 = 
len(X_train)//(GPU_COUNT*IMAGES_PER_GPU)
    
VALIDATION_STEPS
 = 
len(X_val)//(GPU_COUNT*IMAGES_PER_GPU)
 
    
DETECTION_MIN_CONFIDENCE
 = 
0.7
config
 = 
DronesConfig()
config.display()
/<code>

根據電腦的計算能力，必須相應地調整這些變量。否則將面臨卡在第1階段而沒有給出任何錯誤消息的問題。甚至為此問題引發了GitHub問題，並提出了許多解決方案。

MRCNN訓練

mrcnn已經在COCO和I mageNet數據集中進行了訓練。為了將這些預訓練的權重用於轉移學習，我們需要將其下載到我們的環境中（請記住首先定義ROOT_DIR）。

<code>＃以訓練的權重的文件本地路徑
COCO_MODEL_PATH = 
os
.
path
.join(ROOT_DIR, 
"mask_rcnn_coco.h5"
)

if
 
not
 
os
.
path
.exists（COCO_MODEL_PATH）：
    utils.download_trained_weights（COCO_MODEL_PATH）/<code>

創建模型並以預先訓練的權重開始。

<code> 
with
 tf.device(
"/gpu:0"
):
    
model
 = modellib.MaskRCNN(
mode
=
"training"
, config=config,model_dir=MODEL_DIR)
 
init_with = 
"imagenet"
   
if
 init_with == 
"imagenet"
:
    model.load_weights(model.get_imagenet_weights(), by_name=
True 
)
elif init_with == 
"coco"
:
     
     
     
    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'
)  
/<code>

對於本次練習僅訓練最後一層以檢測數據集中的無人機。如果時間允許，還應該通過訓練所有前面的層來微調模型。

<code>
model.train（train_dataset，val_dataset，
            
learning_rate
 = 
config.LEARNING_RATE / 10，
            
epochs
 = 
2，
            
layers
 =
“ all”）
/<code>

訓練mrcnn模型就完成了。可以使用這兩行代碼保存模型的權重。

<code>＃保存權重
model_path = 
os
.
path
.join（MODEL_DIR，“ mask_rcnn_drones.h5”）
model.keras_model.save_weights（model_path）/<code>

MRCNN-推論

要對其他圖像進行推斷，將需要使用自定義Config創建一個新的推斷模型。

<code>
＃創建推理類InferenceConfig（DronesConfig）：
    
GPU_COUNT
 = 
1 
    
IMAGES_PER_GPU
 = 
1 inference_config = InferenceConfig（）＃在推理模式下重新創建模型
model
 = 
modellib.MaskRCNN（mode =“ inference”，
config
 = 
inference_config，model_dir = MODEL_DIR）＃訓練有素的負載權重
model_path
 = 
os.path.join（MODEL_DIR，“ mask_rcnn_drones.h5”）
  = 
True）
/<code>

來自mrcnn的可視化類在這裡派上用場。

<code> 
def
 
get_ax
(rows=
1
, cols=
1
, size=
8
)
:

    _, ax = plt.subplots(rows, cols, figsize=(size*cols, size*rows))
return
 ax
 
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模型。如上圖所見。