無人機檢測: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 beread
prob_list = ['video17_295'
,'video19_1900'
] # cantread
format
txt_list = [ffor
fin
os
.listdir(dir)if
f.endswith(".txt"
)and
f[:-4
]not
in
prob_list] file_list = set([re.match
("\w+(?=.)"
,f)[0
]for
fin
txt_list]) #create
data list as tuple of (jpeg,txt) data_list = []for
fin
file_list: data_list.append((f+".JPEG"
,f+".txt"
))/<code>
接下來要做的事很少;
- 檢查標籤是否存在(某些圖像不包含無人機)
- 讀取並處理圖像
- 讀取和處理邊界框的座標
- 繪製邊界框以進行可視化
<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類。
這是有些棘手的問題,需要對源代碼進行一些閱讀。
這些是需要修改的功能;
- add_class,它確定模型的類數
- add_image,在其中定義image_id和圖像的路徑(如果適用)
- load_image,圖像數據加載到的位置
- load_mask,它獲取有關圖像的蒙版/邊框的信息
<code>class
DronesDataset
(utils
.Dataset
):def
__init__
(
self
,X,y):self
.X = Xself
.y = ysuper
().__init__
()def
load_dataset
(
self
):self
.add_class("dataset"
,1
,"drones"
)for
iin
range(len(self
.X)):self
.add_image("dataset"
,i,path=None)def
load_image
(
self
,image_id) : image =self
.X[image_id]return
imagedef
load_mask
(
self
,image_id): info =self
.image_info[image_id] masks = np.zeros([128
,128
, len(self
.X)], dtype='uint8'
) class_ids = []for
iin
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模型。如上圖所見。