Trained on images labeled by myself.
1from PIL import Image
2import torch
3import torch.nn as nn
4import uform
5
6path = "image.jpg"
7
8# generate 768 dimension embeddings for an image
9uf_model = uform.get_model("unum-cloud/uform-vl-english")
10img = Image.open(path)
11image_data = uf_model.preprocess_image(img)
12text_data = uf_model.preprocess_text(path.name) # filename not necessarily useful, but encode anyway
13memb = uf_model.encode_multimodal(image=image_data, text=text_data)
14memb = memb.detach().numpy()
15assert memb.shape == (1, 768)
16
17# load model, infer and Sigmoid
18class ImageBinaryClassifier(nn.Module):
19 def __init__(self):
20 super().__init__()
21 self.layers = nn.Sequential(
22 nn.Linear(768, 2048),
23 nn.ReLU(),
24 nn.Linear(2048, 2048),
25 nn.ReLU(),
26 nn.Linear(2048, 2048),
27 nn.ReLU(),
28 nn.Linear(2048, 256),
29 nn.ReLU(),
30 nn.Linear(256, 1),
31 )
32 def forward(self, x):
33 return self.layers(x)
34
35saved_model = ImageBinaryClassifier()
36saved_model.load_state_dict(torch.load("model.pt", map_location=torch.device('cpu')))
37saved_model = torch.compile(saved_model)
38saved_model.eval()
39
40prob = float(torch.sigmoid(saved_model(memb)))
