다양한 파일을 읽어와 이미지 분류하기

images 폴더와 labels 폴더에 동일한 이름의 파일이 존재

labels 폴더의 txt 파일 맨 앞에는 대응하는 이미지의 클래스를 기록

import os
import shutil

images_folder = './train/images'
labels_folder = './train/labels'
output_folder = './train/classes'

for label_file in os.listdir(labels_folder):
    if label_file.endswith('.txt'):
        # 동일한 이름의 이미지 파일 경로 생성
        image_name = label_file.replace('.txt', '.jpg')
        image_path = os.path.join(images_folder, image_name)
        
        # 클래스 ID 추출
        with open(os.path.join(labels_folder, label_file), 'r') as file:
            class_id = file.readline().split()[0]  # 첫 번째 줄의 첫 번째 항목(클래스 ID)을 읽음
            
            # 클래스별 폴더 생성 및 이미지 이동
            class_folder = os.path.join(output_folder, class_id)
            if not os.path.exists(class_folder):
                os.makedirs(class_folder)
                
            # 이미지 파일 이동
            if os.path.exists(image_path):
                shutil.move(image_path, os.path.join(class_folder, image_name))

 

 

classes 폴더에는 이미지가, train.csv 파일에는 클래스 정보가 존재

csv 파일은 첫 번째 열에 이미지 파일의 이름이, 두 번째 열에 클래스를 기록

import os
import shutil
import pandas as pd

csv_file_path = './train/train.csv'
images_folder_path = './train/images'
output_folder_path = './train/classes'

for index, row in df.iterrows():
    # 파일명과 클래스명 추출
    file_name = row[0]
    class_name = row[1]

    # 해당 클래스의 폴더 경로 생성
    class_folder_path = os.path.join(output_folder_path, str(class_name))
    if not os.path.exists(class_folder_path):
        os.makedirs(class_folder_path)

    # 이미지 파일의 원본 경로와 목적지 경로
    src_path = os.path.join(images_folder_path, file_name)
    dst_path = os.path.join(class_folder_path, file_name)

    # 이미지 파일 이동
    shutil.move(src_path, dst_path)