자동차 데이터셋 전처리하기
이전에 했던 커스텀 학습에서 문제점을 보완하기 위해, 차량데이터를 새로 찾아보았다.
자동차 데이터셋 Link
http://ai.stanford.edu/~jkrause/cars/car_dataset.html
해당 데이터셋에는 16185장의 차량데이터를 가지고 있고, 196 종의 차량 클래스로 분류되어있다.
라벨형식은 .mat 파일로 matlab을 사용해야만 볼 수 있었다. 다행이 같은 연구실의 대학원생분의 컴퓨터에 matlab을 사용할 수 있어서 라벨형식을 확인할 수 있었다.
다음과 같이 이미지의 경로, x1min, y1min, x2max, y2max, 클래스정보, test 유무 순으로 annotation이 되어있었다. 다행이 ctrl+a 그리고 복사붙여넣기를 통해 txt 파일에 옮길수 있었다.
클래스는 차량의 모델명으로 되어있었지만, 내가 원하는 정보는 아니기 때문에 모두 'car'라는 클래스로 통일하였다.
- 전처리하는 코드
import os
from PIL import Image
output_dir = 'output/'
def convert(size, box):
dw = 1./size[0]
dh = 1./size[1]
x = (box[0] + box[1])/2.0
y = (box[2] + box[3])/2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = round(x*dw,6)
w = round(w*dw,6)
y = round(y*dh,6)
h = round(h*dh,6)
return (x,y,w,h)
f_input = open('label.txt','r')
items=f_input.readlines()
if not os.path.exists(output_dir):
os.makedirs(output_dir)
for item in items:
data=item.split()
#이미지 경로 가져오기
img_path=data[0]
image = Image.open(img_path[1:-1])
width = int(image.size[0])
height = int(image.size[1])
#이미지 이름 가져오기
img=img_path.split('/')
file_name=img[1][:-5]
#바운딩박스 위치 가져오기
xmin = data[1]
xmax = data[3]
ymin = data[2]
ymax = data[4]
#바운딩박스 위치 변환하기
b = (float(xmin), float(xmax), float(ymin), float(ymax))
bb = convert((width, height), b)
#변환한거 저장하기
f = open(output_dir + file_name + '.txt', 'w')
f.write('3' + " " + " ".join([str(a) for a in bb]) + '\n') 내가 학습하는 것은 yolov5이기 때문에 yolo형식 라벨링으로 바꿔주어야 한다. 따라서 해당이미지의 크기를 가져와서 해당 크기에 맞는 라벨 비율을 나타내주어야 yolo형식이 된다. 맨 마지막줄의 python f.write('3' + " " + " ".join([str(a) for a in bb]) + '\n')에서 '3'은 나중에 학습에 사용될 클래스 인덱스이다. 만약 다른 형식의 변환을 원하면 해당 코드의 convert 메소드만 가져다가 쓰면 될 것 같다. 해당 메소드는 이미지크기와 라벨 정보만 주면 yolo형식으로 변환한 후 return한다.
- 변환 결과
총 16185장의 이미지에대해 각각의 라벨을 얻었다.
바로 학습을 시작할려다가 보행자 데이터셋도 추가하여 학습하면 좋을 것 같다는 생각이 들었다. 보행자데이터셋을 얻고 전처리하는 과정은 추가로 글을 작성하겠다.
