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.txtimage_id image_name图像ID到路径的映射bounding_boxes.txtimage_id x y w h边界框坐标左上x,y 宽高part_locs.txtimage_id part_id x y visible关键点坐标及可见性1.2 安装必要依赖推荐使用 conda 创建虚拟环境conda create -n cub python3.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, transformNone, trainTrue): 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, onimg_id) \ .merge(class_df, onimg_id) \ .merge(split_df, onimg_id) # 处理关键点将长格式转为宽格式每张图一行包含所有关键点 self.parts parts_df.pivot(indeximg_id, columnspart_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, dtypetorch.float32) class_id torch.tensor(row[class_id] - 1, dtypetorch.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(brightness0.2, contrast0.2, saturation0.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, transformtrain_transform, trainTrue) dataset_val CUBDataset(path/to/CUB_200_2011, transformval_transform, trainFalse) train_loader torch.utils.data.DataLoader( dataset_train, batch_size32, shuffleTrue, num_workers4) val_loader torch.utils.data.DataLoader( dataset_val, batch_size32, shuffleFalse, num_workers4)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]], s100, labelfPart {i1}, edgecolorswhite) # 添加图例 plt.legend(bbox_to_anchor(1.05, 1), locupper left) plt.title(fClass: {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_samples9): plt.figure(figsize(15, 15)) indices np.random.choice(len(dataset), num_samples, replaceFalse) 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, i1) plt.imshow(img) # 绘制所有可见关键点 for x, y, visible in keypoints: if visible 0.5: plt.scatter(x, y, cred, s10) plt.axis(off) plt.tight_layout() plt.suptitle(Random Samples with Keypoints, y1.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), dtypestr) 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_size32, samplersampler)4.3 多任务学习框架同时利用关键点和类别信息进行多任务学习import torch.nn as nn class MultiTaskModel(nn.Module): def __init__(self, num_classes200): super().__init__() # 共享的特征提取器 self.backbone torchvision.models.resnet50(pretrainedTrue) 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) # 关键点损失只计算visible1的点 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 # 加权求和