Update clip.py

clip.py

@@ -3,70 +3,61 @@ 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"
+class ClipModel:
+    def __init__(self, model_name="OFA-Sys/chinese-clip-vit-base-patch16", model_path=None):
+        # Set device
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Load model and processor
+        self.model = ChineseCLIPModel.from_pretrained(model_name)
+        if model_path is None:
+            script_dir = os.path.dirname(os.path.abspath(__file__))
+            model_path = os.path.join(script_dir, 'artifacts/models', 'best_clip_model.pth')
+        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
+        self.model = self.model.to(self.device)
+        self.model.eval()
+
+        self.processor = ChineseCLIPProcessor.from_pretrained(model_name)
+
+    def clip_result(self, image_path, vocab_path='./chiikawa/word_list.txt', top_k=3):
+        # Load image
+        image = Image.open(image_path)
+
+        # Load Chinese vocabulary
+        with open(vocab_path, 'r', encoding='utf-8') as f:
+            vocab = [line.strip() for line in f.readlines()]
+
+        # Process images and texts
+        batch_size = 16  # Process 16 vocab at a time
+        similarities = []
+
+        # Release unused memory
+        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 = self.processor(
+                    text=batch_vocab,
+                    images=image,
+                    return_tensors="pt",
+                    padding=True
+                ).to(self.device)
+
+                # Inference and similarity calculation
+                outputs = self.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)
+
+        # Merge all similarities
+        similarity = torch.cat(similarities, dim=0)
+
+        # Find top-3 similarities
+        top_k_indices = torch.topk(similarity, top_k).indices.tolist()
+        top_k_words = [vocab[idx] for idx in top_k_indices]
+        
+        # 6. 輸出最接近的前3名中文詞彙
+        return top_k_words
 
-    # 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, 'artifacts/models', 'best_clip_model.pth')
-
-    print("model_path:", model_path)
-
-    # 載入訓練好的模型和處理器
-    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('./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