Spaces:
Sleeping
Sleeping
File size: 586 Bytes
0788e19 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | from .clip import clip
from PIL import Image
import torch.nn as nn
class CLIPModel(nn.Module):
def __init__(self, name="ViT-L/14", num_classes=1):
super(CLIPModel, self).__init__()
self.model, self.preprocess = clip.load(name, device="cpu") # self.preprecess will not be used during training, which is handled in Dataset class
self.fc = nn.Linear(768, num_classes )
def forward(self, x, return_feature=False):
features = self.model.encode_image(x)
if return_feature:
return features
return self.fc(features)
|