Upload clip.py

clip.py

@@ -0,0 +1,70 @@
+import os
+import torch
+from PIL import Image
+from transformers import ChineseCLIPProcessor, ChineseCLIPModel
+
+def clip_result(image_path):
+    # 設置設備
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Get the directory where this script is located
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    
+    # Construct the full path to the file in the subfolder
+    model_path = os.path.join(script_dir, 'model', 'best_clip_model.pth')
+
+    # 載入訓練好的模型和處理器
+    model = ChineseCLIPModel.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model = model.to(device)
+    model.eval()
+
+    processor = ChineseCLIPProcessor.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")
+
+    # 1. 加載圖片
+    # image_path = '/content/drive/MyDrive/幽靈吉伊卡哇.png'
+    image = Image.open(image_path)
+
+    # 2. 加載中文詞彙表
+    with open('/content/drive/MyDrive/chiikawa/word_list.txt', 'r', encoding='utf-8') as f:
+        vocab = [line.strip() for line in f.readlines()]
+
+    # 3. 圖像和文本處理
+    batch_size = 16  # 每次處理16個詞彙
+    similarities = []
+
+    # 釋放未使用的顯存
+    torch.cuda.empty_cache()
+
+    with torch.no_grad():
+        for i in range(0, len(vocab), batch_size):
+            batch_vocab = vocab[i:i + batch_size]
+            inputs = processor(
+                text=batch_vocab,
+                images=image,
+                return_tensors="pt",
+                padding=True
+            ).to(device)
+
+            # 推理並進行相似度計算
+            outputs = model(**inputs)
+            image_embeds = outputs.image_embeds / outputs.image_embeds.norm(dim=-1, keepdim=True)
+            text_embeds = outputs.text_embeds / outputs.text_embeds.norm(dim=-1, keepdim=True)
+            similarity = torch.matmul(image_embeds, text_embeds.T).squeeze(0)
+            similarities.append(similarity)
+
+    # 4. 合併所有相似度
+    similarity = torch.cat(similarities, dim=0)
+
+    # 5. 找到相似度最高的詞彙
+    top_k = 3
+    top_k_indices = torch.topk(similarity, top_k).indices.tolist()
+    top_k_words = [vocab[idx] for idx in top_k_indices]
+
+    # 6. 輸出最接近的前3名中文詞彙
+
+    print("最接近的前3名中文詞彙是:")
+    for rank, word in enumerate(top_k_words, 1):
+        print(f"{rank}. {word}")
+
+    return top_k_words