[Python_ts3] cv2 모듈 이미지 띄우고 키 값 리스트에 저장하기(cv2.imread, cv2.waitKey)

안녕하세요~ 이번에는 cv2 모듈을 통해 이미지를 띄우고, 사용자가 누르는 키 값을 리스트에 저장하는 task를 해보려고 해요. 

os모듈과 반복문을 통해 라벨링된 초음파 이미지 데이터를 창에 띄우면서, 라벨링이 제대로 된 데이터(이미지)인지 확인하기 위한 자동화 툴을 구현해볼거예요. 

이미지를 자동으로 창에 띄우면서 제가 'yes'나 'no'를 누르며 이 이미지의 라벨링이 잘 되었는지 저장해주는 코드인거죠..! 빈 리스트에  'yes'나 'no'값을 저장해나갈거랍니다.😊

 

데이터는 아래와 같이 생겼는데요,

왼쪽은 라벨링이 잘 된 상태, 오른쪽은 라벨링이 잘 되지 않은 상태입니다.

 

os모듈에 대한 포스팅은 이전 포스팅에서 진행했었는데요, 사용법을 참고하고 싶으시면 아래 글을 확인해주세요 :)

https://risingdata.tistory.com/63

[Python_ts1] os.path.split과 os.path.spiltext 알기

https://risingdata.tistory.com/64

[Python_ts2] 여러 폴더 내의 파일 정보 가져오기 (os.walk)

 

---

우선 필요한 모듈을 불러와 준 후, 파일 경로 및 빈 리스트를 만들어줍니다.

os모듈은 상하위 폴더를 전체적으로 읽어들일 때 필요한 모듈이고, cv2모듈은 이미지를 창에 띄울 때 필요한 모듈입니다.

저희는 labeller 리스트, 인덱스 리스트, 'y', 'n' 또는 'NaN'을 담을 리스트를 빈 리스트로 만들어 줄 것입니다.

반복문을 돌면서 해당 정보들이 리스트 안에 담기게 될거예요!

그리고 폴더 내에 이미지 파일이 없는 경우에는 'NaN'값을 입력할 것입니다.

import os
import cv2
import pandas as pd

# 파일 경로
filepath = 'T:/car/Second_upload'

labeller_list = []  # labeller 리스트
index_list = []     # 인덱스 리스트
label_flag = []     # 'y', 'n' 또는 'NaN'을 담을 리스트

 

다음으로 os.walk를 통해 상하위 폴더를 전체적으로 돌건데요,

저희가 띄워야 하는 이미지는 'T:/car/Second_upload\label 1\CA_1000122070\Label_Data' 경로 안에 있는 이미지(bmp파일)를 띄워줄거예요.

 

'for path, _, file in os.walk(filepath):'를 통해 반복문을 돌리면 path는 아래와 같이 나온답니다.

# path = T:/car/Second_upload\label 1\CA_1000122070
#        T:/car/Second_upload\label 1\CA_1000122070\InI
#        T:/car/Second_upload\label 1\CA_1000122070\Label_Data

if문을 통해 path 안에 'Label_Data'가 있는 경우라고 조건문을 걸어서 해당 경로로 들어갈 수 있게 해줍니다.

해당 경로 안에 이미지(bmp파일)이 없는 경우에는 label_flag변수에 'NaN'값을 저장해줄게요. ➡️ append

그리고 labeller와 index 정보도 같이 들어가야 하므로, 이 정보도 넣어줄게요!

for path, _, file in os.walk(filepath):  # _와 file이 있는 path, path 아래에 있는 폴더들, path 아래에 있는 파일들
    if 'Label_Data' in path:
        
      # bmp파일이 없는 경우
      if len(file) == 0:  # .bmp파일이 없는 경우는 경로에 들어가지 않기 때문에 list 길이가 0이다.
                          # 만약, img_path = path + '/{}'.format(file[0])를 할 경우에는 file[0]이 없기 때문에 에러가 뜬다.
                          
         ret = 'NaN'      # label_flag에 'NaN'입력
         label_flag.append(ret)
         print("label doesn't exist, {0}".format(ret))        
         
         # labeller와 index 정보는 그대로 넣어주어야 한다.
         labeller = path.split('\\')[-3]
         labeller_list.append(labeller)
         
         index = path.split('\\')[-2]
         index_list.append(index)

 

bmp 파일이 있는 경우, 해당 경로에 들어가서 cv2모듈을 통해 이미지를 창에 띄워줍니다. ( cv2.imread, cv2.imshow)

그리고 이 창이 사용자가 key를 입력하기까지의 대기 시간을 걸어줍니다. cv2.waitKey(0)이면 무한 대기가 됩니다. 여기에서 key는 'y', 'n'입니다.

 

입력한 key값을 label_flag 변수에 저장해줍니다.

      # bmp 파일이 있는 경우
      else:
         # img_path = T:/car/Second_upload\label 1\CA_1000122070\Label_Data/000.BMP ...
         img_path = path + '/{}'.format(file[0])  # bmp파일이 있으니, file[0]까지 해주어 경로에 들어간다.
         
         # 이미지 띄우기
         img = cv2.imread(img_path)
         cv2.imshow('Label image', img)
         
         # 키 입력 대기 시간 -> 무한 대기
         ret = cv2.waitKey(0)
         
         # labeller와 index 정보
         labeller = img_path.split('\\')[-3]
         labeller_list.append(labeller)

         index = img_path.split('\\')[-2]
         index_list.append(index)
         
         # 입력하는 키 정보('y' or 'n') 넣어주기
         label_flag.append(ret)
         print('pressed key is {0}'.format(ret))

 

그리고 반복문을 빠져나오기 위해 아래와 같은 조건문도 작성해줍니다.

         # esc or x버튼 누르면 반복문 빠져나오기
         if ret == 27 or ret == -1:
               break

 

마지막으로 데이터 프레임을 생성한 뒤, csv파일을 생성해줍니다.

# 데이터 프레임 생성
df = pd.DataFrame({
    'labeller': labeller_list,
    'index': index_list,
    'flag': label_flag
})

# csv파일 생성
df.to_csv('result.csv', index=False)

print("END")

 

 

📌 결과

 

flag변수의 121과 110은 'y'와 'n'의 ASCII 값입니다. 해당 이미지의 labeller, index정보와 라벨링 정상 여부를 판단할 수 있는 flag변수까지 제대로 저장된 것을 확인할 수 있습니다. 

 

 

---

지금까지 긴 글 읽어주셔서 감사합니다!