Moonの日記

AI, プログラミング, ポケモンなど

Google colabでicrawlerを使った画像収集方法

Stable diffusionで使う学習モデルを作る際にGoogle colabでicrawlerを使って画像収集しました。備忘録のため残しておきます。
まずはicrawlerをインストール。

〜
!pip install icrawler
〜

今回は新木優子さんの画像を100枚集めます。

〜
from icrawler.builtin import BingImageCrawler
crawler = BingImageCrawler(storage = {'root_dir' : './drive/MyDrive/ColabNotebooks/araki_yuko'})
crawler.crawl(keyword = '新木優子', max_num = 100)
〜

重複画像を削除します。

〜
import os
from hashlib import md5

# フォルダのパス
folder_path = './drive/MyDrive/ColabNotebooks/araki_yuko'

# ファイルのハッシュ値を格納するディクショナリ
hash_dict = {}

# ファイルを検索してハッシュ値を取得する
for filename in os.listdir(folder_path):
    if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.gif')):
        # 画像ファイルの場合、ハッシュ値を取得する
        with open(os.path.join(folder_path, filename), 'rb') as f:
            filehash = md5(f.read()).hexdigest()

        # ハッシュ値を辞書に追加する
        if filehash in hash_dict:
            # 既に同じハッシュ値が存在する場合、重複した画像とみなす
            print(f"重複したファイルを検出しました: {filename}")
            os.remove(os.path.join(folder_path, filename))
        else:
            hash_dict[filehash] = filename
〜

後で処理しやすいように連番にファイル名変更します。

〜
import os

# フォルダのパス
folder_path = './drive/MyDrive/ColabNotebooks/araki_yuko'

# ファイル名の連番の開始番号
start_number = 1

# ファイル名を変更する
for i, filename in enumerate(os.listdir(folder_path)):
    # 拡張子を取得する
    ext = os.path.splitext(filename)[1]

    # 新しいファイル名を生成する
    new_filename = f"{start_number + i:03d}{ext}"

    # ファイル名を変更する
    os.rename(os.path.join(folder_path, filename), os.path.join(folder_path, new_filename))
〜