Project - [ 유해 요소 객체 탐지 : 데이터 수집 ]

date

Jul 21, 2023

slug

harmful-object-detection-01

author

status

Public

유해 요소 데이터 수집

이미지 수집: https://www.pexels.com/

기준 및 방향을 정하기 위해 먼저 샘플로 총, 칼, 담배, 손가락 욕, 술, 피 20장씩만 수집했다.

Label 및 BBox 정보 저장 방법

파일 다운로드: https://github.com/heartexlabs/labelImg/releases

YOLO 방식 선택 후

Create\nRectBox 로 물체 bbox 지정

Save


{'gun':0, 'knife':1, 'cigarette':2, 'bad fingers':3, 'alcohol':4, 'blood':5}

위에서 alcohol image 를 예시로 사용했으므로 label 값은 4가 된다.

이미지파일명.txt 에 label 및 bbox 정보 저장

수정 사항

원래 predefined_classes 의 default 값은 다음과 같다.

그런데 다음과 같이 dir 을 구성하고 label 을 추가하게 되면


class
├── alcohol
├── bad fingers
├── blood
├── cigarette
├── gun
├── knife

폴더 별로 해당 class 만 predefined_classes 에 추가된다. (classes: dog, person, cat, … hamburger, cavity, alcohol) (classes: dog, person, cat, … hamburger, cavity, bad fingers) (classes: dog, person, cat, … hamburger, cavity, blood) . . . (classes: dog, person, cat, … hamburger, cavity, knife)

그런데 이렇게 되면 나중에 데이터를 합쳐 train 시킬 때, 모든 class 들의 label 값이 15로 같게 되는 문제가 발생한다.

따라서 처음부터 predefined_classes 에 사용할 class 들을 다 넣어주고 시작하는 것이 편하다.

하지만 이미 labeling 및 bboxing 을 한 후에 알게 되었다고 해도, 처음부터 다시 진행하거나 txt 파일마다 일일이 수정할 필요는 없다.

다음과 같이 한 번에 처리할 수 있다.

먼저 label mapping 을 해준다.


# label mapping
labels = ['gun', 'knife', 'cigarette', 'bad fingers', 'alcohol', 'blood']
label_dic = dict()
for i in range(len(labels)):
    label_dic[labels[i]] = i

파일 하나에 대해 수행해 본다.


# 파일 하나에 대해 수행해 보기
file_name = r"C:\Users\mook\Downloads\windows_v1.8.1\data\class\alcohol\pexels-ainis-jankauskas-169391.txt"
file_label = file_name.split('\\')[-2] # file_label = alcohol

with open(file_name, 'r') as file:
    lines = file.readlines() # bbox 가 여러 개인 것들은 여러 줄이 작성되므로 line 단위로 읽기

updated_lines = []
for line in lines:
    line_list = list(line.split())
    line_list[0] = str(label_dic[file_label]) # label: 15 → 4 로 수정
    updated_line = ' '.join(line_list) + '\n' # 수정한 부분과 나머지 다시 합치기
    updated_lines.append(updated_line) # 최종 수정된 line 업데이트

with open(file_name, 'w') as file: # 처음부터 'w' 로 하지 않는 이유는 'w' 로 file 을 열게 되면, 비어있는 txt 가 열리기 때문
    file.writelines(updated_lines) # 최종 수정된 line 작성

os.walk() 에 대해서 알아본다.


# os.walk 함수 살펴보기
import os
root = r"C:\Users\mook\Downloads\windows_v1.8.1\data\class"

for root_dir, sub_dir_names, file_names in os.walk(root): # os.walk(root) : (해당 경로, 하위 디렉토리, 파일 목록) 튜플 형태로 반환
    print(
        f'Root-dir: {root_dir}\n'
        f'Sub-dirs: {sub_dir_names}\n'
        f'Files: {file_names}\n\n'
    )

(참고: https://www.youtube.com/watch?v=DzzXB9XhBg4)

폴더를 순회하면서 모든 txt 파일에 대해서 수행한다. (전체 코드)


import os

labels = ['gun', 'knife', 'cigarette', 'bad fingers', 'alcohol', 'blood']
label_dic = dict()
for i in range(len(labels)):
    label_dic[labels[i]] = i

# 폴더를 순회하면서 모든 txt 파일에 대해서 수행
root = r"C:\Users\mook\Downloads\windows_v1.8.1\data\class"

for root_dir, sub_dir_names, file_names in os.walk(root):
    for file_name in file_names:
        if file_name.endswith('.txt'):
            file_label = root_dir.split('\\')[-1]
            file_path = os.path.join(root_dir, file_name)
            
            with open(file_path, 'r') as file:
                lines = file.readlines()
            updated_lines = []
            for line in lines:
                line_list = list(line.split())
                line_list[0] = str(label_dic[file_label])
                updated_line = ' '.join(line_list) + '\n'
                updated_lines.append(updated_line)
            
            with open(file_path, 'w') as file:
                file.writelines(updated_lines)

위의 코드를 실행하면 각 폴더마다 해당 class 에 mapping 한 숫자로 label 값이 수정된다.