统计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}")