# 图像分类

图像分类数据集用于训练模型以对整个图像进行分类。这些数据集支持广泛的应用，例如识别濒危野生动物物种或在医学图像中筛查疾病。本指南将展示如何对图像分类数据集应用变换。

在开始之前，请确保已安装最新版本的 `albumentations` 和 `cv2`：

```bash
pip install -U albumentations opencv-python
```

本指南使用 [Beans](https://huggingface.co/datasets/beans) 数据集，根据叶片图像识别豆类植物病害的类型。

加载数据集并查看一个示例：

```py
>>> from datasets import load_dataset

>>> dataset = load_dataset("AI-Lab-Makerere/beans")
>>> dataset["train"][10]
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=500x500 at 0x7F8D2F4D7A10>,
 'image_file_path': '/root/.cache/huggingface/datasets/downloads/extracted/b0a21163f78769a2cf11f58dfc767fb458fc7cea5c05dccc0144a2c0f0bc1292/train/angular_leaf_spot/angular_leaf_spot_train.204.jpg',
 'labels': 0}
```

数据集包含三个字段：

* `image`：一个 PIL 图像对象。
* `image_file_path`：图像文件的路径。
* `labels`：图像的标签或类别。

接下来，查看一张图像：

    

现在使用 `albumentations` 应用一些数据增强。你将随机裁剪图像、水平翻转它并调整其亮度。

```py
>>> import cv2
>>> import albumentations
>>> import numpy as np

>>> transform = albumentations.Compose([
...     albumentations.RandomCrop(width=256, height=256),
...     albumentations.HorizontalFlip(p=0.5),
...     albumentations.RandomBrightnessContrast(p=0.2),
... ])
```

创建一个函数，将转换应用于图像：

```py
>>> def transforms(examples):
...     examples["pixel_values"] = [
...         transform(image=np.array(image))["image"] for image in examples["image"]
...     ]
... 
...     return examples
```

使用 [set_transform()](/docs/datasets/v5.2.0/zh/package_reference/main_classes#datasets.Dataset.set_transform) 函数对数据集的批次进行即时转换，以减少磁盘空间占用：

```py
>>> dataset.set_transform(transforms)
```

你可以通过索引第一个示例的 `pixel_values` 来验证转换是否成功：

```py
>>> import numpy as np
>>> import matplotlib.pyplot as plt

>>> img = dataset["train"][0]["pixel_values"]
>>> plt.imshow(img)
```

    
    

> [!TIP]
> 现在你已经了解了如何为图像分类处理数据集，接下来学习[如何训练图像分类模型](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/image_classification.ipynb)并使用它进行推理。