[AI] Custom Dataset 작성

Dataset 작성 중
품목, 품종, 크기로 3단계로 할 수 있어서 어떻게 분류하는 게 좋을지 고민.

def data_split(path="./data/org", split_predictions=0.1):
    # train_dict, valid_dict로 데이터를 담을 변수를 만들어주고
    train_dict = {}
    valid_dict = {}
    file_paths = []

    for root, dirs, files in os.walk(path):
        # 파일을 다 돌면서 file_paths를 구성해주고
        # 10%는 valid_index로,
        valid_index = np.random.randint(low=0, high=len(files), size=int(len(files) * split_predictions))
        # 나머지는 train_index로 만들어준다.
        train_index = list(set(range(0, len(files))) - set(valid_index))

        for file in files:
            file_paths.append(os.path.join(root, file))

            # label = str(file).split('_')
            train_dict[str(file).split('_')[0]] = [file_paths[-1] for i in train_index]
            valid_dict[str(file).split('_')[0]] = [file_paths[-1] for i in valid_index]

    # train_dict, valid_dict를 반환하여 데이터셋을 구성한다.
    return train_dict, valid_dict


# data_split을 호출하여 train_split_data와 valid_split_data에 각각 분할한 데이터를 받아온다.
train_split_data, valid_split_data = data_split()

# 출력해보면 몇 개씩 나누어졌는지 볼 수 있다.
print("Train data size : ", [len(l) for l in train_split_data.values()])
print("Valid data size : ", [len(l) for l in valid_split_data.values()])

Train data size : [2153, 2153, 2153, 2153, 2153, 2153, 2153, 2153, 2153, 2153]
Valid data size : [238, 238, 238, 238, 238, 238, 238, 238, 238, 238]

잘 나누어졌다.