[Python] LoRA 学習用タグファイルからタグを一括削除

はじめに

LoRAを学習させる時に画像毎のタグファイルが必要になります。
ひとつひとつ削除するのが面倒なので一括削除するスクリプトを作りました。

自分用に作ったので適当な作りですが、短いスクリプトなのでまぁ大丈夫でしょう。
Python触って一ヶ月くらいなのでもっと良い書き方があれば教えてください。

• 元のファイルを上書きします。
• Dataset Tag Editor みたいに親切ではないのでバックアップは取ってません。
• Tagger とかでスパっと出力できますし「間違えたら出力しなおせばいいや」という精神で。

スクリプト

以下の内容を removeTag.py として保存します。
IMG_PATH、REMOVE_TAG_FILE 自分の環境に合わせて変更してください。

images\**\*.txt というのは images\ 以下のサブフォルダも対象にするという意味です。

removeTag.py

# タグファイルから特定のタグを削除するだけのツール

import glob
import re

# 学習画像フォルダの１個上のフォルダ
IMG_PATH = "{学習ファイルのフォルダ}\images\**\*.txt"
# 削除対象を記述したファイル
REMOVE_TAG_FILE = "{学習ファイルのフォルダ}\remove_tag_list.txt"

_targetTags = []


# ---------------
# 除外タグに一致しないものを書き出す
def removeTag(filename, tags):
    removedTags = []

    for tag in tags:
        if not tag in _targetTags:
            removedTags.append(tag)

    return removedTags


# ---------------
# カンマと改行で分割してリスト化
def fileToList(filename):
    handl = open(filename, "r")
    list = re.split(" *[,\n]+ *", handl.read())
    handl.close()
    return list


# ---------------
def main():
    global _targetTags

    # 除外タグ一覧読み込み
    _targetTags = fileToList(REMOVE_TAG_FILE)

    # タグファイルを開く
    for filename in glob.glob(IMG_PATH):
        print("\n" + filename + "\n")

        tags = fileToList(filename)
        removedTags = removeTag(filename, tags)

        tagFile = open(filename, "w")
        tagFile.write(",".join(removedTags))
        tagFile.close()
        print(",".join(removedTags))


main()

削除対象タグ記述ファイル

remove_tag_list.txt というファイルに削除対象を記述します。
以下が記述例です。

remove_tag_list.txt

school swimsuit, swimsuit,swimsuit under clothes,
blue one-piece swimsuit,
one-piece swimsuit
old school swimsuit
sailor school uniform,

• カンマ「,」か改行でタグを区切ります
• 前後にある空白は無視します

使い方

> python removeTag.py

出力結果

slskmz, 1girl, solo, swimsuit, sailor school uniform

👇

slskmz,1girl,solo