影像分類

影像分類資料集用於訓練模型對整個影像進行分類。這些資料集可用於各種應用，例如識別瀕危野生物種或在醫學影像中篩查疾病。本指南將向您展示如何對影像分類資料集應用轉換。

在開始之前，請確保您已安裝最新版本的 albumentations 和 cv2

pip install -U albumentations opencv-python

本指南使用 Beans 資料集，根據豆類植物葉子的影像來識別其疾病型別。

載入資料集並檢視一個示例

>>> 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 應用一些增強。您將隨機裁剪影像、水平翻轉並調整其亮度。

>>> 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),
... ])

建立一個函式來對影像應用轉換

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

使用 set_transform() 函式即時將轉換應用於資料集的批次，以減少磁碟空間佔用

>>> dataset.set_transform(transforms)

您可以透過索引第一個示例的 pixel_values 來驗證轉換是否成功

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

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

既然您知道如何處理影像分類資料集，接下來可以學習如何訓練影像分類模型並將其用於推理。

< > 在 GitHub 上更新