AI を使って類似する画像を調べる
Python プログラム
インストール
pip install git+https://github.com/openai/CLIP.git
類似していない画像を調べてコピーするプログラムです
import os
import shutil
import torch
import clip
from PIL import Image
def copy_not_similar_images(query_image_path, image_folder, output_folder, threshold=0.9):
"""
クエリ画像と各画像のコサイン類似度をCLIPで計算し、
類似度が threshold 未満(=類似していない)画像を output_folder にコピーします。
Parameters:
query_image_path (str): クエリ画像のファイルパス
image_folder (str): 検索対象の画像フォルダのパス
output_folder (str): 類似していない画像を出力(コピー)するフォルダのパス
threshold (float): 類似度の閾値(0~1の範囲)
"""
# 使用可能なデバイス(GPUがあればGPUを使用)
device = "cuda" if torch.cuda.is_available() else "cpu"
# CLIPモデルと前処理関数のロード
model, preprocess = clip.load("ViT-B/32", device=device)
# クエリ画像の読み込み・前処理
query_image = Image.open(query_image_path)
query_input = preprocess(query_image).unsqueeze(0).to(device)
with torch.no_grad():
query_embedding = model.encode_image(query_input)
query_embedding /= query_embedding.norm(dim=-1, keepdim=True)
# 出力フォルダが存在しない場合は作成
if not os.path.exists(output_folder):
os.makedirs(output_folder)
# フォルダ内の画像を処理
for filename in os.listdir(image_folder):
if filename.lower().endswith(('png', 'jpg', 'jpeg', 'bmp', 'gif')):
image_path = os.path.join(image_folder, filename)
try:
image = Image.open(image_path)
image_input = preprocess(image).unsqueeze(0).to(device)
with torch.no_grad():
image_embedding = model.encode_image(image_input)
image_embedding /= image_embedding.norm(dim=-1, keepdim=True)
# コサイン類似度の計算
similarity = (query_embedding @ image_embedding.T).item()
# しきい値未満=類似していない画像をコピー
if similarity < threshold:
shutil.copy2(image_path, output_folder)
print(f"Copied {filename}: similarity {similarity:.2f}")
except Exception as e:
print(f"画像 {filename} の処理中にエラーが発生しました: {e}")
if __name__ == "__main__":
# 各種パスを指定(raw stringでエスケープの問題を回避)
query_image_path = r"C:\data\1.png"
image_folder = r"C:\data\tmp2"
output_folder = r"D:\WinPython\notebooks\output_folder"
# 類似度のしきい値。0.9未満なら「類似していない」と判断
threshold = 0.9
copy_not_similar_images(query_image_path, image_folder, output_folder, threshold)
ディスカッション
コメント一覧
まだ、コメントがありません