31 lines
1.0 KiB
Python
31 lines
1.0 KiB
Python
import torch
|
|
import numpy as np
|
|
|
|
use_cuda = torch.cuda.is_available()
|
|
|
|
def image_embeddings_direct(image, model, processor):
|
|
inputs = processor(images=image, return_tensors='pt')['pixel_values']
|
|
if use_cuda:
|
|
inputs = inputs.to('cuda')
|
|
result = model.get_image_features(pixel_values=inputs).cpu().detach().numpy()
|
|
return (result / np.linalg.norm(result)).squeeze(axis=0)
|
|
|
|
# binary classifier that consumes CLIP embeddings
|
|
class Classifier(torch.nn.Module):
|
|
def __init__(self, input_size, hidden_size, output_size):
|
|
super().__init__()
|
|
self.fc1 = torch.nn.Linear(input_size, hidden_size)
|
|
self.fc2 = torch.nn.Linear(hidden_size, hidden_size//2)
|
|
self.fc3 = torch.nn.Linear(hidden_size//2, output_size)
|
|
self.relu = torch.nn.ReLU()
|
|
self.sigmoid = torch.nn.Sigmoid()
|
|
|
|
def forward(self, x):
|
|
x = self.fc1(x)
|
|
x = self.relu(x)
|
|
x = self.fc2(x)
|
|
x = self.relu(x)
|
|
x = self.fc3(x)
|
|
x = self.sigmoid(x)
|
|
return x
|