LV-MHP-v2 数据集解析

本代码用于可视化 LV-MHP-v2 数据集。

注意： 该数据集的标注文件的标签是最后一个通道，并且是一个 rgb 的三通道图片，所有直接以灰度图像的格式读取，会导致标签错误。

import cv2
import numpy as np
import os

# 读取正常图像
color_image = cv2.imread("val/images/7.jpg")

# 读取RGB图像
image = cv2.imread('val/parsing_annos/7_02_01.png', cv2.IMREAD_COLOR)
# 提取最后一个通道
label_image = image[:, :, 2]

# 获取类别数量
num_classes = np.max(label_image) + 1
print(num_classes)

# 类别名称与颜色映射
label_colors_dict = {
	0: (255, 255, 255), # Background - 白色
	1: (255, 0, 0), # Cap/hat - 红色
	2: (0, 255, 0), # Helmet - 绿色
	3: (0, 0, 255), # Face - 蓝色
	4: (255, 255, 0), # Hair - 黄色
	5: (255, 0, 255), # Left-arm - 紫色
	6: (0, 255, 255), # Right-arm - 青色
	7: (255, 0, 128), # Left-hand - 粉色
	8: (128, 0, 255), # Right-hand - 深紫色
	9: (0, 255, 128), # Protector - 青绿色
	10: (255, 128, 0), # Bikini/bra - 橙色
	11: (128, 255, 0), # Jacket/windbreaker/hoodie - 亮绿色
	12: (0, 128, 255), # Tee-shirt - 深蓝色
	13: (255, 0, 0), # Polo-shirt - 红色
	14: (0, 255, 0), # Sweater - 绿色
	15: (0, 0, 255), # Singlet - 蓝色
	16: (255, 255, 0), # Torso-skin - 黄色
	17: (255, 0, 255), # Pants - 紫色
	18: (0, 255, 255), # Shorts/swim-shorts - 青色
	19: (255, 0, 128), # Skirt - 粉色
	20: (128, 0, 255), # Stockings - 深紫色
	21: (0, 255, 128), # Socks - 青绿色
	22: (255, 128, 0), # Left-boot - 橙色
	23: (128, 255, 0), # Right-boot - 亮绿色
	24: (0, 128, 255), # Left-shoe - 深蓝色
	25: (255, 0, 0), # Right-shoe - 红色
	26: (0, 255, 0), # Left-highheel - 绿色
	27: (0, 0, 255), # Right-highheel - 蓝色
	28: (255, 255, 0), # Left-sandal - 黄色
	29: (255, 0, 255), # Right-sandal - 紫色
	30: (0, 255, 255), # Left-leg - 青色
	31: (255, 0, 128), # Right-leg - 粉色
	32: (128, 0, 255), # Left-foot - 深紫色
	33: (0, 255, 128), # Right-foot - 青绿色
	34: (255, 0, 0), # Coat - 红色
	35: (0, 255, 0), # Dress - 绿色
	36: (0, 0, 255), # Robe - 蓝色
	37: (255, 255, 0), # Jumpsuit - 黄色
	38: (255, 0, 255), # Other-full-body-clothes - 紫色
	39: (0, 255, 255), # Headwear - 青色
	40: (255, 0, 128), # Backpack - 粉色
	41: (128, 0, 255), # Ball - 深紫色
	42: (0, 255, 128), # Bats - 青绿色
	43: (255, 128, 0), # Belt - 橙色
	44: (0, 128, 255), # Bottle - 深蓝色
	45: (255, 0, 0), # Carrybag - 红色
	46: (0, 255, 0), # Cases - 绿色
	47: (0, 0, 255), # Sunglasses - 蓝色
	48: (255, 255, 0), # Eyewear - 黄色
	49: (255, 0, 255), # Glove - 紫色
	50: (0, 255, 255), # Scarf - 青色
	51: (255, 128, 0), # Umbrella - 橙色
	52: (128, 255, 0), # Wallet/purse - 亮绿色
	53: (0, 128, 255), # Watch - 深蓝色
	54: (255, 0, 0), # Wristband - 红色
	55: (0, 255, 0), # Tie - 绿色
	56: (0, 0, 255), # Other-accessory - 蓝色
	57: (255, 255, 0), # Other-upper-body-clothes - 黄色
	58: (255, 0, 255), # Other-lower-body-clothes - 紫色
	59: (128, 128, 128), # Mask - 灰色
}

# 创建彩色标注图像
label_image_color = cv2.cvtColor(label_image, cv2.COLOR_GRAY2BGR)

# Get unique class IDs and their counts
class_ids, class_counts = np.unique(label_image, return_counts=True)

# 将每个类别的颜色叠加在彩色标注图像上
for class_id in range(1, num_classes):
	class_pixels = np.where(label_image == class_id)
	label_image_color[class_pixels] = label_colors_dict[class_id]

# 将标注图像叠加在正常图像上
blended_image = cv2.addWeighted(color_image, 0.7, label_image_color, 0.6, 0)

# 显示或保存叠加后的图像
cv2.imshow("Blended Image", blended_image)
cv2.waitKey(0)
# cv2.imwrite("path/to/save/blended_image.png", blended_image)
cv2.destroyAllWindows()

切出来目标区域用于训练分类模型代码：

# -*- encoding: utf-8 -*-
'''
@File    :   parsing.py
@Time    :   2023/07/05 17:53:23
@Author  :   
@Version :   1.0
@Contact :   
@License :   
@Desc    :   把分割的区域切出来用于分类图像。
'''

# here put the import lib
import cv2
import numpy as np
import os
from tqdm import tqdm
import multiprocessing
from functools import partial

def create_directory(path):
	if not os.path.exists(path):  # 判断目录是否存在
		os.makedirs(path)  # 递归创建目录
		print("Directory created:", path)

def save_class_regions(label_path, image_dir, class_ids, out_dir):
	# 读取图像和标签图像
	# 读取RGB图像
	image = cv2.imread(label_path, cv2.IMREAD_COLOR)
	# 提取最后一个通道
	label_image = image[:, :, 2]
	label_filename = os.path.basename(label_path)
	# 构建对应的 JPG 文件路径
	jpg_filename = label_filename.split('_')[0]  + '.jpg'
	image_path = os.path.join(image_dir, jpg_filename)
	image = cv2.imread(image_path)

	# 对每个指定的类别进行处理
	for class_id in class_ids:
		if class_id > np.max(label_image) or np.sum(label_image == class_id) == 0:
			continue

		# 创建一个与原图像大小相同的全黑图像
		class_region_seg = np.zeros_like(image)

		# 将类别区域的像素值设置为原图像对应位置的像素值
		class_region_seg[label_image == class_id] = image[label_image == class_id]

		# 创建一个二值图像，将目标类别设置为白色，其他类别设置为黑色
		class_mask = np.where(label_image == class_id, 255, 0).astype(np.uint8)

		# 使用类别掩码获得类别区域的轮廓
		contours, _ = cv2.findContours(class_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

		# 计算最大边界框的初始值
		x_min = y_min = float('inf')
		x_max = y_max = 0

		# 遍历每个轮廓
		for contour in contours:
			# 找到当前轮廓的边界框
			x, y, w, h = cv2.boundingRect(contour)

			# 更新最大边界框的坐标
			x_min = min(x_min, x)
			y_min = min(y_min, y)
			x_max = max(x_max, x + w)
			y_max = max(y_max, y + h)

		# 切割原始图像中的区域
		class_region = image[y_min:y_max, x_min:x_max]

		# 获取 seg 区域
		seg_class_region = class_region_seg[y_min:y_max, x_min:x_max]

		# 生成保存图像的文件名
		out_detect_dir = os.path.join(out_dir, 'detect', class_ids[class_id])
		create_directory(out_detect_dir)
		detect_path = os.path.join(out_detect_dir, jpg_filename)
		# 保存类别区域图像
		cv2.imwrite(detect_path, class_region)

		out_detect_dir = os.path.join(out_dir, 'seg', class_ids[class_id])
		create_directory(out_detect_dir)
		seg_path = os.path.join(out_detect_dir, jpg_filename)

		# # 保存 seg 类别区域图像
		cv2.imwrite(seg_path, seg_class_region)

def get_selected_classes():
	selected_classes = {
		1: 'Cap_hat',
		4: 'Hair',
		11: 'Jacket_windbreaker_hoodie',
		12: 'Tee-shirt',
		13: 'Polo-shirt',
		14: 'Sweater',
		15: 'Singlet',
		19: 'Skirt',
		34: 'Coat',
		35: 'Dress',
		37: 'Jumpsuit',
		38: 'Other-full-body-clothes',
		47: 'Sunglasses',
		57: 'Other-upper-body-clothes'
	}
	return selected_classes

if __name__ == "__main__":
	selected_classes = get_selected_classes()
	skirt_class_id = selected_classes[19]
	print(skirt_class_id)  # 输出为 Skirt
	image_dir = 'val/images'
	label_dir = 'val/parsing_annos'
	out_dir = 'val_cls'

	label_files = []
	for entry in os.scandir(label_dir):
		if entry.is_file() and entry.name.endswith('.png'):
			label_files.append(entry.path)

	print("数据读取完成！")

	with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool: # 使用CPU核心数作为进程数
		total_files = len(label_files)
		process_func = partial(save_class_regions, image_dir=image_dir, class_ids=selected_classes, out_dir=out_dir)
		with tqdm(total=total_files, desc='Processing files') as pbar:
			for _ in pool.imap_unordered(process_func, label_files):
				pbar.update(1)