CUB-200-2011 数据集实战：3步完成 PyTorch 数据加载与 15 个关键点可视化
在细粒度图像分类领域，CUB-200-2011 数据集因其丰富的标注信息和挑战性而备受研究者青睐。本文将带你从零开始，快速掌握该数据集的核心使用方法，重点解决三个工程难题：如何高效加载多模态标注数据、如何处理复杂的文件映射关系，以及如何实现关键点的精准可视化。
1. 数据准备与环境搭建
1.1 数据集结构解析
下载并解压 CUB-200-2011 数据集后，你会看到如下目录结构：
CUB_200_2011/ ├── images/ # 200个子目录，每个对应一种鸟类 ├── parts/ # 关键点标注 │ ├── parts.txt │ ├── part_locs.txt │ └── part_click_locs.txt ├── attributes/ # 属性标注 ├── images.txt # 图像ID与路径映射 ├── bounding_boxes.txt # 边界框坐标 ├── train_test_split.txt # 官方划分 └── image_class_labels.txt # 图像类别标签关键文件说明：
| 文件 | 每行格式 | 描述 |
|---|---|---|
| images.txt | <image_id> <image_name> | 图像ID到路径的映射 |
| bounding_boxes.txt | <image_id> x y w h | 边界框坐标（左上x,y + 宽高） |
| part_locs.txt | <image_id> <part_id> x y visible> | 关键点坐标及可见性 |
1.2 安装必要依赖
推荐使用 conda 创建虚拟环境：
conda create -n cub python=3.8 conda activate cub pip install torch torchvision matplotlib pandas2. 构建 PyTorch 数据管道
2.1 数据集类设计
我们创建一个继承自torch.utils.data.Dataset的类，完整处理所有标注信息：
import os import torch from PIL import Image import pandas as pd class CUBDataset(torch.utils.data.Dataset): def __init__(self, root, transform=None, train=True): self.root = root self.transform = transform # 加载所有映射文件 self._load_metadata() # 根据train_test_split筛选数据 self._filter_split(train) def _load_metadata(self): # 加载图像路径映射 images_df = pd.read_csv(os.path.join(self.root, 'images.txt'), sep=' ', names=['img_id', 'img_path']) # 加载边界框 bbox_df = pd.read_csv(os.path.join(self.root, 'bounding_boxes.txt'), sep=' ', names=['img_id', 'x', 'y', 'w', 'h']) # 加载类别标签 class_df = pd.read_csv(os.path.join(self.root, 'image_class_labels.txt'), sep=' ', names=['img_id', 'class_id']) # 加载训练测试划分 split_df = pd.read_csv(os.path.join(self.root, 'train_test_split.txt'), sep=' ', names=['img_id', 'is_train']) # 加载关键点 parts_df = pd.read_csv(os.path.join(self.root, 'parts/part_locs.txt'), sep=' ', names=['img_id', 'part_id', 'x', 'y', 'visible']) # 合并所有信息 self.metadata = images_df.merge(bbox_df, on='img_id') \ .merge(class_df, on='img_id') \ .merge(split_df, on='img_id') # 处理关键点：将长格式转为宽格式（每张图一行，包含所有关键点） self.parts = parts_df.pivot(index='img_id', columns='part_id', values=['x', 'y', 'visible']) def _filter_split(self, train): split_flag = 1 if train else 0 self.metadata = self.metadata[self.metadata['is_train'] == split_flag] def __len__(self): return len(self.metadata) def __getitem__(self, idx): row = self.metadata.iloc[idx] img_id = row['img_id'] # 加载图像 img_path = os.path.join(self.root, 'images', row['img_path']) img = Image.open(img_path).convert('RGB') # 获取边界框并裁剪 bbox = (row['x'], row['y'], row['x']+row['w'], row['y']+row['h']) img = img.crop(bbox) # 获取关键点（相对边界框的坐标） parts = self.parts.loc[img_id] keypoints = [] for part_id in range(1, 16): # 共15个关键点 x = parts['x'][part_id] - row['x'] y = parts['y'][part_id] - row['y'] visible = parts['visible'][part_id] keypoints.append([x, y, visible]) # 转换为tensor keypoints = torch.tensor(keypoints, dtype=torch.float32) class_id = torch.tensor(row['class_id'] - 1, dtype=torch.long) # 转为0-based # 应用变换 if self.transform: img = self.transform(img) return img, keypoints, class_id2.2 数据增强策略
针对细粒度分类任务，推荐使用以下变换组合：
from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.8, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])2.3 创建数据加载器
dataset_train = CUBDataset('path/to/CUB_200_2011', transform=train_transform, train=True) dataset_val = CUBDataset('path/to/CUB_200_2011', transform=val_transform, train=False) train_loader = torch.utils.data.DataLoader( dataset_train, batch_size=32, shuffle=True, num_workers=4) val_loader = torch.utils.data.DataLoader( dataset_val, batch_size=32, shuffle=False, num_workers=4)3. 关键点可视化实现
3.1 关键点信息解析
CUB-200-2011 定义了15个鸟类关键点：
- 喙尖 (beak tip)
- 喙基部 (beak base)
- 左眼 (left eye)
- 右眼 (right eye)
- 头顶 (crown)
- 颈部 (nape)
- 左翼 (left wing)
- 右翼 (right wing)
- 尾部 (tail)
- 左脚 (left foot)
- 右脚 (right foot)
- 胸部 (breast)
- 背部 (back)
- 左腿 (left leg)
- 右腿 (right leg)
3.2 可视化代码实现
import matplotlib.pyplot as plt import numpy as np def visualize_sample(dataset, index): img, keypoints, class_id = dataset[index] # 反归一化图像 img = img.numpy().transpose((1, 2, 0)) mean = np.array([0.485, 0.456, 0.406]) std = np.array([0.229, 0.224, 0.225]) img = std * img + mean img = np.clip(img, 0, 1) # 创建绘图 plt.figure(figsize=(10, 10)) plt.imshow(img) # 绘制关键点 colors = plt.cm.hsv(np.linspace(0, 1, 15)) for i, (x, y, visible) in enumerate(keypoints): if visible > 0.5: # 只绘制可见点 plt.scatter(x, y, c=[colors[i]], s=100, label=f'Part {i+1}', edgecolors='white') # 添加图例 plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left') plt.title(f'Class: {class_id.item()+1}') plt.axis('off') plt.tight_layout() plt.show() # 示例：可视化训练集第一个样本 visualize_sample(dataset_train, 0)3.3 批量可视化技巧
对于研究分析，我们常需要批量查看关键点分布：
def plot_keypoints_distribution(dataset, num_samples=9): plt.figure(figsize=(15, 15)) indices = np.random.choice(len(dataset), num_samples, replace=False) for i, idx in enumerate(indices): img, keypoints, _ = dataset[idx] img = img.numpy().transpose((1, 2, 0)) img = std * img + mean img = np.clip(img, 0, 1) plt.subplot(3, 3, i+1) plt.imshow(img) # 绘制所有可见关键点 for x, y, visible in keypoints: if visible > 0.5: plt.scatter(x, y, c='red', s=10) plt.axis('off') plt.tight_layout() plt.suptitle('Random Samples with Keypoints', y=1.02) plt.show() plot_keypoints_distribution(dataset_train)4. 高级应用与性能优化
4.1 数据加载加速技巧
当处理大规模数据时，可以采用以下优化策略：
- 预加载关键点数据：将关键点信息预先转换为numpy数组存储
- 使用内存映射文件：对于大型标注文件，使用
np.memmap - 并行加载：设置
num_workers为CPU核心数
优化后的数据集类初始化：
def _load_metadata_optimized(self): # 使用更高效的文件读取方式 self.images = np.loadtxt(os.path.join(self.root, 'images.txt'), dtype=str) self.bboxes = np.loadtxt(os.path.join(self.root, 'bounding_boxes.txt')) # 预加载关键点到内存 parts = np.loadtxt(os.path.join(self.root, 'parts/part_locs.txt')) self.keypoints = np.zeros((len(self.images), 15, 3)) # (N, 15, 3) for img_id, part_id, x, y, visible in parts: self.keypoints[int(img_id)-1, int(part_id)-1] = [x, y, visible]4.2 自定义数据采样策略
针对类别不平衡问题，实现加权随机采样：
from torch.utils.data import WeightedRandomSampler # 计算每个类别的样本数 class_counts = np.bincount([dataset_train[i][2] for i in range(len(dataset_train))]) class_weights = 1. / class_counts sample_weights = class_weights[[dataset_train[i][2] for i in range(len(dataset_train))]] sampler = WeightedRandomSampler(sample_weights, len(sample_weights)) balanced_loader = DataLoader(dataset_train, batch_size=32, sampler=sampler)4.3 多任务学习框架
同时利用关键点和类别信息进行多任务学习：
import torch.nn as nn class MultiTaskModel(nn.Module): def __init__(self, num_classes=200): super().__init__() # 共享的特征提取器 self.backbone = torchvision.models.resnet50(pretrained=True) in_features = self.backbone.fc.in_features self.backbone.fc = nn.Identity() # 移除原始全连接层 # 分类头 self.classifier = nn.Linear(in_features, num_classes) # 关键点回归头 self.keypoint_regressor = nn.Sequential( nn.Linear(in_features, 512), nn.ReLU(), nn.Linear(512, 15*2) # 预测15个点的(x,y)坐标 ) def forward(self, x): features = self.backbone(x) # 分类输出 class_logits = self.classifier(features) # 关键点输出 (batch_size, 15, 2) keypoints = self.keypoint_regressor(features).view(-1, 15, 2) return class_logits, keypoints训练时需要定义复合损失函数：
def multitask_loss(class_logits, keypoints_pred, targets): class_target, keypoints_target = targets # 分类损失 cls_loss = F.cross_entropy(class_logits, class_target) # 关键点损失（只计算visible=1的点） visible = keypoints_target[:, :, 2] > 0.5 kp_loss = F.mse_loss( keypoints_pred[visible], keypoints_target[:, :, :2][visible] ) return cls_loss + 0.1 * kp_loss # 加权求和
