Datasets documentation

의미 분할

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

의미 분할

Semantic segmentation 데이터셋은 이미지 내 모든 픽셀을 분류하도록 모델을 학습하는 데 사용됩니다. 이러한 데이터셋을 활용하면 이미지 배경 제거, 이미지 스타일링, 자율주행용 장면 이해 등 다양한 애플리케이션을 구현할 수 있습니다. 이 가이드에서는 이미지 세그멘테이션 데이터셋에 변환을 적용하는 방법을 보여줍니다.

시작하기 전에 albumentations와 cv2의 최신 버전이 설치되어 있는지 확인하세요:

pip install -U albumentations opencv-python

Albumentations는 컴퓨터 비전을 위한 데이터 증강을 수행하는 Python 라이브러리입니다. 이미지 분류, 객체 감지, 분할, 키포인트 추정과 같은 다양한 컴퓨터 비전 작업을 지원합니다.

이 가이드는 Scene Parsing 데이터셋을 사용하여 이미지를 하늘, 도로, 사람, 침대와 같은 의미론적 범주와 연관된 서로 다른 이미지 영역으로 분할하고 분석하는 데 사용합니다.

데이터셋의 train 분할을 로드하고 예시를 확인해 보세요:

>>> from datasets import load_dataset

>>> dataset = load_dataset("scene_parse_150", split="train")
>>> index = 10
>>> dataset[index]
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=683x512 at 0x7FB37B0EC810>,
 'annotation': <PIL.PngImagePlugin.PngImageFile image mode=L size=683x512 at 0x7FB37B0EC9D0>,
 'scene_category': 927}

데이터셋에는 세 가지 필드가 있습니다:

  • image: PIL 이미지 객체.
  • annotation: 이미지의 분할 마스크
  • scene_category: 이미지의 라벨 또는 장면 카테고리(예: “kitchen” 또는 “office”).

다음으로, 다음 명령어로 이미지를 확인해 보세요:

>>> dataset[index]["image"]

마찬가지로, 해당 분할 마스크를 확인할 수 있습니다:

>>> dataset[index]["annotation"]

분할 마스크에 색상 팔레트를 추가하고 원본 이미지 위에 겹쳐서 데이터셋을 시각화할 수도 있습니다:

색상 팔레트를 정의한 후, 오버레이를 시각화할 준비가 되었습니다.

>>> import matplotlib.pyplot as plt

>>> def visualize_seg_mask(image: np.ndarray, mask: np.ndarray):
...    color_seg = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
...    palette = np.array(create_ade20k_label_colormap())
...    for label, color in enumerate(palette):
...        color_seg[mask == label, :] = color
...    color_seg = color_seg[..., ::-1]  # convert to BGR

...    img = np.array(image) * 0.5 + color_seg * 0.5  # plot the image with the segmentation map
...    img = img.astype(np.uint8)

...    plt.figure(figsize=(15, 10))
...    plt.imshow(img)
...    plt.axis("off")
...    plt.show()


>>> visualize_seg_mask(
...     np.array(dataset[index]["image"]),
...     np.array(dataset[index]["annotation"])
... )

albumentations로 몇 가지 증강을 적용해 보겠습니다. 먼저 이미지를 리사이즈하고 밝기를 조정합니다.

>>> import albumentations

>>> transform = albumentations.Compose(
...     [
...         albumentations.Resize(256, 256),
...         albumentations.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.5),
...     ]
... )

이미지에 변환을 적용하는 함수를 생성합니다:

>>> def transforms(examples):
...     transformed_images, transformed_masks = [], []
...
...     for image, seg_mask in zip(examples["image"], examples["annotation"]):
...         image, seg_mask = np.array(image), np.array(seg_mask)
...         transformed = transform(image=image, mask=seg_mask)
...         transformed_images.append(transformed["image"])
...         transformed_masks.append(transformed["mask"])
...
...     examples["pixel_values"] = transformed_images
...     examples["label"] = transformed_masks
...     return examples

set_transform() 함수를 사용하여 데이터셋의 배치에 변환을 온더플라이로 적용하여 디스크 공간을 덜 사용하도록 할 수 있습니다:

>>> dataset.set_transform(transforms)

pixel_values와 label에 인덱싱하여 예시를 확인함으로써 변환이 제대로 작동했는지 검증할 수 있습니다:

>>> image = np.array(dataset[index]["pixel_values"])
>>> mask = np.array(dataset[index]["label"])

>>> visualize_seg_mask(image, mask)

이 가이드에서는 albumentations를 사용하여 데이터셋을 증강했습니다. torchvision을 사용하여 유사한 변환을 적용하는 것도 가능합니다.

>>> from torchvision.transforms import Resize, ColorJitter, Compose

>>> transformation_chain = Compose([
...     Resize((256, 256)),
...     ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1)
... ])
>>> resize = Resize((256, 256))

>>> def train_transforms(example_batch):
...     example_batch["pixel_values"] = [transformation_chain(x) for x in example_batch["image"]]
...     example_batch["label"] = [resize(x) for x in example_batch["annotation"]]
...     return example_batch

>>> dataset.set_transform(train_transforms)

>>> image = np.array(dataset[index]["pixel_values"])
>>> mask = np.array(dataset[index]["label"])

>>> visualize_seg_mask(image, mask)

이제 의미 분할을 위한 데이터셋을 처리하는 방법을 알게 되셨으니, 의미 분할 모델을 학습하는 방법과 추론에 이를 활용하는 방법을 배워보겠습니다.

Update on GitHub