#!/usr/bin/env python # -*- encoding: utf-8 -*- import os from typing import List import cv2 import torch import numpy as np from PIL import Image import torchvision.transforms as transforms from app.models.schp import networks from app.models.schp.utils.transforms import get_affine_transform, transform_logits # 数据集设置 dataset_settings = { 'lip': { 'input_size': [473, 473], 'num_classes': 20, 'label': ['Background', 'Hat', 'Hair', 'Glove', 'Sunglasses', 'Upper-clothes', 'Dress', 'Coat', 'Socks', 'Pants', 'Jumpsuits', 'Scarf', 'Skirt', 'Face', 'Left-arm', 'Right-arm', 'Left-leg', 'Right-leg', 'Left-shoe', 'Right-shoe'] }, 'atr': { 'input_size': [512, 512], 'num_classes': 18, 'label': ['Background', 'Hat', 'Hair', 'Sunglasses', 'Upper-clothes', 'Skirt', 'Pants', 'Dress', 'Belt', 'Left-shoe', 'Right-shoe', 'Face', 'Left-leg', 'Right-leg', 'Left-arm', 'Right-arm', 'Bag', 'Scarf'] }, 'pascal': { 'input_size': [512, 512], 'num_classes': 7, 'label': ['Background', 'Head', 'Torso', 'Upper Arms', 'Lower Arms', 'Upper Legs', 'Lower Legs'], } } def get_color_by_label(label: str) -> List[int]: """ 根据标签名称获取对应的RGB颜色值 Args: label (str): 标签名称，如 'Face', 'Hair' 等 Returns: List[int]: RGB颜色值列表，格式为 [R, G, B]，值范围0-255 """ # LIP数据集标签 labels = ['Background', 'Hat', 'Hair', 'Glove', 'Sunglasses', 'Upper-clothes', 'Dress', 'Coat', 'Socks', 'Pants', 'Jumpsuits', 'Scarf', 'Skirt', 'Face', 'Left-arm', 'Right-arm', 'Left-leg', 'Right-leg', 'Left-shoe', 'Right-shoe'] # 检查标签是否存在 if label not in labels: return [] # 获取标签索引 label_index = labels.index(label) # 生成调色板 palette = get_palette(len(labels)) # 获取对应颜色的RGB值并返回列表 r = palette[label_index * 3 + 0] g = palette[label_index * 3 + 1] b = palette[label_index * 3 + 2] return [r, g, b] def get_palette(num_cls): """返回用于可视化分割掩码的颜色映射""" n = num_cls palette = [0] * (n * 3) for j in range(0, n): lab = j palette[j * 3 + 0] = 0 palette[j * 3 + 1] = 0 palette[j * 3 + 2] = 0 i = 0 while lab: palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i)) palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i)) palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i)) i += 1 lab >>= 3 return palette def _box2cs(box, aspect_ratio): """将边界框转换为中心点和尺度""" x, y, w, h = box[:4] return _xywh2cs(x, y, w, h, aspect_ratio) def _xywh2cs(x, y, w, h, aspect_ratio): """将xywh格式转换为中心点和尺度""" center = np.zeros((2), dtype=np.float32) center[0] = x + w * 0.5 center[1] = y + h * 0.5 if w > aspect_ratio * h: h = w * 1.0 / aspect_ratio elif w < aspect_ratio * h: w = h * aspect_ratio scale = np.array([w, h], dtype=np.float32) return center, scale class HumanParsingModel: def __init__(self, model_path, dataset='atr', device=None): """ 初始化人体解析模型 Args: model_path: 预训练模型路径 dataset: 数据集类型 ('lip', 'atr', 'pascal') device: 计算设备 (None表示自动选择) """ self.dataset = dataset self.device = device if device else torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 获取数据集设置 self.num_classes = dataset_settings[dataset]['num_classes'] self.input_size = dataset_settings[dataset]['input_size'] self.label = dataset_settings[dataset]['label'] self.aspect_ratio = self.input_size[1] * 1.0 / self.input_size[0] self.input_size_array = np.asarray(self.input_size) # 初始化模型 self.model = networks.init_model('resnet101', num_classes=self.num_classes, pretrained=None) # 加载预训练权重 state_dict = torch.load(model_path, map_location=self.device)['state_dict'] from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] if k.startswith('module.') else k # 移除 'module.' 前缀 new_state_dict[name] = v self.model.load_state_dict(new_state_dict) # 将模型移动到指定设备并设置为评估模式 self.model.to(self.device) self.model.eval() # 图像预处理变换 self.transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229]) ]) # 获取调色板 self.palette = get_palette(self.num_classes) print(f"模型已加载，使用设备: {self.device}") print(f"数据集: {dataset}, 类别数: {self.num_classes}") def process_single_image(model, input_image): """ 处理单张图片 Args: model: HumanParsingModel实例 input_image: 输入图片，可以是： - numpy数组 (H, W, C) BGR格式 - PIL Image对象 - 图片文件路径字符串 Returns: PIL Image对象，包含分割结果的彩色图像 """ # 处理不同类型的输入 if isinstance(input_image, str): # 如果是文件路径 img = cv2.imread(input_image, cv2.IMREAD_COLOR) if img is None: raise ValueError(f"无法读取图片: {input_image}") elif isinstance(input_image, Image.Image): # 如果是PIL Image，转换为BGR numpy数组 img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR) elif isinstance(input_image, np.ndarray): # 如果是numpy数组，直接使用 img = input_image.copy() else: raise ValueError("输入图片格式不支持，请使用numpy数组、PIL Image或文件路径") h, w, _ = img.shape # 获取人体中心点和尺度 person_center, s = _box2cs([0, 0, w - 1, h - 1], model.aspect_ratio) r = 0 # 获取仿射变换矩阵 trans = get_affine_transform(person_center, s, r, model.input_size_array) # 应用仿射变换 input_tensor = cv2.warpAffine( img, trans, (int(model.input_size[1]), int(model.input_size[0])), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0) ) # 预处理 input_tensor = model.transform(input_tensor) input_tensor = input_tensor.unsqueeze(0) # 添加batch维度 input_tensor = input_tensor.to(model.device) # 模型推理 with torch.no_grad(): output = model.model(input_tensor) # 上采样到输入尺寸 upsample = torch.nn.Upsample(size=model.input_size, mode='bilinear', align_corners=True) upsample_output = upsample(output[0][-1][0].unsqueeze(0)) upsample_output = upsample_output.squeeze() upsample_output = upsample_output.permute(1, 2, 0) # CHW -> HWC # 变换回原始图像尺寸 logits_result = transform_logits( upsample_output.data.cpu().numpy(), person_center, s, w, h, input_size=model.input_size ) # 获取分割结果 parsing_result = np.argmax(logits_result, axis=2) # 转换为PIL图像并应用调色板 output_img = Image.fromarray(np.asarray(parsing_result, dtype=np.uint8)) output_img.putpalette(model.palette) return output_img # 使用示例函数 def parse_human_image(input_image, model_path, dataset='lip', device=None): """ 便捷函数：解析单张人体图像 Args: input_image: 输入图片 (numpy数组、PIL Image或文件路径) model_path: 预训练模型路径 dataset: 数据集类型 ('lip', 'atr', 'pascal') device: 计算设备 Returns: PIL Image对象，包含分割结果 """ # 创建模型实例 model = HumanParsingModel(model_path, dataset, device) # 处理图像 result = process_single_image(model, input_image) return result # 使用示例 if __name__ == '__main__': # 示例1：使用便捷函数 model_path = r"D:\work\PycharmProjects\PythonProject\checkpoints\exp-schp-201908261155-lip.pth" input_image_path = r"D:\work\PycharmProjects\PythonProject\img1.jpg" # result_image = parse_human_image(input_image_path, model_path, dataset='atr') # result_image.save(r"D:\work\PycharmProjects\PythonProject\output_result.png") # 示例2：使用类的方式（推荐用于批量处理） model = HumanParsingModel(model_path, dataset='lip') # 处理多张图片 image_paths = [r"D:\work\PycharmProjects\PythonProject\img1.jpg", r"D:\work\PycharmProjects\PythonProject\img2.jpg", r"D:\work\PycharmProjects\PythonProject\img3.jpg"] for i, img_path in enumerate(image_paths): result = process_single_image(model, img_path) result.save(f"result_{i}.png")