개요

Python 언어를 사용하여 딕셔너리를 yaml 파일로 저장하는 과정에서 본인이 딕셔너리에 넣은 아이템들의 순서가 그대로 반영되지 않아 삽질하는 과정을 겪었다. 자세하게는 인공지능 데이터 셋을 COCO 형식에서 YOLO 형식으로 변환하는 과정에서 겪은 일이다. 이를 해결해나가는 과정을 기술한다.

환경

  • Ubuntu 20.04
  • python 3.8.15
  • PyYAML 6.0

문제

일반적인 딕셔너리 저장

코드는 다음과 같다.

import yaml


meta_data = dict()
meta_data['path'] = './'
meta_data['train'] = './train/images'
meta_data['val'] = './val/images'

meta_data['nc'] = 2
meta_data['names'] = ['cls1', 'cls2']

print(meta_data)

with open('./result.yaml', 'w', encoding='utf8') as f:
    yaml.dump(meta_data, f)

이 코드를 통해 의도한 바는 result.yamlpath, train, val, nc, names 순서대로 내용들이 적히길 바라는 것이다.

그러나 result.yaml 파일을 열어보면 다음과 같았다.

names:
- cls1
- cls2
nc: 2
path: ./
train: ./train/images
val: ./val/images

본인이 의도한 순서대로 적히지 않는 것을 볼 수 있다.

OrderedDict 사용

python 3.7 버전 이후로는 딕셔너리의 정렬이 자동으로 되는 것으로 알고 있었지만 원인이 딕셔너리에 있다고 생각해서 collections 라이브러리의 OrderedDict 모듈을 사용해보았다.

코드는 다음과 같다.

import yaml
from collections import OrderedDict


meta_data = OrderedDict()
meta_data['path'] = './'
meta_data['train'] = './train/images'
meta_data['val'] = './val/images'

meta_data['nc'] = 2
meta_data['names'] = ['cls1', 'cls2']

print(meta_data)

with open('./result.yaml', 'w', encoding='utf8') as f:
    yaml.dump(meta_data, f)

print 문의 결과로 출력되는 딕셔너리의 결과가 아래와 같이 달랐지만 result.yaml 의 내용물은 똑같이 정렬되어있지 않았다.

OrderedDict([('path', './'), ('train', './train/images'), ('val', './val/images'), ('nc', 2), ('names', ['cls1', 'cls2'])])

yaml dump 옵션 추가(해결)

찾아보니 yaml.dump() 메소드의 파라미터 중 다음과 같은 항목이 있었다.

sort_keys: bool = …) -> None

다음과 같이 수정하여 실행해보았다.

with open('./result.yaml', 'w', encoding='utf8') as f:
    yaml.dump(meta_data, f, sort_keys=False)

OrderedDictdump 시 제대로 되지 않는것 같다.

!!python/object/apply:collections.OrderedDict
- - - path
    - ./
  - - train
    - ./train/images
  - - val
    - ./val/images
  - - nc
    - 2
  - - names
    - - cls1
      - cls2

OrderedDict를 일반 딕셔너리로 변경한 최종 코드는 다음과 같다.

import yaml


meta_data = dict()
meta_data['path'] = './'
meta_data['train'] = './train/images'
meta_data['val'] = './val/images'

meta_data['nc'] = 2
meta_data['names'] = ['cls1', 'cls2']

with open('./result.yaml', 'w', encoding='utf8') as f:
    yaml.dump(meta_data, f, sort_keys=False)

결과로 생성된 result.yaml 파일이 제대로 출력되는 모습을 확인할 수 있다.

path: ./
train: ./train/images
val: ./val/images
nc: 2
names:
- cls1
- cls2

참고링크: stackoverflow - pyyaml dict dump ordered

댓글남기기