Pandas クレンジング ラベル付き学習データの分割

やったこと

0から9のラベル付き学習データを任意のラベルで分割する関数を作成しました。
今回処理するテキストのサンプルを示します。

sample.txt

data1   0
data2   1
data3   2
data4   3
data5   4
data6   5
data7   6
data8   7
data9   8
data10  9

データとそのラベルが後ろについています。特にheader等もありません。
今回は特定のラベルを複数選択し、sample.txtを分割します。
以下、ソースコードです。

conditional_list_split.py

# -*- coding: utf-8 -*-
import glob
import os
import argparse
from PIL import Image
import cv2
import matplotlib.pyplot as plt
import numpy
import re
import pandas as pd


def conditional_list_split(lbl, path):
    lbls = lbl.split(',')
    df = pd.read_table(path, sep=' ', encoding='utf-8',header=None)

    df.columns = ['data', 'label']
    df_slct = df.where(df['label'].isin(lbls)).dropna(how='all')
    df_no_slct = df.where(~df['label'].isin(lbls)).dropna(how='all')

    df_slct['label'] = df_slct['label'].astype('int64')
    df_no_slct['label'] = df_no_slct['label'].astype('int64')
    dirname = os.path.dirname(path)
    slct_path = dirname + '\\' + os.path.splitext(os.path.basename(path))[0] + '_target.txt'
    no_slct_path = dirname + '\\' + os.path.splitext(os.path.basename(path))[0] + '_no_target.txt'
    df_slct.to_csv(slct_path, header=False, index=False, sep=' ', encoding='utf-8')
    df_no_slct.to_csv(no_slct_path, header=False, index=False, sep=' ', encoding='utf-8')


def main():
    parser = argparse.ArgumentParser(description='conditional_list_split')
    parser.add_argument('--conditions', required=True, default=None)
    parser.add_argument('--datas', help='Path to label list file')
    args = parser.parse_args()
    labels = args.conditions
    conditional_list_split(labels, args.datas)


if __name__ == '__main__':
    main()

使い方

コマンドライン.txt

python conditional_list_split.py --conditions 0,1,2 --datas sample.txt

実行結果例

sample_target.txt

data1   0
data2   1
data3   2

sample_no_target.txt

data4   3
data5   4
data6   5
data7   6
data8   7
data9   8
data10  9