<bdo id="4g88a"><xmp id="4g88a">
  • <legend id="4g88a"><code id="4g88a"></code></legend>

    【目標檢測】Fast R-CNN算法實現

    一、前言

    2014年,Ross Girshick提出RCNN,成為目標檢測領域的開山之作。一年后,借鑒空間金字塔池化思想,Ross Girshick推出設計更為巧妙的Fast RCNN(https://github.com/rbgirshick/fast-rcnn),極大地提高了檢測速度。Fast RCNN的提出解決了RCNN結構固有的三個弊端:

    1. 繁瑣的多階段訓練:RCNN在訓練時,首先需要在推薦區域上微調卷積網絡,然后利用提取的卷積特征針對每個類別訓練一個SVM分類器,最后還需要基于卷積特征進行邊界框回歸訓練;
    2. 高空間和時間成本:RCNN在訓練SVM和回歸器時,需在磁盤上對推薦區域的卷積特征進行讀寫,內存和時間消耗較為嚴重;
    3. 檢測速度慢:檢測時,需要對每個推薦區域進行特征提取,計算重復性高,導致檢測速度很慢。

    同前面RCNN實現一樣(見 http://www.diao-diao.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的實現流程:

    1. 候選區域生成:利用ss方法為每一幀圖片生成數量不固定的候選區域,利用IoU結果對候選區域進行標注,同時記錄標簽和邊界框信息;
    2. 訓練集數據準備:構建可迭代數據集類,為模型訓練提供原始圖像、標簽、推薦區域邊界框和邊界框偏移值信息;
    3. 模型訓練:利用ROI池化提取候選區域特征,利用多任務損失同時訓練分類和回歸模型;
    4. 模型預測:通過ss方法生成推薦區域,利用模型結果對推薦區域位置進行修正,再利用非極大值抑制剔除冗余邊界框,獲得最終目標框。

    1. 候選區域生成

    Fast RCNNt同RCNN一樣采用選擇性搜索(selective search,后面簡稱為ss)的辦法產生候選區域,ss方法的詳細思路同樣請參考 http://www.diao-diao.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")
    View Code

    通過以上方法,會在./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
    View Code

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

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

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

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

    通過上述流程,本文實現實現了一個可迭代的數據集類,每次迭代返回一個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
    View Code

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

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

    其中

    # 自適應最大值池化將推薦區域特征池化到固定大小
    self.pool = nn.AdaptiveMaxPool2d((self.pool_size, self.pool_size))
    View Code

    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)
    )
    View Code

    3.2 多任務損失

    Fast RCNN損失函數由分類損失(CrossEntropy Loss)和回歸損失(SmoothL1 Loss)兩部分構成。其中分類模型需要區分候選區域類別,以判定物體還是背景,因此需要對所有輸出進行計算損失?;貧w模型只需要對物體邊界框進行矯正,因此只計算非背景區域的損失。

    多任務損失代碼如下:

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

    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)
        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)
    View Code

    4. 模型預測

    通過上述流程,訓練好了Fast RCNN模型??梢暂斎雸D片對目標進行預測,但是有三點值得注意:

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

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

    4.3 非極大值抑制

    和RCNN一樣,Fast RCNN也通過ss方法產生大量的候選區域(~2k)。雖然分類器能夠去除大量背景區域,但是仍然有較多的目標區域,會得較多的目標檢測框。這些檢測框大部分都是重疊的,需要進行篩選。非極大值抑制(Non-Maximum Supression,后續稱之為 nms)就是一種候選框選取方法。

    4.3.1 nms算法流程

    非極大值抑制通過目標置信度對邊界框進行篩選,其具體流程如下:

     ?。?)初始化輸出列表 out_bboxes = [ ];

     ?。?)獲取輸入邊界框 in_bboxes 對應的目標類別置信度;

     ?。?)將照置信度由高到低的方式對邊界框進行排序;

     ?。?)選取置信度最高的邊界框 bbox_max,將其添加到 out_bboxes 中,并計算它與其他邊界框的IoU結果;

     ?。?)IoU大于閾值的表明兩區域較為接近,邊界框重疊性較高,將這些框和bbox_max從 in_bboxes 中移除;

     ?。?)重復執行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
    View Code

    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)
    View Code

    4.5 預測結果

    如下為Fast RCNN在花朵數據集上預測結果展示,左圖中當同一圖中存在多個相同類別目標且距離較近時,邊界框并沒有很好地檢測出每個個體。

    三、算法缺點

    相較于RCNN算法,Fast RCNN算法極大的縮短了檢測時間,但是整個過程仍需要使用ss方法生成候選區域,總體時間消耗仍然不適用于實時檢測任務。

    四、數據和代碼

    本文中數據和詳細工程代碼實現請移步:https://github.com/jchsun1/Fast-RCNN

    Reference


     本文至此結束,如有疑惑歡迎留言交流。 

    posted @ 2023-09-20 17:18  萬象為賓客  閱讀(667)  評論(5編輯  收藏  舉報
    免费视频精品一区二区_日韩一区二区三区精品_aaa在线观看免费完整版_世界一级真人片
    <bdo id="4g88a"><xmp id="4g88a">
  • <legend id="4g88a"><code id="4g88a"></code></legend>