【目标检测】统计YOLO实例分割标注Polygon多边形像素&BBox

原创已于 2024-04-18 14:27:32 修改 · 1.1k 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#目标检测 #YOLO #计算机视觉 #深度学习 #人工智能

于 2023-06-13 13:17:14 首次发布

该代码片段用于统计YOLO实例分割格式的数据集中，各个类别的目标框数量以及对应的多边形像素总数。它遍历指定目录下的标注文本文件，读取每一行的内容，解析出类别标签和坐标信息，然后计算像素数。最后，输出各类别的目标框计数和像素计数。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

统计YOLO实例分割格式各类别多边形像素数和目标框数量：

from glob import glob
import re
import cv2
from PIL import Image, ImageDraw
from tqdm import tqdm
import numpy as np
import os

# 使用glob模式匹配，避免目录切分问题
all_files = glob(r'YOLO/merge/labels/*.txt')
image_path = 'YOLO/merge/images'

bbox_sta = {str(i): 0 for i in range(10)}
mask_sta = {str(i): 0 for i in range(10)}

for file in tqdm(all_files, desc='Processing', unit='item'):
    file_name = file.split('\\')[-1].split('.')[0]
    img_path = f'{image_path}/{file_name}.jpg'
    
    # 判断图像文件是否存在
    if not os.path.exists(img_path):
        continue

    img = cv2.imread(img_path)
    height, width, _ = img.shape

    with open(file, 'r') as f:
        for l in f.readlines():
            line = l.split()
            class_label = line[0]
            bbox_sta[class_label] += 1

            # 解析坐标点
            points = [(int(float(line[i]) * width), int(float(line[i+1]) * height)) for i in range(1, len(line), 2)]

            # 创建掩模并计算像素数
            mask_image = Image.new("L", (width, height), 0)
            ImageDraw.Draw(mask_image).polygon(points, outline=1, fill=1)
            mask_array = np.array(mask_image)
            mask_sta[class_label] += int(np.sum(mask_array))

print("------------------")
print("Class BBox Counts:")
for class_label, count in bbox_sta.items():
    print(f"Class {class_label}: {count}")
print("------------------")
print("Class Pixel Counts:")
for class_label, count in mask_sta.items():
    print(f"Class {class_label}: {count}")