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)

AI を使わない 類似画像検索 SimiPix

SimiPix

Python

Posted by eightban