clip_util.py

import torch
import clip

# @torch.no_grad()
# def build_normalized_text_features(model, classnames, templates, device):
#     """
#     Build a zero-shot classifier weight matrix from class names and text templates.
#     Each class is represented by the mean of its template embeddings (L2-normalized).

#     **!! THIS FUNCTION IS NOT MEANT TO GENERATE TRAINABLE SAMPLES**
    
#     Returns:
#         zeroshot_weights: [embed_dim, num_classes] tensor
#     """
#     zeroshot_weights = []
#     for classname in classnames:
#         if callable(templates):
#             # Single template function
#             texts = [templates(classname)]
#         elif isinstance(templates, list) and callable(templates[0]):
#             # List of template functions
#             texts = [t(classname) for t in templates]
#         else:
#             texts = [templates(classname)]

#         texts = clip.tokenize(texts).to(device)
#         class_embeddings = model(None, texts)
#         # average all templates
#         class_embeddings = class_embeddings / class_embeddings.norm(dim=-1, keepdim=True)
#         class_embedding = class_embeddings.mean(dim=0)
#         # normalize
#         class_embedding = class_embedding / class_embedding.norm()
#         zeroshot_weights.append(class_embedding)

#     zeroshot_weights = torch.stack(zeroshot_weights, dim=1)  # [embed_dim, num_classes]
#     # 直接就是列向量
#     return zeroshot_weights

def build_normalized_text_features_from_text_embedding(model, texts: torch.Tensor):
    """
    Build a zero-shot classifier weight matrix from class text clip embeddings.
    
    Returns:
        class_features: [embed_dim, num_classes] tensor
    """
    class_features = model(None, texts)
    # normalize
    class_features = class_features / class_features.norm(dim=-1, keepdim=True)
    return class_features

def build_normalized_image_features_from_image_embedding(model, images: torch.Tensor):
    """
    Build a zero-shot classifier weight matrix from class image clip embeddings.
    
    Returns:
        image_features: [embed_dim, num_classes] tensor
    """
    image_features = model(images, None)
    # normalize
    image_features = image_features / image_features.norm(dim=-1, keepdim=True)
    return image_features