目标检测脚本之统计数据集中某个类别某个尺寸的目标数量

   
# 计算label标签中各个尺寸的数量分布
import os
import shutil
import cv2
from tqdm import tqdm

def calculate(image_dir,label_dir):
    # 10以下尺寸数量
    small_10_count = 0      # 10 以下尺寸数量
    # 10-32 尺寸数量
    small_10_32_count = 0    #  10-32 尺寸数量
    # 大于 32尺寸数量
    other_object_count = 0   # 大于 32尺寸数量

    # 获取标签文件列表
    label_file_list = os.listdir(label_dir)
    # 遍历标签文件列表
    for label_file in tqdm(label_file_list):
        # 拼接图片路径
        image_path = os.path.join(image_dir,label_file.replace("txt","jpg"))
        # 拼接标签路径
        label_path = os.path.join(label_dir,label_file)
        # 如果图片路径不存在，则跳过当前循环
        if not os.path.exists(image_path):
            continue
        # 读取图片
        image = cv2.imread(image_path)
        # 获取图片的宽和高
        image_w,image_h = image.shape[1],image.shape[0]

        with open(label_path,"r",encoding='utf-8') as f:
            # 读取标签文件内容，按行存储到labels列表中
            labels = f.readlines()
            # 遍历标签列表
            for label in labels:
                # 将标签内容按空格分割成items列表
                items = label.split(" ")
                # 如果items列表的第一个元素为'0'，则执行以下判断和计数操作
                if items[0] == '0':
                    # 如果图片的宽小于10或高小于10，则small_10_count加1
                    if (float(items[2]) * image_w < 10) or (float(items[3]) * image_h < 10):
                        small_10_count +=1
                    # 如果图片的宽小于32或高小于32，则small_10_32_count加1
                    elif (float(items[2]) * image_w < 32.0) or (float(items[3]) * image_h < 32.0):
                        small_10_32_count +=1
                    # 其他情况则other_object_count加1
                    else:
                        other_object_count +=1
            
    return small_10_count,small_10_32_count,other_object_count                            


if __name__ == "__main__":
    image_dir = '' # 图片地址
    label_dir = '' # 标签地址
    
   
    small_10_counts,small_10_32_count,other_counts = calculate(image_dir,label_dir)
    
    print(small_10_counts,small_10_32_count,other_counts)