一、前言
2014年，Ross Girshick提出RCNN，成為目標檢測領域的開山之作。一年後，借鑑空間金字塔池化思想，Ross Girshick推出設計更為巧妙的Fast RCNN（https://github.com/rbgirshick/fast-rcnn），極大地提高了檢測速度。Fast RCNN的提出解決了RCNN結構固有的三個弊端：
- 繁瑣的多階段訓練：RCNN在訓練時，首先需要在推薦區域上微調卷積網路，然後利用提取的卷積特徵針對每個類別訓練一個SVM分類器，最後還需要基於卷積特徵進行邊界框迴歸訓練；
- 高空間和時間成本：RCNN在訓練SVM和迴歸器時，需在磁碟上對推薦區域的卷積特徵進行讀寫，記憶體和時間消耗較為嚴重；
- 檢測速度慢：檢測時，需要對每個推薦區域進行特徵提取，計算重複性高，導致檢測速度很慢。
同前面RCNN實現一樣（見 https://www.cnblogs.com/Haitangr/p/17690028.html），本文將基於Pytorch框架，實現Fast RCNN演算法，完成對17flowes資料集的花朵目標檢測任務。
二、Fast RCNN演算法實現
如下為RCNN演算法和Fast RCNN演算法流程對比圖：
RCNN演算法實現過程中，需要將生成的所有推薦區域（~2k）縮放到同一大小後，全部走一遍卷積網路CNN，以提取相應特徵進行邊界框預測，這個過程極為耗時。同時，由於這2k張圖片均來源於同一張輸入，卷積網路會進行大量重複性計算。Fast RCNN則完全不同，其輸入圖片只進行一次CNN計算，以獲得整幅影像的特徵。而推薦區域的特徵則直接利用區域池化技術，根據相應的邊界框在全圖特徵上進行提取，大大降低了計算成本。此外，Fast RCNN採用多工損失對分類模型和迴歸模型同時進行最佳化，避免了繁瑣的多階段訓練過程。
下面是本文中Fast RCNN的實現流程：
- 候選區域生成：利用ss方法為每一幀圖片生成數量不固定的候選區域，利用IoU結果對候選區域進行標註，同時記錄標籤和邊界框資訊；
- 訓練集資料準備：構建可迭代資料集類，為模型訓練提供原始影像、標籤、推薦區域邊界框和邊界框偏移值資訊；
- 模型訓練：利用ROI池化提取候選區域特徵，利用多工損失同時訓練分類和迴歸模型；
- 模型預測：透過ss方法生成推薦區域，利用模型結果對推薦區域位置進行修正，再利用非極大值抑制剔除冗餘邊界框，獲得最終目標框。
1. 候選區域生成
Fast RCNNt同RCNN一樣採用選擇性搜尋（selective search，後面簡稱為ss）的辦法產生候選區域，ss方法的詳細思路同樣請參考 https://www.cnblogs.com/Haitangr/p/17690028.html 。不同的是，Fast RCNN只需要記錄候選區域的標籤（物體/背景）、候選區域的邊界框位置和對應的真實邊界框位置，而不需要儲存推薦區域影像。具體程式碼實現如下：
# SelectiveSearch.py import os import numpy as np import pandas as pd import cv2 as cv import shutil from Utils import cal_IoU from skimage import io from multiprocessing import Process import threading import matplotlib.pyplot as plt import matplotlib.patches as patches from Config import * class SelectiveSearch: def __init__(self, root, max_pos_regions: int = None, max_neg_regions: int = None, threshold=0.5): """ 採用ss方法生成候選區域檔案 :param root: 訓練/驗證資料集所在路徑 :param max_pos_regions: 每張圖片最多產生的正樣本候選區域個數, None表示不進行限制 :param max_neg_regions: 每張圖片最多產生的負樣本候選區域個數, None表示不進行限制 :param threshold: IoU進行正負樣本區分時的閾值 """ self.source_root = os.path.join(root, 'source') self.ss_root = os.path.join(root, 'ss') self.csv_path = os.path.join(self.source_root, "gt_loc.csv") self.max_pos_regions = max_pos_regions self.max_neg_regions = max_neg_regions self.threshold = threshold self.info = None @staticmethod def cal_proposals(img) -> np.ndarray: """ 計算後續區域座標 :param img: 原始輸入影像 :return: candidates, 候選區域座標矩陣n*4維, 每列分別對應[x, y, w, h] """ # 生成候選區域 ss = cv.ximgproc.segmentation.createSelectiveSearchSegmentation() ss.setBaseImage(img) ss.switchToSelectiveSearchFast() proposals = ss.process() candidates = set() # 對區域進行限制 for region in proposals: rect = tuple(region) if rect in candidates: continue candidates.add(rect) candidates = np.array(list(candidates)) return candidates def save(self, num_workers=1, method="thread"): """ 生成目標區域並儲存 :param num_workers: 程式或執行緒數 :param method: 多程式-process或者多執行緒-thread :return: None """ self.info = pd.read_csv(self.csv_path, header=0, index_col=None) index = self.info.index.to_list() span = len(index) // num_workers # 多程式生成影像 if "process" in method.lower(): print("=" * 8 + "開始多程式生成候選區域影像" + "=" * 8) processes = [] for i in range(num_workers): if i != num_workers - 1: p = Process(target=self.save_proposals, kwargs={'index': index[i * span: (i + 1) * span]}) else: p = Process(target=self.save_proposals, kwargs={'index': index[i * span:]}) p.start() processes.append(p) for p in processes: p.join() # 多執行緒生成影像 elif "thread" in method.lower(): print("=" * 8 + "開始多執行緒生成候選區域影像" + "=" * 8) threads = [] for i in range(num_workers): if i != num_workers - 1: thread = threading.Thread(target=self.save_proposals, kwargs={'index': index[i * span: (i + 1) * span]}) else: thread = threading.Thread(target=self.save_proposals, kwargs={'index': index[i * span: (i + 1) * span]}) thread.start() threads.append(thread) for thread in threads: thread.join() else: print("=" * 8 + "開始生成候選區域影像" + "=" * 8) self.save_proposals(index=index) return None def save_proposals(self, index, show_fig=False): """ 生成候選區域圖片並儲存相關資訊 :param index: 檔案index :param show_fig: 是否展示後續區域劃分結果 :return: None """ for row in index: name = self.info.iloc[row, 0] label = self.info.iloc[row, 1] # gt值為[x, y, w, h] gt_box = self.info.iloc[row, 2:].values im_path = os.path.join(self.source_root, name) img = io.imread(im_path) # 計算推薦區域座標矩陣[x, y, w, h] proposals = self.cal_proposals(img=img) # 計算proposals與gt的IoU結果 IoU = cal_IoU(proposals, gt_box) # 根據IoU閾值將proposals影像劃分到正負樣本集 p_boxes = proposals[np.where(IoU >= self.threshold)] n_boxes = proposals[np.where(IoU < self.threshold)] # 展示proposals結果 if show_fig: fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(6, 6)) ax.imshow(img) for (x, y, w, h) in p_boxes: rect = patches.Rectangle((x, y), w, h, fill=False, edgecolor='red', linewidth=1) ax.add_patch(rect) for (x, y, w, h) in n_boxes: rect = patches.Rectangle((x, y), w, h, fill=False, edgecolor='green', linewidth=1) ax.add_patch(rect) plt.show() # 根據影像名稱建立資料夾, 儲存原始圖片/真實邊界框/推薦區域邊界框/推薦區域標籤資訊 folder = name.split("/")[-1].split(".")[0] save_root = os.path.join(self.ss_root, folder) os.makedirs(save_root, exist_ok=True) # 儲存原始影像 im_save_path = os.path.join(save_root, folder + ".jpg") io.imsave(fname=im_save_path, arr=img, check_contrast=False) # loc.csv用於儲存邊界框資訊 loc_path = os.path.join(save_root, "ss_loc.csv") # 記錄正負樣本資訊 locations = [] header = ["label", "px", "py", "pw", "ph", "gx", "gy", "gw", "gh"] num_p = num_n = 0 for p_box in p_boxes: num_p += 1 locations.append([label, *p_box, *gt_box]) if self.max_pos_regions is None: continue if num_p >= self.max_pos_regions: break # 記錄負樣本資訊, 負樣本為背景, label置為0 for n_box in n_boxes: num_n += 1 locations.append([0, *n_box, *gt_box]) if self.max_neg_regions is None: continue if num_n >= self.max_neg_regions: break print("{name}: {num_p}個正樣本, {num_n}個負樣本".format(name=name, num_p=num_p, num_n=num_n)) pf = pd.DataFrame(locations) pf.to_csv(loc_path, header=header, index=False) if __name__ == '__main__': data_root = "./data" ss_root = os.path.join(data_root, "ss") if os.path.exists(ss_root): print("正在刪除{}目錄下原有資料".format(ss_root)) shutil.rmtree(ss_root) print("正在利用選擇性搜尋方法建立資料集: {}".format(ss_root)) select = SelectiveSearch(root=data_root, max_pos_regions=MAX_POSITIVE, max_neg_regions=MAX_NEGATIVE, threshold=IOU_THRESH) select.save(num_workers=os.cpu_count(), method="thread")
透過以上方法，會在./data/ss目錄下為每一張圖片生成一個資料夾，資料夾記憶體放原始影像（.jpg檔案）和邊界框資訊（ss_loc.csv檔案）。其中邊界框資訊檔案結構如圖，依次存放標籤（label）、推薦區域邊界框（gx, gy, gw, gh）和真實邊界框（gx, gy, gw, gh），這些資料將用於後續模型分類和迴歸。
2. 訓練集資料準備
從前面的流程圖可以看出，Fast RCNN輸入有兩個，分別是影像和推薦區域邊界框，前者用於計算特徵圖，後者用於在特徵圖上進行目標區域特徵提取。此外，在多工損失中，還需要區域標籤進行分類損失計算，需要區域邊界框偏移值進行迴歸損失計算。因此，每一幀訓練影像需要包含：原始影像、標籤、推薦區域邊界框和邊界框偏移值。本文對影像進行了隨機翻轉和固定尺寸縮放，以達到資料增強的目的，相應的邊界框等資訊也需要在資料增強過程中重新計算。下面透過程式碼介紹資料準備過程中的一些細節處理。
2.1 資料增強
2.1.1 隨機水平翻轉
假設輸入圖片尺寸為（rows, cols, 3），邊界框為（x, y, w, h），在進行水平翻轉時，原邊界框左上角座標點（x, y）會被翻轉到右上角（rows, cols - 1 - x），翻轉後左上角座標應為（rows, cols - 1 - x - w），而w和h不改變。
隨機水平翻轉程式碼如下：
def random_horizontal_flip(self, im: np.ndarray, ss_boxes: np.ndarray, gt_boxes: np.ndarray): """ 隨機水平翻轉影像 :param im: 輸入影像 :param ss_boxes: 推薦區域邊界框 :param gt_boxes: 邊界框真值 :return: 翻轉後影像和邊界框結果 """ if random.uniform(0, 1) < self.prob_horizontal_flip: rows, cols = im.shape[:2] # 左右翻轉影像 im = np.fliplr(im) # 邊界框位置重新計算 ss_boxes[:, 0] = cols - 1 - ss_boxes[:, 0] - ss_boxes[:, 2] gt_boxes[:, 0] = cols - 1 - gt_boxes[:, 0] - gt_boxes[:, 2] else: pass return im, ss_boxes, gt_boxes
2.1.2 隨機垂直翻轉
假設輸入圖片尺寸為（rows, cols, 3），邊界框為（x, y, w, h），在進行垂直翻轉時，原邊界框左上角座標點（x, y）會被翻轉到左下角（rows - 1 - y, cols），翻轉後左上角座標應為（rows - 1 - y - h, cols），而w和h不改變。
隨機垂直翻轉程式碼如下：
def random_vertical_flip(self, im: np.ndarray, ss_boxes: np.ndarray, gt_boxes: np.ndarray): """ 隨機垂直翻轉影像 :param im: 輸入影像 :param ss_boxes: 推薦區域邊界框 :param gt_boxes: 邊界框真值 :return: 翻轉後影像和邊界框結果 """ if random.uniform(0, 1) < self.prob_vertical_flip: rows, cols = im.shape[:2] # 上下翻轉影像 im = np.flipud(im) # 重新計算邊界框位置 ss_boxes[:, 1] = rows - 1 - ss_boxes[:, 1] - ss_boxes[:, 3] gt_boxes[:, 1] = rows - 1 - gt_boxes[:, 1] - gt_boxes[:, 3] else: pass return im, ss_boxes, gt_boxes
2.1.3 影像縮放
本文中採用vgg16提取影像特徵，在影像輸入vgg16模型前，需要縮放到固定大小（文中採用 512 * 512）。假設輸入圖片尺寸為（rows, cols, 3），邊界框為（x, y, w, h），縮放後尺寸為（im_width, im_height, 3），由於縮放過程中x和w是等比例縮放，y和h也是等比例縮放，則縮放後邊界框為（x * im_width / cols, y * im_height / rows, w * im_width / cols, h * im_height / rows）。
影像縮放程式碼如下：
def resize(im: np.ndarray, im_width: int, im_height: int, ss_boxes: np.ndarray, gt_boxes: Union[np.ndarray, None]): """ 對影像進行縮放 :param im: 輸入影像 :param im_width: 目標影像寬度 :param im_height: 目標影像高度 :param ss_boxes: 推薦區域邊界框->[n, 4] :param gt_boxes: 真實邊界框->[n, 4] :return: 影像和兩種邊界框經過縮放後的結果 """ rows, cols = im.shape[:2] # 影像縮放 im = cv.resize(src=im, dsize=(im_width, im_height), interpolation=cv.INTER_CUBIC) # 計算縮放過程中(x, y, w, h)尺度縮放比例 scale_ratio = np.array([im_width / cols, im_height / rows, im_width / cols, im_height / rows]) # 邊界框也等比例縮放 ss_boxes = (ss_boxes * scale_ratio).astype("int") if gt_boxes is None: return im, ss_boxes gt_boxes = (gt_boxes * scale_ratio).astype("int") return im, ss_boxes, gt_boxes
2.2 邊界框偏移值計算
Fast RCNN中邊界框偏移值採用比例和對數方式計算，避免了邊界框數值大小對訓練的影響。
邊界框偏移值程式碼如下：
def calc_offsets(ss_boxes: np.ndarray, gt_boxes: np.ndarray) -> np.ndarray: """ 計算候選區域與真值間的位置偏移 :param ss_boxes: 候選邊界框 :param gt_boxes: 真值 :return: 邊界框偏移值 """ offsets = np.zeros_like(ss_boxes, dtype="float32") # 基於比例計算偏移值可以不受位置大小的影響 offsets[:, 0] = (gt_boxes[:, 0] - ss_boxes[:, 0]) / ss_boxes[:, 2] offsets[:, 1] = (gt_boxes[:, 1] - ss_boxes[:, 1]) / ss_boxes[:, 3] # 使用log計算w/h的偏移值, 避免值過大 offsets[:, 2] = np.log(gt_boxes[:, 2] / ss_boxes[:, 2]) offsets[:, 3] = np.log(gt_boxes[:, 3] / ss_boxes[:, 3]) return offsets
2.3 迭代資料獲取
我們需要構建包含原始影像（images）、標籤（labels）、推薦區域邊界框（ss_boxes）和邊界框偏移值（offsets）的資料集，假設每個batch有N個原始輸入影像，其中第 i 個影像產生Mi（i=1, 2, ..., N）個推薦區域。由於每個圖象生成的推薦區域數量不固定，那麼相應的標籤、邊界框、偏移值維度也不統一，無法直接繼承torchvision.transforms.Dataset從而構建資料集，因為Dataset要求資料具有相同型別和形狀。因此本文透過batch_size大小來控制每個batch資料量，將同一batch資料存在一個列表中，後續再迭代提取。
資料獲取程式碼如下：
def get_fdata(self): """ 資料集準備 :return: 資料列表 """ fdata = [] if self.shuffle: random.shuffle(self.flist) for num in range(self.num_batch): # 按照batch大小讀取資料 cur_flist = self.flist[num * self.batch_size: (num + 1) * self.batch_size] # 記錄當前batch的影像/推薦區域標籤/邊界框/位置偏移 cur_ims, cur_labels, cur_ss_boxes, cur_offsets = [], [], [], [] for img_path, doc_path in cur_flist: # 讀取影像 img = io.imread(img_path) # 讀取邊界框並堆積打亂框順序 ss_info = pd.read_csv(doc_path, header=0, index_col=None) ss_info = ss_info.sample(frac=1).reset_index(drop=True) labels = ss_info.label.to_list() ss_boxes = ss_info.iloc[:, 1: 5].values gt_boxes = ss_info.iloc[:, 5: 9].values # 資料歸一化 img = self.normalize(im=img) # 隨機翻轉資料增強 img, ss_boxes, gt_boxes = self.random_horizontal_flip(im=img, ss_boxes=ss_boxes, gt_boxes=gt_boxes) img, ss_boxes, gt_boxes = self.random_vertical_flip(im=img, ss_boxes=ss_boxes, gt_boxes=gt_boxes) # 將影像縮放到統一大小 img, ss_boxes, gt_boxes = self.resize(im=img, im_width=self.im_width, im_height=self.im_height, ss_boxes=ss_boxes, gt_boxes=gt_boxes) # 計算最終座標偏移值 offsets = self.calc_offsets(ss_boxes=ss_boxes, gt_boxes=gt_boxes) # 轉換為tensor im_tensor = torch.tensor(np.transpose(img, (2, 0, 1))) ss_boxes_tensor = torch.tensor(data=ss_boxes) cur_ims.append(im_tensor) cur_labels.extend(labels) cur_ss_boxes.append(ss_boxes_tensor) cur_offsets.extend(offsets) # 每個batch資料放一起方便後續訓練呼叫 cur_ims = torch.stack(cur_ims) cur_labels = torch.tensor(cur_labels) cur_offsets = torch.tensor(np.array(cur_offsets)) fdata.append([cur_ims, cur_labels, cur_ss_boxes, cur_offsets]) return fdata
透過上述流程，本文實現實現了一個可迭代的資料集類，每次迭代返回一個batch的資料，用於模型訓練使用，完整程式碼如下：
class GenDataSet: def __init__(self, root, im_width, im_height, batch_size, shuffle=False, prob_vertical_flip=0.5, prob_horizontal_flip=0.5): """ 初始化GenDataSet :param root: 資料路徑 :param im_width: 目標圖片寬度 :param im_height: 目標圖片高度 :param batch_size: 批資料大小 :param shuffle: 是否隨機打亂批資料 :param prob_vertical_flip: 隨機垂直翻轉機率 :param prob_horizontal_flip: 隨機水平翻轉機率 """ self.root = root self.im_width, self.im_height = (im_width, im_height) self.batch_size = batch_size self.shuffle = shuffle self.flist = self.get_flist() self.num_batch = self.calc_num_batch() self.prob_vertical_flip = prob_vertical_flip self.prob_horizontal_flip = prob_horizontal_flip def get_flist(self) -> list: """ 獲取原始影像和推薦區域邊界框 :return: [影像, 邊界框]列表 """ flist = [] for roots, dirs, files in os.walk(self.root): for file in files: if not file.endswith(".jpg"): continue img_path = os.path.join(roots, file) doc_path = os.path.join(roots, "ss_loc.csv") if not os.path.exists(doc_path): continue flist.append((img_path, doc_path)) return flist def calc_num_batch(self) -> int: """ 計算batch數量 :return: 批資料數量 """ total = len(self.flist) if total % self.batch_size == 0: num_batch = total // self.batch_size else: num_batch = total // self.batch_size + 1 return num_batch @staticmethod def normalize(im: np.ndarray) -> np.ndarray: """ 將影像資料歸一化 :param im: 輸入影像->uint8 :return: 歸一化影像->float32 """ if im.dtype != np.uint8: raise TypeError("uint8 img is required.") else: im = im / 255.0 im = im.astype("float32") return im @staticmethod def calc_offsets(ss_boxes: np.ndarray, gt_boxes: np.ndarray) -> np.ndarray: """ 計算候選區域與真值間的位置偏移 :param ss_boxes: 候選邊界框 :param gt_boxes: 真值 :return: 邊界框偏移值 """ offsets = np.zeros_like(ss_boxes, dtype="float32") # 基於比例計算偏移值可以不受位置大小的影響 offsets[:, 0] = (gt_boxes[:, 0] - ss_boxes[:, 0]) / ss_boxes[:, 2] offsets[:, 1] = (gt_boxes[:, 1] - ss_boxes[:, 1]) / ss_boxes[:, 3] # 使用log計算w/h的偏移值, 避免值過大 offsets[:, 2] = np.log(gt_boxes[:, 2] / ss_boxes[:, 2]) offsets[:, 3] = np.log(gt_boxes[:, 3] / ss_boxes[:, 3]) return offsets @staticmethod def resize(im: np.ndarray, im_width: int, im_height: int, ss_boxes: np.ndarray, gt_boxes: Union[np.ndarray, None]): """ 對影像進行縮放 :param im: 輸入影像 :param im_width: 目標影像寬度 :param im_height: 目標影像高度 :param ss_boxes: 推薦區域邊界框->[n, 4] :param gt_boxes: 真實邊界框->[n, 4] :return: 影像和兩種邊界框經過縮放後的結果 """ rows, cols = im.shape[:2] # 影像縮放 im = cv.resize(src=im, dsize=(im_width, im_height), interpolation=cv.INTER_CUBIC) # 計算縮放過程中(x, y, w, h)尺度縮放比例 scale_ratio = np.array([im_width / cols, im_height / rows, im_width / cols, im_height / rows]) # 邊界框也等比例縮放 ss_boxes = (ss_boxes * scale_ratio).astype("int") if gt_boxes is None: return im, ss_boxes gt_boxes = (gt_boxes * scale_ratio).astype("int") return im, ss_boxes, gt_boxes def random_horizontal_flip(self, im: np.ndarray, ss_boxes: np.ndarray, gt_boxes: np.ndarray): """ 隨機水平翻轉影像 :param im: 輸入影像 :param ss_boxes: 推薦區域邊界框 :param gt_boxes: 邊界框真值 :return: 翻轉後影像和邊界框結果 """ if random.uniform(0, 1) < self.prob_horizontal_flip: rows, cols = im.shape[:2] # 左右翻轉影像 im = np.fliplr(im) # 邊界框位置重新計算 ss_boxes[:, 0] = cols - 1 - ss_boxes[:, 0] - ss_boxes[:, 2] gt_boxes[:, 0] = cols - 1 - gt_boxes[:, 0] - gt_boxes[:, 2] else: pass return im, ss_boxes, gt_boxes def random_vertical_flip(self, im: np.ndarray, ss_boxes: np.ndarray, gt_boxes: np.ndarray): """ 隨機垂直翻轉影像 :param im: 輸入影像 :param ss_boxes: 推薦區域邊界框 :param gt_boxes: 邊界框真值 :return: 翻轉後影像和邊界框結果 """ if random.uniform(0, 1) < self.prob_vertical_flip: rows, cols = im.shape[:2] # 上下翻轉影像 im = np.flipud(im) # 重新計算邊界框位置 ss_boxes[:, 1] = rows - 1 - ss_boxes[:, 1] - ss_boxes[:, 3] gt_boxes[:, 1] = rows - 1 - gt_boxes[:, 1] - gt_boxes[:, 3] else: pass return im, ss_boxes, gt_boxes def get_fdata(self): """ 資料集準備 :return: 資料列表 """ fdata = [] if self.shuffle: random.shuffle(self.flist) for num in range(self.num_batch): # 按照batch大小讀取資料 cur_flist = self.flist[num * self.batch_size: (num + 1) * self.batch_size] # 記錄當前batch的影像/推薦區域標籤/邊界框/位置偏移 cur_ims, cur_labels, cur_ss_boxes, cur_offsets = [], [], [], [] for img_path, doc_path in cur_flist: # 讀取影像 img = io.imread(img_path) # 讀取邊界框並堆積打亂框順序 ss_info = pd.read_csv(doc_path, header=0, index_col=None) ss_info = ss_info.sample(frac=1).reset_index(drop=True) labels = ss_info.label.to_list() ss_boxes = ss_info.iloc[:, 1: 5].values gt_boxes = ss_info.iloc[:, 5: 9].values # 資料歸一化 img = self.normalize(im=img) # 隨機翻轉資料增強 img, ss_boxes, gt_boxes = self.random_horizontal_flip(im=img, ss_boxes=ss_boxes, gt_boxes=gt_boxes) img, ss_boxes, gt_boxes = self.random_vertical_flip(im=img, ss_boxes=ss_boxes, gt_boxes=gt_boxes) # 將影像縮放到統一大小 img, ss_boxes, gt_boxes = self.resize(im=img, im_width=self.im_width, im_height=self.im_height, ss_boxes=ss_boxes, gt_boxes=gt_boxes) # 計算最終座標偏移值 offsets = self.calc_offsets(ss_boxes=ss_boxes, gt_boxes=gt_boxes) # 轉換為tensor im_tensor = torch.tensor(np.transpose(img, (2, 0, 1))) ss_boxes_tensor = torch.tensor(data=ss_boxes) cur_ims.append(im_tensor) cur_labels.extend(labels) cur_ss_boxes.append(ss_boxes_tensor) cur_offsets.extend(offsets) # 每個batch資料放一起方便後續訓練呼叫 cur_ims = torch.stack(cur_ims) cur_labels = torch.tensor(cur_labels) cur_offsets = torch.tensor(np.array(cur_offsets)) fdata.append([cur_ims, cur_labels, cur_ss_boxes, cur_offsets]) return fdata def __len__(self): # 以batch數量定義資料集大小 return self.num_batch def __iter__(self): self.fdata = self.get_fdata() self.index = 0 return self def __next__(self): if self.index >= self.num_batch: raise StopIteration # 生成當前batch資料 value = self.fdata[self.index] self.index += 1 return value
3. 模型訓練
Fast RCNN的訓練流程是：CNN獲取特徵圖 → ROI_POOL提取候選區域特徵 → 獲取分類器和迴歸器結果 → 多工損失引數調優。可知，Fast RCNN模型結構中需要依次實現影像特徵提取器features、ROI池化、分類器classifier和迴歸器regressor，訓練過程中需要構建多工損失函式。下文詳細介紹相關結構和程式碼。
3.1 Fast RCNN模型結構
3.1.1 特徵提取器
本文中採用vgg16_bn作為模型特徵提取器，直接呼叫torchvison.models中預訓練模型即可，程式碼如下：
# 採用vgg16_bn作為backbone self.features = models.vgg16_bn(pretrained=True).features
3.1.2 ROI 池化
ROI池化的作用是在特徵圖上進行候選區域特徵的抽取，同時將抽取的特徵縮放到固定大小，方便全連線層型別的分類器和迴歸器使用。其具體實現過程如下：
以本文介紹的Fast RCNN模型為例，其輸入影像張量大小為 [3, 512, 512]，經過vgg16_bn特徵提取得到輸出feature維度為 [512, 16, 16]。過程中資料經過5次2*2的MaxPool，特徵進行了32倍縮放，相應地原邊界框位置和大小也會進行等比例縮放。
假設輸入模型的某一邊界框 ss_box = (50, 72, 260, 318)，經過等比例縮放後對應到特徵圖上的候選邊界框 ss_box' = (50/32, 72/32, 260/32, 318/32) = (1.56, 2.25, 8.13, 9.94)。ss_box'值存在小數，此時的處理方法是直接向下取整，得到ss_box' = (1, 2, 8, 9）,也就是說特徵圖上對應的 roi_feature = features[:, 2: 2 + 9, 1: 1 + 8]，對應維度為 [512, 9, 8]。該特徵是需要輸入分類器和迴歸器的，由於兩者是全連線層結構，輸入特徵尺寸是固定的，因此需要使用一定方法將其縮放到對應尺寸。本文直接採用 torch.nn.AdaptiveMaxPool2d進行縮放，當然你也可以採取其他方式，比如插值方法。
ROI池化具體實現程式碼如下：
def roi_pool(self, im_features: Tensor, ss_boxes: list): """ 提取推薦區域特徵圖並縮放到固定大小 :param im_features: backbone輸出的影像特徵->[batch, channel, rows, cols] :param ss_boxes: 推薦區域邊界框資訊->[batch, num, 4] :return: 推薦區域特徵 """ roi_features = [] for im_idx, im_feature in enumerate(im_features): im_boxes = ss_boxes[im_idx] for box in im_boxes: # 輸入全圖經過backbone後空間位置需進行縮放, 利用空間縮放比例將box位置對應到feature上 fx, fy, fw, fh = [int(p / self.spatial_scale) for p in box] # 縮放後維度不足1個pixel, 是由於int取整導致, 仍取1個pixel防止維度為0 if fw == 0: fw = 1 if fh == 0: fh = 1 # 在特徵圖上提取候選區域對應的區域特徵 roi_feature = im_feature[:, fy: fy + fh, fx: fx + fw] # 將區域特徵池化到固定大小 roi_feature = self.pool(roi_feature) # 將池化後特徵展開方便後續送入分類器和迴歸器 roi_feature = roi_feature.view(-1) roi_features.append(roi_feature) # 轉換成tensor roi_features = torch.stack(roi_features) return roi_features
其中
# 自適應最大值池化將推薦區域特徵池化到固定大小 self.pool = nn.AdaptiveMaxPool2d((self.pool_size, self.pool_size))
3.1.3 分類器和迴歸器
分類器和迴歸器是全連線層結構，由於vgg16_bn輸出通道為512，經過區域池化後特徵長寬固定，因此將其定義如下：
# 分類器, 輸入為vgg16的512通道特徵經過roi_pool的結果 self.classifier = nn.Sequential( nn.Linear(in_features=512 * self.pool_size * self.pool_size, out_features=32), nn.ReLU(), nn.Dropout(p=drop), nn.Linear(in_features=32, out_features=num_classes) ) # 迴歸器, 輸入為vgg16的512通道特徵經過roi_pool的結果 self.regressor = nn.Sequential( nn.Linear(in_features=512 * self.pool_size * self.pool_size, out_features=64), nn.ReLU(), nn.Dropout(p=drop), nn.Linear(in_features=64, out_features=4) )
3.2 多工損失
Fast RCNN損失函式由分類損失（CrossEntropy Loss）和迴歸損失（SmoothL1 Loss）兩部分構成。其中分類模型需要區分候選區域類別，以判定物體還是背景，因此需要對所有輸出進行計算損失。迴歸模型只需要對物體邊界框進行矯正，因此只計算非背景區域的損失。
多工損失程式碼如下：
def multitask_loss(output: tuple, labels: Tensor, offsets: Tensor, criterion: list, alpha: float = 1.0): """ 計算多工損失 :param output: 模型輸出 :param labels: 邊界框標籤 :param offsets: 邊界框偏移值 :param criterion: 損失函式 :param alpha: 損失函式 :return: """ output_cls, output_reg = output # 計算分類損失 loss_cls = criterion[0](output_cls, labels) # 計算正樣本的迴歸損失 output_reg_valid = output_reg[labels != 0] offsets_valid = offsets[labels != 0] loss_reg = criterion[1](output_reg_valid, offsets_valid) # 損失加權 loss = loss_cls + alpha * loss_reg return loss
3.3 訓練程式碼
Fast RCNN無需RCNN那樣分階段的繁瑣訓練，直接同時訓練特徵提取器、分類器和迴歸器。
模型訓練階段程式碼如下：
# Train.py import os import matplotlib.pyplot as plt import numpy as np import torch from torch import nn from torch.optim.lr_scheduler import StepLR from Utils import FastRCNN, GenDataSet from torch import Tensor from Config import * def multitask_loss(output: tuple, labels: Tensor, offsets: Tensor, criterion: list, alpha: float = 1.0): """ 計算多工損失 :param output: 模型輸出 :param labels: 邊界框標籤 :param offsets: 邊界框偏移值 :param criterion: 損失函式 :param alpha: 損失函式 :return: """ output_cls, output_reg = output # 計算分類損失 loss_cls = criterion[0](output_cls, labels) # 計算正樣本的迴歸損失 output_reg_valid = output_reg[labels != 0] offsets_valid = offsets[labels != 0] loss_reg = criterion[1](output_reg_valid, offsets_valid) # 損失加權 loss = loss_cls + alpha * loss_reg return loss def train(data_set, network, num_epochs, optimizer, scheduler, criterion, device, train_rate=0.8): """ 模型訓練 :param data_set: 訓練資料集 :param network: 網路結構 :param num_epochs: 訓練輪次 :param optimizer: 最佳化器 :param scheduler: 學習率排程器 :param criterion: 損失函式 :param device: CPU/GPU :param train_rate: 訓練集比例 :return: None """ os.makedirs('./model', exist_ok=True) network = network.to(device) best_loss = np.inf print("=" * 8 + "開始訓練模型" + "=" * 8) # 計算訓練batch數量 batch_num = len(data_set) train_batch_num = round(batch_num * train_rate) # 記錄訓練過程中每一輪損失和準確率 train_loss_all, val_loss_all, train_acc_all, val_acc_all = [], [], [], [] for epoch in range(num_epochs): # 記錄train/val分類準確率和總損失 num_train_acc = num_val_acc = num_train_loss = num_val_loss = 0 train_loss = val_loss = 0.0 train_corrects = val_corrects = 0 for step, batch_data in enumerate(data_set): # 讀取資料 ims, labels, ss_boxes, offsets = batch_data ims = ims.to(device) labels = labels.to(device) ss_boxes = [ss.to(device) for ss in ss_boxes] offsets = offsets.to(device) # 模型輸入為全圖和推薦區域邊界框, 即[ims: Tensor, ss_boxes: list[Tensor]] inputs = [ims, ss_boxes] if step < train_batch_num: # train network.train() output = network(inputs) loss = multitask_loss(output=output, labels=labels, offsets=offsets, criterion=criterion) optimizer.zero_grad() loss.backward() optimizer.step() # 計算每個batch分類正確的數量和loss label_hat = torch.argmax(output[0], dim=1) train_corrects += (label_hat == labels).sum().item() num_train_acc += labels.size(0) # 計算每個batch總損失 train_loss += loss.item() * ims.size(0) num_train_loss += ims.size(0) else: # validation network.eval() with torch.no_grad(): output = network(inputs) loss = multitask_loss(output=output, labels=labels, offsets=offsets, criterion=criterion) # 計算每個batch分類正確的數量和loss和 label_hat = torch.argmax(output[0], dim=1) val_corrects += (label_hat == labels).sum().item() num_val_acc += labels.size(0) val_loss += loss.item() * ims.size(0) num_val_loss += ims.size(0) scheduler.step() # 記錄loss和acc變化曲線 train_loss_all.append(train_loss / num_train_loss) val_loss_all.append(val_loss / num_val_loss) train_acc_all.append(100 * train_corrects / num_train_acc) val_acc_all.append(100 * val_corrects / num_val_acc) print("Epoch:[{:0>3}|{}] train_loss:{:.3f} train_acc:{:.2f}% val_loss:{:.3f} val_acc:{:.2f}%".format( epoch + 1, num_epochs, train_loss_all[-1], train_acc_all[-1], val_loss_all[-1], val_acc_all[-1] )) # 儲存模型 if val_loss_all[-1] < best_loss: best_loss = val_loss_all[-1] save_path = os.path.join("./model", "model.pth") torch.save(network, save_path) # 繪製訓練曲線 fig_path = os.path.join("./model/", "train_curve.png") plt.subplot(121) plt.plot(range(num_epochs), train_loss_all, "r-", label="train") plt.plot(range(num_epochs), val_loss_all, "b-", label="val") plt.title("Loss") plt.legend() plt.subplot(122) plt.plot(range(num_epochs), train_acc_all, "r-", label="train") plt.plot(range(num_epochs), val_acc_all, "b-", label="val") plt.title("Acc") plt.legend() plt.tight_layout() plt.savefig(fig_path) plt.close() return None if __name__ == "__main__": if not os.path.exists("./data/ss"): raise FileNotFoundError("資料不存在, 請先執行SelectiveSearch.py生成目標區域") device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = FastRCNN(num_classes=CLASSES, in_size=IM_SIZE, pool_size=POOL_SIZE, spatial_scale=SCALE, device=device) exit(0) criterion = [nn.CrossEntropyLoss(), nn.SmoothL1Loss()] optimizer = torch.optim.Adam(params=model.parameters(), lr=LR) scheduler = StepLR(optimizer, step_size=STEP, gamma=GAMMA) model_root = "./model" os.makedirs(model_root, exist_ok=True) # 在生成的ss資料上進行預訓練 train_root = "./data/ss" train_set = GenDataSet(root=train_root, im_width=IM_SIZE, im_height=IM_SIZE, batch_size=BATCH_SIZE, shuffle=True) train(data_set=train_set, network=model, num_epochs=EPOCHS, optimizer=optimizer, scheduler=scheduler, criterion=criterion, device=device)
4. 模型預測
透過上述流程，訓練好了Fast RCNN模型。可以輸入圖片對目標進行預測，但是有三點值得注意：
a. 模型輸出包含分類和迴歸兩部分結果，其中迴歸器輸出為邊界框的偏移值，需要利用偏移值對推薦區域位置先進性修正；
b. 模型輸入被resize到了 [3, 512, 512] ，預測得到的邊界框位置也是相對於縮放後影像，需要重新對映到原圖中；
c. 預測結果中可能存在很多冗餘邊界框，需要設計去除。
4.1 預測邊界框位置修正
2.2中已經介紹邊界框偏移值是採用比例和對數方式計算，現在只需反向操作即可根據偏移值和推薦邊界框計算出修正後的結果。
邊界框位置修正程式碼如下：
def rectify_bbox(ss_boxes: np.ndarray, offsets: np.ndarray) -> np.ndarray: """ 修正邊界框位置 :param ss_boxes: 邊界框 :param offsets: 邊界框偏移值 :return: 位置修正後的邊界框 """ # 和Utils.GenDataSet.calc_offsets過程相反 ss_boxes[:, 0] = ss_boxes[:, 2] * offsets[:, 0] + ss_boxes[:, 0] ss_boxes[:, 1] = ss_boxes[:, 3] * offsets[:, 1] + ss_boxes[:, 1] ss_boxes[:, 2] = np.exp(offsets[:, 2]) * ss_boxes[:, 2] ss_boxes[:, 3] = np.exp(offsets[:, 3]) * ss_boxes[:, 3] boxes = ss_boxes.astype("int") return boxes
4.2 預測邊界框位置對映
模型輸入資料是經過resize方式得到的，在邊界框對映回原圖時，也只需要計算縮放比例反向對映即可。
邊界框位置對映程式碼如下：
def map_bbox_to_img(boxes: np.ndarray, src_img: np.ndarray, im_width: int, im_height: int): """ 根據縮放比例將邊界框對映回原圖 :param boxes: 縮放後影像上的邊界框 :param src_img: 原始影像 :param im_width: 縮放後影像寬度 :param im_height: 縮放後影像高度 :return: boxes->對映到原影像上的邊界框 """ rows, cols = src_img.shape[:2] scale_ratio = np.array([cols / im_width, rows / im_height, cols / im_width, rows / im_height]) boxes = (boxes * scale_ratio).astype("int") return boxes
4.3 非極大值抑制
和RCNN一樣，Fast RCNN也透過ss方法產生大量的候選區域（~2k）。雖然分類器能夠去除大量背景區域，但是仍然有較多的目標區域，會得較多的目標檢測框。這些檢測框大部分都是重疊的，需要進行篩選。非極大值抑制（Non-Maximum Supression，後續稱之為 nms）就是一種候選框選取方法。
4.3.1 nms演算法流程
非極大值抑制透過目標置信度對邊界框進行篩選，其具體流程如下：
（1）初始化輸出列表 out_bboxes = [ ]；
（2）獲取輸入邊界框 in_bboxes 對應的目標類別置信度；
（3）將照置信度由高到低的方式對邊界框進行排序；
（4）選取置信度最高的邊界框 bbox_max，將其新增到 out_bboxes 中，並計算它與其他邊界框的IoU結果；
（5）IoU大於閾值的表明兩區域較為接近，邊界框重疊性較高，將這些框和bbox_max從 in_bboxes 中移除；
（6）重複執行3~5過程，直到 in_bboxes 為空，此時 out_bboxes 即最終保留的邊界框結果。
非極大值抑制的程式碼實現如下：
def nms(bboxes: np.ndarray, scores: np.ndarray, threshold: float) -> np.ndarray: """ 非極大值抑制去除冗餘邊界框 :param bboxes: 目標邊界框 :param scores: 目標得分 :param threshold: 閾值 :return: keep->保留下的有效邊界框 """ # 獲取邊界框和分數 x1 = bboxes[:, 0] y1 = bboxes[:, 1] x2 = bboxes[:, 0] + bboxes[:, 2] - 1 y2 = bboxes[:, 1] + bboxes[:, 3] - 1 # 計算面積 areas = (x2 - x1 + 1) * (y2 - y1 + 1) # 逆序排序 order = scores.argsort()[::-1] keep = [] while order.size > 0: # 取分數最高的一個 i = order[0] keep.append(bboxes[i]) if order.size == 1: break # 計算相交區域 xx1 = np.maximum(x1[i], x1[order[1:]]) xx2 = np.minimum(x2[i], x2[order[1:]]) yy1 = np.maximum(y1[i], y1[order[1:]]) yy2 = np.minimum(y2[i], y2[order[1:]]) # 計算IoU inter = np.maximum(0.0, xx2 - xx1 + 1) * np.maximum(0.0, yy2 - yy1 + 1) iou = inter / (areas[i] + areas[order[1:]] - inter) # 保留IoU小於閾值的bbox idx = np.where(iou <= threshold)[0] order = order[idx + 1] keep = np.array(keep) return keep
4.3.2 nms結果
如下，左圖為未採用nms邊界框結果，右圖為nms處理後結果，表明nms能夠較好地去除冗餘邊界框。
4.4 預測程式碼
結合上述流程，可以完成Fast RCNN對目標的預測，其預測階段程式碼如下：
# Predict.py import os import torch import numpy as np import skimage.io as io from Utils import GenDataSet, draw_box from torch.nn.functional import softmax from SelectiveSearch import SelectiveSearch as ss from Config import * def rectify_bbox(ss_boxes: np.ndarray, offsets: np.ndarray) -> np.ndarray: """ 修正邊界框位置 :param ss_boxes: 邊界框 :param offsets: 邊界框偏移值 :return: 位置修正後的邊界框 """ # 和Utils.GenDataSet.calc_offsets過程相反 ss_boxes = ss_boxes.astype("float32") ss_boxes[:, 0] = ss_boxes[:, 2] * offsets[:, 0] + ss_boxes[:, 0] ss_boxes[:, 1] = ss_boxes[:, 3] * offsets[:, 1] + ss_boxes[:, 1] ss_boxes[:, 2] = np.exp(offsets[:, 2]) * ss_boxes[:, 2] ss_boxes[:, 3] = np.exp(offsets[:, 3]) * ss_boxes[:, 3] boxes = ss_boxes.astype("int") return boxes def map_bbox_to_img(boxes: np.ndarray, src_img: np.ndarray, im_width: int, im_height: int): """ 根據縮放比例將邊界框對映回原圖 :param boxes: 縮放後影像上的邊界框 :param src_img: 原始影像 :param im_width: 縮放後影像寬度 :param im_height: 縮放後影像高度 :return: boxes->對映到原影像上的邊界框 """ rows, cols = src_img.shape[:2] scale_ratio = np.array([cols / im_width, rows / im_height, cols / im_width, rows / im_height]) boxes = (boxes * scale_ratio).astype("int") return boxes def nms(bboxes: np.ndarray, scores: np.ndarray, threshold: float) -> np.ndarray: """ 非極大值抑制去除冗餘邊界框 :param bboxes: 目標邊界框 :param scores: 目標得分 :param threshold: 閾值 :return: keep->保留下的有效邊界框 """ # 獲取邊界框和分數 x1 = bboxes[:, 0] y1 = bboxes[:, 1] x2 = bboxes[:, 0] + bboxes[:, 2] - 1 y2 = bboxes[:, 1] + bboxes[:, 3] - 1 # 計算面積 areas = (x2 - x1 + 1) * (y2 - y1 + 1) # 逆序排序 order = scores.argsort()[::-1] keep = [] while order.size > 0: # 取分數最高的一個 i = order[0] keep.append(bboxes[i]) if order.size == 1: break # 計算相交區域 xx1 = np.maximum(x1[i], x1[order[1:]]) xx2 = np.minimum(x2[i], x2[order[1:]]) yy1 = np.maximum(y1[i], y1[order[1:]]) yy2 = np.minimum(y2[i], y2[order[1:]]) # 計算IoU inter = np.maximum(0.0, xx2 - xx1 + 1) * np.maximum(0.0, yy2 - yy1 + 1) iou = inter / (areas[i] + areas[order[1:]] - inter) # 保留IoU小於閾值的bbox idx = np.where(iou <= threshold)[0] order = order[idx + 1] keep = np.array(keep) return keep def predict(network, im, im_width, im_height, device, nms_thresh=None, save_name=None): """ 模型預測 :param network: 模型結構 :param im: 輸入影像 :param im_width: 模型輸入影像寬度 :param im_height: 模型輸入影像長度 :param device: CPU/GPU :param nms_thresh: 非極大值抑制閾值 :param save_name: 儲存檔名 :return: None """ network.eval() # 生成推薦區域 ss_boxes_src = ss.cal_proposals(img=im) # 資料歸一化 im_norm = GenDataSet.normalize(im=im) # 將影像縮放固定大小, 並將邊界框對映到縮放後影像上 im_rsz, ss_boxes_rsz = GenDataSet.resize(im=im_norm, im_width=im_width, im_height=im_height, ss_boxes=ss_boxes_src, gt_boxes=None) im_tensor = torch.tensor(np.transpose(im_rsz, (2, 0, 1))).unsqueeze(0).to(device) ss_boxes_tensor = torch.tensor(ss_boxes_rsz).to(device) # 模型輸入為[img: Tensor, ss_boxes: list(Tensor)] inputs = [im_tensor, [ss_boxes_tensor]] with torch.no_grad(): outputs = network(inputs) # 計算各個類別的分類得分 scores = softmax(input=outputs[0], dim=1) scores = scores.cpu().numpy() # 獲取位置偏移 offsets = outputs[1] offsets = offsets.cpu().numpy() # 根據模型計算出的offsets對推薦區域邊界框位置進行修正 out_boxes = rectify_bbox(ss_boxes=ss_boxes_rsz, offsets=offsets) # 將邊界框位置對映回原始影像 out_boxes = map_bbox_to_img(boxes=out_boxes, src_img=im, im_width=im_width, im_height=im_height) # 邊界框篩選 predicted_boxes = [] for i in range(1, CLASSES): # 獲取當前類別目標得分 cur_obj_scores = scores[:, i] # 只選取置信度滿足閾值要求的預測框 idx = cur_obj_scores >= CONFIDENCE_THRESH valid_scores = cur_obj_scores[idx] valid_out_boxes = out_boxes[idx] # 遍歷物體類別, 對每個類別的邊界框預測結果進行非極大值抑制 if nms_thresh is not None: used_boxes = nms(bboxes=valid_out_boxes, scores=valid_scores, threshold=nms_thresh) else: used_boxes = valid_out_boxes # 可能存在值不符合要求的情況, 需要剔除 for j in range(used_boxes.shape[0]): if used_boxes[j, 0] < 0 or used_boxes[j, 1] < 0 or used_boxes[j, 2] <= 0 or used_boxes[j, 3] <= 0: continue predicted_boxes.append(used_boxes[j]) draw_box(img=im, boxes=predicted_boxes, save_name=save_name) return None if __name__ == "__main__": device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model_path = "./model/model.pth" model = torch.load(model_path, map_location=device) test_root = "./data/" for roots, dirs, files in os.walk(test_root): for file in files: if not file.endswith(".jpg"): continue save_name = file.split(".")[0] im_path = os.path.join(roots, file) im = io.imread(im_path) predict(network=model, im=im, im_width=IM_SIZE, im_height=IM_SIZE, device=device, nms_thresh=NMS_THRESH, save_name=save_name)
4.5 預測結果
如下為Fast RCNN在花朵資料集上預測結果展示，左圖中當同一圖中存在多個相同類別目標且距離較近時，邊界框並沒有很好地檢測出每個個體。
三、演算法缺點
相較於RCNN演算法，Fast RCNN演算法極大的縮短了檢測時間，但是整個過程仍需要使用ss方法生成候選區域，總體時間消耗仍然不適用於實時檢測任務。
四、資料和程式碼
本文中資料和詳細工程程式碼實現請移步：https://github.com/jchsun1/Fast-RCNN
Reference：
本文結束，祝君好運。