Tensorflow2.0 實現 YOLOv3(八):test.py

文章說明

本系列文章旨在對 Github 上 malin9402 提供的代碼進行說明,在這篇文章中,我們會對 YOLOv3 項目中的 test.py 文件進行說明。

如果只是想運行 Github 上的代碼,可以參考對 YOLOv3 代碼的說明一文。

導入需要的庫

import cv2
import os
import shutil
import numpy as np
import tensorflow as tf
import core.utils as utils
from core.config import cfg
from core.yolov3 import YOLOv3, decode

設置參數

INPUT_SIZE = 416  # 測試集中圖片的尺寸
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))  # 類別的數目
CLASSES = utils.read_class_names(cfg.YOLO.CLASSES)  # 類別的名稱索引

創建保存路徑

# 路徑名稱
predicted_dir_path = '../mAP/predicted'
ground_truth_dir_path = '../mAP/ground-truth'
# 若路徑下有文件,則刪除
if os.path.exists(predicted_dir_path):
    shutil.rmtree(predicted_dir_path)
if os.path.exists(ground_truth_dir_path):
    shutil.rmtree(ground_truth_dir_path)
if os.path.exists(cfg.TEST.DECTECTED_IMAGE_PATH):
    shutil.rmtree(cfg.TEST.DECTECTED_IMAGE_PATH)
# 創建路徑
os.mkdir(predicted_dir_path)
os.mkdir(ground_truth_dir_path)
os.mkdir(cfg.TEST.DECTECTED_IMAGE_PATH)

構建模型

# 確定模型輸入
input_layer = tf.keras.layers.Input([INPUT_SIZE, INPUT_SIZE, 3])
# 確定模型輸出
feature_maps = YOLOv3(input_layer)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
    bbox_tensor = decode(fm, i)  # 將檢測框解碼到原圖上
    bbox_tensors.append(bbox_tensor)
# 構建模型
model = tf.keras.Model(input_layer, bbox_tensors)
# 加載模型參數
model.load_weights("./yolov3")

測試

1、首先要從 .txt 文件中導入每一張圖片的真實框信息

with open(cfg.TEST.ANNOT_PATH, 'r') as annotation_file:

2、然後用 for 循環來遍歷每一個圖片中的信息

for num, line in enumerate(annotation_file):

假設此時是第 200 張圖片,那麼此時的 line 爲:

line = 'E:\\Pycharm\\code\\Jupyter\\tensorflow2.0\\My_net\\YOLO_v3\\data\\dataset\\test\\000200.jpg 141,305,197,361,7 87,23,143,79,0 46,114,102,170,1 16,55,44,83,5 31,253,87,309,5 257,305,341,389,6 317,163,401,247,4 164,118,248,202,0\n'

3、使用 strip() 去除首尾空格,用 split() 按照指定的符號將字符串 line 分割成一個列表中的不同元素

annotation = line.strip().split()
annotation = ['E:\\Pycharm\\code\\Jupyter\\tensorflow2.0\\My_net\\YOLO_v3\\data\\dataset\\test\\000200.jpg',
			  '141,305,197,361,7',
			  '87,23,143,79,0',
			  '46,114,102,170,1',
			  '16,55,44,83,5',
			  '31,253,87,309,5',
			  '257,305,341,389,6',
			  '317,163,401,247,4',
			  '164,118,248,202,0']

4、將圖片路徑單獨提取出來

image_path = annotation[0]
image_path = 'E:\\Pycharm\\code\\Jupyter\\tensorflow2.0\\My_net\\YOLO_v3\\data\\dataset\\test\\000200.jpg'

5、用 split() 提取圖片名稱

image_name = image_path.split('\\')[-1]

這裏可能有的同學是用 ‘/’ 分割,可以自己嘗試一下。

image_name = '000200.jpg'

6、讀取圖片

image = cv2.imread(image_path)  # BGR 圖像,像素在 0~255 之間
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 將 BGR 轉成 RGB
image_size = image.shape[:2]  # 這張圖片的尺寸

7、用 map 函數將真實框信息從字符串類型變爲數值類型,然後放到一個列表中

bbox_data_gt = np.array([list(map(int, box.split(','))) for box in annotation[1:]])
bbox_data_gt = array([[141, 305, 197, 361,   7],
				      [ 87,  23, 143,  79,   0],
				      [ 46, 114, 102, 170,   1],
				      [ 16,  55,  44,  83,   5],
				      [ 31, 253,  87, 309,   5],
				      [257, 305, 341, 389,   6],
				      [317, 163, 401, 247,   4],
				      [164, 118, 248, 202,   0]])

8、將 bbox_data_gt 中的框座標和分類分開

if len(bbox_data_gt) == 0:
	bboxes_gt = []
	classes_gt = []
else:
    bboxes_gt, classes_gt = bbox_data_gt[:, :4], bbox_data_gt[:, 4]
bboxes_gt = array([[141, 305, 197, 361],
			       [ 87,  23, 143,  79],
			       [ 46, 114, 102, 170],
			       [ 16,  55,  44,  83],
			       [ 31, 253,  87, 309],
			       [257, 305, 341, 389],
			       [317, 163, 401, 247],
			       [164, 118, 248, 202]])
classes_gt = array([7, 0, 1, 5, 5, 6, 4, 0])

9、生成200份存儲真實框信息的文本路徑

ground_truth_path = os.path.join(ground_truth_dir_path, str(num) + '.txt')  # os.path.join拼接路徑名稱
ground_truth_path = './mAP/ground-truth\\199.txt'

10、統計這張圖片中真實框的數量

num_bbox_gt = len(bboxes_gt)
num_bbox_gt = 8

11、將圖片上真實標註框信息寫入剛得到的 .txt 文件中

with open(ground_truth_path, 'w') as f:
    for i in range(num_bbox_gt):
        class_name = CLASSES[classes_gt[i]]
        # 將列表中各個元素分別賦值給多個變量
        xmin, ymin, xmax, ymax = list(map(str, bboxes_gt[i]))
        # ' '.join() 將列表裏面的元素都合成一個,以空格符作爲間隔 '\n'表示換行
        bbox_mess = ' '.join([class_name, xmin, ymin, xmax, ymax]) + '\n'
        f.write(bbox_mess)
        print('\t' + str(bbox_mess).strip())  # '\t' 橫向製表符

此時第 200 張圖片對應的 199.txt 中的內容是:

7 141 305 197 361
0 87 23 143 79
1 46 114 102 170
5 16 55 44 83
5 31 253 87 309
6 257 305 341 389
4 317 163 401 247
0 164 118 248 202

12、生成200份存儲預測框信息的文本路徑

predict_result_path = os.path.join(predicted_dir_path, str(num) + '.txt')
predict_result_path = './mAP/predicted\\199.txt'

13、圖像預處理

  • 圖像比例縮放
  • 填充
  • 均一化
image_data = utils.image_preporcess(np.copy(image), [INPUT_SIZE, INPUT_SIZE])
image_data = image_data[np.newaxis, ...].astype(np.float32)

14、模型預測輸出

pred_bbox = model.predict(image_data)
pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]  # shape [3, -1, 85],3 是因爲有 3 個 feature map
pred_bbox = tf.concat(pred_bbox, axis=0)  # shape [-1, 85]

15、確定預測框的信息

  • 確定預測框在原始圖片位置
  • 確定超出邊界的預測框索引
  • 確定分數大於一定閾值的預測框索引
  • 確定概率最大所對應類別索引
  • 確定滿足三個索引的預測框
# 分數=置信值*分類最大概率值
bboxes = utils.postprocess_boxes(pred_bbox, image_size, INPUT_SIZE, cfg.TEST.SCORE_THRESHOLD)  #bboxes格式爲[預測框的數量,預測框位置+分數+類別] shape爲[-1,6]

16、極大值抑制(預測框冗餘處理)

  • 找出擁有該類別最大分數的預測框
  • 存儲該預測框
  • 計算該預測框與其他預測框的 IOU
  • 根據 IOU 相關條件刪除預測框
  • 剩下的預測框繼續執行上述四個步驟,直至沒有預測框
bboxes = utils.nms(bboxes, cfg.TEST.IOU_THRESHOLD, method='nms')

17、繪製預測框

if cfg.TEST.DECTECTED_IMAGE_PATH is not None:
	image = utils.draw_bbox(image, bboxes)
	cv2.imwrite(cfg.TEST.DECTECTED_IMAGE_PATH+image_name, image)

18、將預測結果寫在存儲預測框信息的文本中

with open(predict_result_path, 'w') as f:
    for bbox in bboxes:
        coor = np.array(bbox[:4], dtype=np.int32)
        score = bbox[4]
        class_ind = int(bbox[5])
        class_name = CLASSES[class_ind]
        score = '%.4f' % score
        xmin, ymin, xmax, ymax = list(map(str, coor))
        bbox_mess = ' '.join([class_name, score, xmin, ymin, xmax, ymax]) + '\n'
        f.write(bbox_mess)
        print('\t' + str(bbox_mess).strip())

完整代碼

import cv2
import os
import shutil
import numpy as np
import tensorflow as tf
import core.utils as utils
from core.config import cfg
from core.yolov3 import YOLOv3, decode


INPUT_SIZE   = 416
NUM_CLASS    = len(utils.read_class_names(cfg.YOLO.CLASSES))
CLASSES      = utils.read_class_names(cfg.YOLO.CLASSES)

predicted_dir_path = './mAP/predicted'
ground_truth_dir_path = './mAP/ground-truth'
if os.path.exists(predicted_dir_path): shutil.rmtree(predicted_dir_path)
if os.path.exists(ground_truth_dir_path): shutil.rmtree(ground_truth_dir_path)
if os.path.exists(cfg.TEST.DECTECTED_IMAGE_PATH): shutil.rmtree(cfg.TEST.DECTECTED_IMAGE_PATH)

os.mkdir(predicted_dir_path)
os.mkdir(ground_truth_dir_path)
os.mkdir(cfg.TEST.DECTECTED_IMAGE_PATH)

# Build Model
input_layer  = tf.keras.layers.Input([INPUT_SIZE, INPUT_SIZE, 3])
feature_maps = YOLOv3(input_layer)

bbox_tensors = []
for i, fm in enumerate(feature_maps):
    bbox_tensor = decode(fm, i)
    bbox_tensors.append(bbox_tensor)

model = tf.keras.Model(input_layer, bbox_tensors)
model.load_weights("./yolov3")

with open(cfg.TEST.ANNOT_PATH, 'r') as annotation_file:
    for num, line in enumerate(annotation_file):
        annotation = line.strip().split()
        image_path = annotation[0]
        image_name = image_path.split('/')[-1]
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        bbox_data_gt = np.array([list(map(int, box.split(','))) for box in annotation[1:]])

        if len(bbox_data_gt) == 0:
            bboxes_gt=[]
            classes_gt=[]
        else:
            bboxes_gt, classes_gt = bbox_data_gt[:, :4], bbox_data_gt[:, 4]
        ground_truth_path = os.path.join(ground_truth_dir_path, str(num) + '.txt')

        print('=> ground truth of %s:' % image_name)
        num_bbox_gt = len(bboxes_gt)
        with open(ground_truth_path, 'w') as f:
            for i in range(num_bbox_gt):
                class_name = CLASSES[classes_gt[i]]
                xmin, ymin, xmax, ymax = list(map(str, bboxes_gt[i]))
                bbox_mess = ' '.join([class_name, xmin, ymin, xmax, ymax]) + '\n'
                f.write(bbox_mess)
                print('\t' + str(bbox_mess).strip())
        print('=> predict result of %s:' % image_name)
        predict_result_path = os.path.join(predicted_dir_path, str(num) + '.txt')
        # Predict Process
        image_size = image.shape[:2]
        image_data = utils.image_preporcess(np.copy(image), [INPUT_SIZE, INPUT_SIZE])
        image_data = image_data[np.newaxis, ...].astype(np.float32)

        pred_bbox = model.predict(image_data)
        pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
        pred_bbox = tf.concat(pred_bbox, axis=0)
        bboxes = utils.postprocess_boxes(pred_bbox, image_size, INPUT_SIZE, cfg.TEST.SCORE_THRESHOLD)
        bboxes = utils.nms(bboxes, cfg.TEST.IOU_THRESHOLD, method='nms')


        if cfg.TEST.DECTECTED_IMAGE_PATH is not None:
            image = utils.draw_bbox(image, bboxes)
            cv2.imwrite(cfg.TEST.DECTECTED_IMAGE_PATH+image_name, image)

        with open(predict_result_path, 'w') as f:
            for bbox in bboxes:
                coor = np.array(bbox[:4], dtype=np.int32)
                score = bbox[4]
                class_ind = int(bbox[5])
                class_name = CLASSES[class_ind]
                score = '%.4f' % score
                xmin, ymin, xmax, ymax = list(map(str, coor))
                bbox_mess = ' '.join([class_name, score, xmin, ymin, xmax, ymax]) + '\n'
                f.write(bbox_mess)
                print('\t' + str(bbox_mess).strip())
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章