import torch import clip import numpy as np import logging import traceback from typing import List, Dict, Tuple, Optional, Union, Any from PIL import Image class CLIPModelManager: """ 專門管理 CLIP 模型相關的操作，包括模型載入、設備管理、圖像和文本的特徵編碼等核心功能 """ def __init__(self, model_name: str = "ViT-B/16", device: str = None): """ 初始化 CLIP 模型管理器 Args: model_name: CLIP模型名稱，默認為"ViT-B/16" device: 運行設備，None則自動選擇 """ self.logger = logging.getLogger(__name__) self.model_name = model_name # 設置運行設備 if device is None: self.device = "cuda" if torch.cuda.is_available() else "cpu" else: self.device = device self.model = None self.preprocess = None self._initialize_model() def _initialize_model(self): """ 初始化CLIP模型 """ try: self.logger.info(f"Initializing CLIP model ({self.model_name}) on {self.device}") self.model, self.preprocess = clip.load(self.model_name, device=self.device) self.logger.info("Successfully loaded CLIP model") except Exception as e: self.logger.error(f"Error loading CLIP model: {e}") self.logger.error(traceback.format_exc()) raise def encode_image(self, image_input: torch.Tensor) -> torch.Tensor: """ 編碼圖像特徵 Args: image_input: 預處理後的圖像張量 Returns: torch.Tensor: 標準化後的圖像特徵 """ try: with torch.no_grad(): image_features = self.model.encode_image(image_input) image_features = image_features / image_features.norm(dim=-1, keepdim=True) return image_features except Exception as e: self.logger.error(f"Error encoding image features: {e}") self.logger.error(traceback.format_exc()) raise def encode_text_batch(self, text_prompts: List[str], batch_size: int = 128) -> torch.Tensor: """ 批量編碼文本特徵，避免CUDA內存問題 Args: text_prompts: 文本提示列表 batch_size: 批處理大小 Returns: torch.Tensor: 標準化後的文本特徵 """ if not text_prompts: return None try: with torch.no_grad(): features_list = [] for i in range(0, len(text_prompts), batch_size): batch_prompts = text_prompts[i:i+batch_size] text_tokens = clip.tokenize(batch_prompts).to(self.device) batch_features = self.model.encode_text(text_tokens) batch_features = batch_features / batch_features.norm(dim=-1, keepdim=True) features_list.append(batch_features) # 連接所有批次 if len(features_list) > 1: text_features = torch.cat(features_list, dim=0) else: text_features = features_list[0] return text_features except Exception as e: self.logger.error(f"Error encoding text features: {e}") self.logger.error(traceback.format_exc()) raise def encode_single_text(self, text_prompts: List[str]) -> torch.Tensor: """ 編碼單個文本批次的特徵 Args: text_prompts: 文本提示列表 Returns: torch.Tensor: 標準化後的文本特徵 """ try: with torch.no_grad(): text_tokens = clip.tokenize(text_prompts).to(self.device) text_features = self.model.encode_text(text_tokens) text_features = text_features / text_features.norm(dim=-1, keepdim=True) return text_features except Exception as e: self.logger.error(f"Error encoding single text batch: {e}") self.logger.error(traceback.format_exc()) raise def calculate_similarity(self, image_features: torch.Tensor, text_features: torch.Tensor) -> np.ndarray: """ 計算圖像和文本特徵之間的相似度 Args: image_features: 圖像特徵張量 text_features: 文本特徵張量 Returns: np.ndarray: 相似度分數數組 """ try: with torch.no_grad(): similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1) similarity = similarity.cpu().numpy() if self.device == "cuda" else similarity.numpy() return similarity except Exception as e: self.logger.error(f"Error calculating similarity: {e}") self.logger.error(traceback.format_exc()) raise def preprocess_image(self, image: Union[Image.Image, np.ndarray]) -> torch.Tensor: """ 預處理圖像以供CLIP模型使用 Args: image: PIL圖像或numpy數組 Returns: torch.Tensor: 預處理後的圖像張量 """ try: if not isinstance(image, Image.Image): if isinstance(image, np.ndarray): image = Image.fromarray(image) else: raise ValueError("Unsupported image format. Expected PIL Image or numpy array.") image_input = self.preprocess(image).unsqueeze(0).to(self.device) return image_input except Exception as e: self.logger.error(f"Error preprocessing image: {e}") self.logger.error(traceback.format_exc()) raise def process_image_region(self, image: Union[Image.Image, np.ndarray], box: List[float]) -> torch.Tensor: """ 處理圖像的特定區域 Args: image: 原始圖像 box: 邊界框 [x1, y1, x2, y2] Returns: torch.Tensor: 區域圖像的特徵 """ try: # 確保圖像是PIL格式 if not isinstance(image, Image.Image): if isinstance(image, np.ndarray): image = Image.fromarray(image) else: raise ValueError("Unsupported image format. Expected PIL Image or numpy array.") # 裁剪區域 x1, y1, x2, y2 = map(int, box) cropped_image = image.crop((x1, y1, x2, y2)) # 預處理並編碼 image_input = self.preprocess_image(cropped_image) image_features = self.encode_image(image_input) return image_features except Exception as e: self.logger.error(f"Error processing image region: {e}") self.logger.error(traceback.format_exc()) raise def batch_process_regions(self, image: Union[Image.Image, np.ndarray], boxes: List[List[float]]) -> torch.Tensor: """ 批量處理多個圖像區域 Args: image: 原始圖像 boxes: 邊界框列表 Returns: torch.Tensor: 所有區域的圖像特徵 """ try: # ensure PIL format if not isinstance(image, Image.Image): if isinstance(image, np.ndarray): image = Image.fromarray(image) else: raise ValueError("Unsupported image format. Expected PIL Image or numpy array.") if not boxes: return torch.empty(0) # 裁剪並預處理所有區域 cropped_inputs = [] for box in boxes: x1, y1, x2, y2 = map(int, box) cropped_image = image.crop((x1, y1, x2, y2)) processed_image = self.preprocess(cropped_image).unsqueeze(0) cropped_inputs.append(processed_image) # 批量處理 batch_tensor = torch.cat(cropped_inputs).to(self.device) image_features = self.encode_image(batch_tensor) return image_features except Exception as e: self.logger.error(f"Error batch processing regions: {e}") self.logger.error(traceback.format_exc()) raise def is_model_loaded(self) -> bool: """ 檢查模型是否已成功載入 Returns: bool: 模型載入狀態 """ return self.model is not None and self.preprocess is not None def get_device(self) -> str: """ 獲取當前設備 Returns: str: 設備名稱 """ return self.device def get_model_name(self) -> str: """ 獲取模型名稱 Returns: str: 模型名稱 """ return self.model_name