YOLOv5小目標切圖檢測的思路與方法

2022-12-21 14:00:30

前言

當我們在檢測較大解析度的圖片時，對小目標的檢測效果一直是較差的，所以就有了下面幾種方法：

將圖片壓縮成大尺寸進行訓練（想法：沒視訊記憶體，搞不來）
新增小檢測頭（想法：P5模型還有點用，P6模型完全沒用）
新增一些檢測模型和玄學機制（想法：你要是寫論文就去看看知*吧，只需要在最後面加一句：已達到工業檢測要求）
切圖檢測（想法：比較耗時，過程也比較繁瑣，可以嘗試）

切圖檢測

思路：

將原圖切成你想要的數量
將切成的小圖進行訓練，得到模型
將你需要檢測的圖片切成小圖，用模型檢測，並得到每張圖目標位置的資訊，儲存在對應圖片的txt檔案
將所有txt檔案融合，得到1個txt檔案，並在原圖上顯示

一：切塊

# -*- coding:utf-8 -*-
import os
import matplotlib.pyplot as plt
import cv2
import numpy as np
 
 
def divide_img(img_path, img_name, save_path):
    imgg = img_path + img_name
    img = cv2.imread(imgg)
    #   img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    h = img.shape[0]
    w = img.shape[1]
    n = int(np.floor(h * 1.0 / 1000)) + 1
    m = int(np.floor(w * 1.0 / 1000)) + 1
    print('h={},w={},n={},m={}'.format(h, w, n, m))
    dis_h = int(np.floor(h / n))
    dis_w = int(np.floor(w / m))
    num = 0
    for i in range(n):
        for j in range(m):
            num += 1
            print('i,j={}{}'.format(i, j))
            sub = img[dis_h * i:dis_h * (i + 1), dis_w * j:dis_w * (j + 1), :]
            cv2.imwrite(save_path + '{}_{}.bmp'.format(name, num), sub)
 
 
if __name__ == '__main__':
 
    img_path = r'G:1/'
    save_path = r'G:3/'
    img_list = os.listdir(img_path)
    for name in img_list:
        divide_img(img_path, name, save_path)

使用模型檢測後得到：

二：融合txt檔案

import os
from cv2 import cv2
 
# 儲存所有圖片的寬高
# todo: img_info={'name': [w_h, child_w_h, mix_row_w_h， mix_col_w_h]}
img_info = {}
all_info = {}
 
 
# 初始化img_info
def init(big_images_path, mix_percent, rows, cols):
    image_names = os.listdir(big_images_path)
    for img_name in image_names:
        big_path = big_images_path + '\' + img_name
        # print(big_path)
        img = cv2.imread(big_path)
        size = img.shape[0:2]
        w = size[1]
        h = size[0]
        child_width = int(w) // cols
        child_height = int(h) // rows
 
        mix_row_width = int(child_width * mix_percent * 2)
        mix_row_height = child_height
 
        mix_col_width = child_width
        mix_col_height = int(child_height * mix_percent * 2)
        # 根據img儲存w和h
        img_info[img_name.split('.')[0]] = [w, h, child_width, child_height, mix_row_width, mix_row_height,
                                            mix_col_width, mix_col_height]
 
 
# 讀取所有檢測出來的 小圖片的label
def get_label_info(labels_path, mix_percent, rows, cols):
    labels = os.listdir(labels_path)
    for label in labels:
        # print(label)
        # todo: type: 0正常, 1row, 2col
        # 判斷該label屬於哪一張圖片
        cur_label_belong = label.split('_')[0]
        cur_big_width = img_info[cur_label_belong][0]
        cur_big_height = img_info[cur_label_belong][1]
        # 融合區域距離邊界的一小部分寬高
        cur_row_width_step = img_info[cur_label_belong][2] * (1 - mix_percent)
        cur_col_height_step = img_info[cur_label_belong][3] * (1 - mix_percent)
        # 檔名給予資料
        # child_type = []
        # child_num = []
        # label內容給予資料
        child_class_index = []
        child_x = []
        child_y = []
        child_width = []
        child_height = []
 
        type = -1
        num = -1
        class_index = -1
        x = 0.0
        y = 0.0
        width = 0.0
        height = 0.0
 
        # print(f'{label}')
        # 讀取所有需要的資料
        f = open(labels_path + '\' + label, 'r')
        lines = f.read()
        # print(lines)
        f.close()
        contents = lines.split('n')[:-1]
        # print(contents)
        for content in contents:
            content = content.split(' ')
            # print(content)
            class_index = int(content[0])
            x = float(content[1])
            y = float(content[2])
            width = float(content[3])
            height = float(content[4])
            pass
            # print(class_index, x, y, width, height)
            assert class_index != -1 or x != -1.0 or y != -1.0 or width != -1.0 or height != -1.0, 
                f'class_index:{class_index}, x:{x}, y:{y}, width:{width}, height:{height}'
            # 轉換成 資料 座標, 並根據不同的num進行處理
            num = label.split('_')[-1].split('.')[0]  # 圖片尾號 命名： xxxx_x.jpg  xxxx_mix_row_xx.jpg xxxx_mix_col_xx.jpg
            cur_img_width = 0
            cur_img_height = 0
            distance_x = 0
            distance_y = 0
            small_image_width = img_info[cur_label_belong][2]
            small_image_height = img_info[cur_label_belong][3]
            if label.find('mix_row') != -1:
                # type = 1.
                distance_x = int(num) % (cols-1)
                distance_y = int(num) // (rows-1)
                cur_img_width = img_info[cur_label_belong][4]
                cur_img_height = img_info[cur_label_belong][5]
                # row x 加上step
                x = x * cur_img_width + cur_row_width_step + distance_x * small_image_width
                y = y * cur_img_height + distance_y * cur_img_height
            elif label.find('mix_col') != -1:
                # type = 2
                distance_x = int(num) % cols
                distance_y = int(num) // rows
                cur_img_width = img_info[cur_label_belong][6]
                cur_img_height = img_info[cur_label_belong][7]
                # col y 加上step
                print(f'x:{x}, y:{y}, cur_img_width:{cur_img_width}, cur_img_height:{cur_img_height}')
                x = x * cur_img_width + distance_x * cur_img_width
                y = y * cur_img_height + cur_col_height_step + distance_y * small_image_height
                print(f'x:{x}, y:{y}, height:{cur_col_height_step}')
            else:
                # type = 0
                distance_x = int(num) % cols
                distance_y = int(num) // rows
                cur_img_width = img_info[cur_label_belong][2]
                cur_img_height = img_info[cur_label_belong][3]
                # 小圖片內， 無需加上 step
                x = x * cur_img_width + distance_x * cur_img_width
                y = y * cur_img_height + distance_y * cur_img_height
            assert cur_img_width != 0 or cur_img_height != 0 or distance_x != 0 or distance_y != 0, 
                f'cur_img_width:{cur_img_width}, cur_img_height:{cur_img_height}, distance_x:{distance_x}, distance_y:{distance_y}'
            assert x < cur_big_width and y < cur_big_height, f'{label}, {content}nw:{cur_big_width}, h:{cur_big_height}, x:{x}, y:{y}'
            width = width * cur_img_width
            height = height * cur_img_height
            assert x != 0.0 or y != 0.0 or width != 0.0 or height != 0.0, f'x:{x}, y:{y}, width:{width}, height:{height}'
            # child_type.append(type)
            # child_num.append(num)
            child_class_index.append(class_index)
            child_x.append(x)
            child_y.append(y)
            child_width.append(width)
            child_height.append(height)
        # todo: 所有資訊 根據 cur_label_belong 儲存在all_info中
        for index, x, y, width, height in zip(child_class_index, child_x, child_y, child_width, child_height):
            if cur_label_belong not in all_info:
                all_info[cur_label_belong] = [[index, x, y, width, height]]
            else:
                all_info[cur_label_belong].append([index, x, y, width, height])
        child_class_index.clear()
        child_x.clear()
        child_y.clear()
        child_width.clear()
        child_height.clear()
 
 
# print((all_info['0342']))
# todo: 轉成 yolo 格式， 儲存
def save_yolo_label(yolo_labels_path):
    for key in all_info:
        # img_path = r'G:Unitycode_projectother_projectdatajointbig_images' + '\' + key + '.JPG'
        # img = cv2.imread(img_path)
        yolo_label_path = yolo_labels_path + '\' + key + '.txt'
        cur_big_width = img_info[key][0]
        cur_big_height = img_info[key][1]
        content = ''
        i = 0
        for index, x, y, width, height in all_info[key]:
            # print(all_info[key][i])
            x = x / cur_big_width
            y = y / cur_big_height
            width = width / cur_big_width
            height = height / cur_big_height
            assert x < 1.0 and y < 1.0 and width < 1.0 and height < 1.0, f'{key} {i}n{all_info[key][i]}nx:{x}, y:{y}, width:{width}, height:{height}'
            content += f'{index} {x} {y} {width} {height}n'
            i += 1
        with open(yolo_label_path, 'w') as f:
            f.write(content)
 
 
def joint_main(big_images_path=r'G:3',
               labels_path=r'G:5',
               yolo_labels_path=r'G:6',
               mix_percent=0.2,
               rows=4,
               cols=4):
    print(f'融合圖片, 原圖片路徑：{big_images_path}n小圖檢測的txt結果路徑：{labels_path}n資料融合後txt結果路徑：{yolo_labels_path}')
    init(big_images_path, mix_percent, rows, cols)
    get_label_info(labels_path, mix_percent, rows, cols)
    save_yolo_label(yolo_labels_path)
 
joint_main()

三：原圖顯示

# -*- coding: utf-8 -*-
import os
from PIL import Image
from PIL import ImageDraw, ImageFont
from cv2 import cv2
 
 
def draw_images(images_dir, txt_dir, box_dir, font_type_path):
    font = ImageFont.truetype(font_type_path, 50)
    if not os.path.exists(box_dir):
        os.makedirs(box_dir)
    # num = 0
 
    # 設定顏色
    all_colors = ['red', 'green', 'yellow', 'blue', 'pink', 'black', 'skyblue', 'brown', 'orange', 'purple', 'gray',
                  'lightpink', 'gold', 'brown', 'black']
    colors = {}
 
    for file in os.listdir(txt_dir):
        print(file)
        image = os.path.splitext(file)[0].replace('xml', 'bmp') + '.bmp'
        # 轉換成cv2讀取，防止圖片載入錯誤
        img = cv2.imread(images_dir + '/' + image)
        TURN = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(TURN)
        # img.show()
 
        if img.mode == "P":
            img = img.convert('RGB')
 
        w, h = img.size
        tag_path = txt_dir + '/' + file
        with open(tag_path) as f:
            for line in f:
                line_parts = line.split(' ')
                # 根據不同的 label 儲存顏色
                if line_parts[0] not in colors.keys():
                    colors[line_parts[0]] = all_colors[len(colors.keys())]
                color = colors[line_parts[0]]
 
                draw = ImageDraw.Draw(img)
                x = (float(line_parts[1]) - 0.5 * float(line_parts[3])) * w
                y = (float(line_parts[2]) - 0.5 * float(line_parts[4])) * h
                xx = (float(line_parts[1]) + 0.5 * float(line_parts[3])) * w
                yy = (float(line_parts[2]) + 0.5 * float(line_parts[4])) * h
                draw.rectangle([x - 10, y - 10, xx, yy], fill=None, outline=color, width=5)
                # num += 1
            del draw
            img.save(box_dir + '/' + image)
        # print(file, num)
    # print(colors)
 
 
def draw_main(box_dir=r'G:5',
              txt_dir=r'G:6',
              image_source_dir=r'G:3'):
    font_type_path = 'C:/Windows/Fonts/simsun.ttc'
    print(f'標註框, 資料來源: {txt_dir}n 被標註圖片: {image_source_dir}n 結果儲存路徑: {box_dir}')
    draw_images(image_source_dir, txt_dir, box_dir, font_type_path)
 
 
draw_main()

效果對比：（上YOLOv5檢測，下YOLOv5+切圖檢測）