基于PyTorch的服装颜色识别系统设计与优化

李昦

1. 项目背景与核心价值

去年帮学弟调试这个毕设项目时，我意识到服装颜色识别这个看似简单的任务，在实际落地时会遇到各种预料之外的挑战。传统计算机视觉方法在应对不同光照条件、拍摄角度和面料材质时表现极不稳定，而基于CNN的解决方案在实验室环境下准确率能轻松突破95%，这背后的技术细节值得深入探讨。

这个项目本质上是要构建一个能自动识别T恤颜色的分类系统，核心难点在于：

颜色在RGB空间中的非线性分布特性
不同材质对光线的反射差异
现实场景中的光照干扰问题
相似色系（如深蓝与藏青）的细微区分

采用PyTorch框架的优势在于其动态计算图和丰富的预训练模型库，特别适合在校生快速实现原型开发。我在实际测试中发现，即使使用相同的网络结构，不同的数据预处理方式会导致最终准确率产生15%以上的波动。

2. 技术方案设计要点

2.1 数据采集与标注规范

原始数据集的质量直接决定模型上限。我们采用自建数据集+公开数据集结合的方式：

自拍300张不同光照下的纯色T恤照片（建议使用手机固定支架）
从DeepFashion数据集中筛选500张相关图像
标注时采用潘通色卡作为颜色基准，将颜色归类为12个主色系

关键细节：所有照片必须包含标准色卡参照物，后期可通过白平衡校正消除设备色差。标注时遇到渐变款T恤时，以主视觉面积超过60%的颜色为准。

2.2 网络架构选型对比

测试了三种经典CNN结构的性能表现（batch_size=32, epoch=50）：

模型	参数量	Top-1准确率	训练时间
ResNet18	11M	89.2%	2.1h
MobileNetV3	5.4M	86.7%	1.3h
自定义CNN	0.8M	82.4%	0.8h

最终选择轻量化的自定义CNN结构，在保证实时性的同时满足毕业设计精度要求。网络包含：

4个卷积块（Conv2D+BN+ReLU+MaxPool）
2个全连接层（含Dropout=0.3）
输出层使用Softmax激活

python复制class ColorCNN(nn.Module):
    def __init__(self, num_classes=12):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(2),
            # ...后续层省略
        )
        self.classifier = nn.Sequential(
            nn.Linear(16*28*28, 128),
            nn.Dropout(0.3),
            nn.Linear(128, num_classes)
        )
    
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)

2.3 颜色空间转换技巧

RGB色彩空间对光照变化过于敏感，实践中发现转换到HSV空间后模型鲁棒性显著提升：

python复制def rgb_to_hsv(img_tensor):
    # R,G,B分量归一化到[0,1]
    r, g, b = img_tensor.unbind(dim=-3)
    max_val = torch.max(img_tensor, dim=-3).values
    min_val = torch.min(img_tensor, dim=-3).values
    
    # 计算H分量
    h = torch.zeros_like(max_val)
    diff = max_val - min_val
    mask = diff != 0
    h[mask & (max_val == r)] = (60 * ((g - b)/diff) + 360) % 360
    h[mask & (max_val == g)] = (60 * ((b - r)/diff) + 120) % 360
    h[mask & (max_val == b)] = (60 * ((r - g)/diff) + 240) % 360
    
    # 计算S分量
    s = torch.where(max_val == 0, torch.zeros_like(max_val), diff / max_val)
    
    return torch.stack([h/360.0, s, max_val], dim=-3)

3. 关键实现细节

3.1 数据增强策略

针对服装图像的独特性，采用组合增强方案：

几何变换：随机旋转（±15°）、水平翻转
颜色扰动：亮度调整（±20%）、饱和度变化（±30%）
模拟干扰：添加高斯噪声（σ=0.01）、随机遮挡

python复制train_transform = transforms.Compose([
    transforms.RandomRotation(15),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, saturation=0.3),
    transforms.ToTensor(),
    transforms.Lambda(rgb_to_hsv),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

3.2 损失函数优化

标准CrossEntropyLoss在颜色分类中会遇到类别不平衡问题，改进方案：

引入Focal Loss解决简单样本主导问题
加入Label Smoothing防止过拟合
自定义权重处理相似色混淆

python复制class WeightedFocalLoss(nn.Module):
    def __init__(self, alpha=None, gamma=2):
        super().__init__()
        # 相似色权重矩阵
        self.alpha = torch.tensor([...]) 
        self.gamma = gamma

    def forward(self, inputs, targets):
        BCE_loss = F.cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-BCE_loss)
        loss = (self.alpha[targets] * (1-pt)**self.gamma * BCE_loss).mean()
        return loss

3.3 模型部署优化

使用LibTorch将模型导出为C++可调用格式，实现10ms级单图推理速度：

python复制# 导出TorchScript
model.eval()
example = torch.rand(1, 3, 224, 224)
traced_script = torch.jit.trace(model, example)
traced_script.save("color_classifier.pt")

# C++端调用示例
torch::jit::script::Module module = torch::jit::load("color_classifier.pt");
at::Tensor output = module.forward({input_tensor}).toTensor();