Views
No views yet
1import transformers
2from huggingface_hub import hf_hub_download
3from PIL import Image
4import faiss
5import timm
6import torch
7import torch.nn as nn
8import joblib
9from torchvision import transforms
10
11if torch.cuda.is_available():
12 device = torch.device('cuda')
13else: device = torch.device('cpu')
14
15'''
16linear - knn --> 0 Real - 1 Fake
17svm --> -1 Real - 1 Fake
18'''
19class VITContrastiveHF(nn.Module):
20 def __init__(self, repo_name, classificator_type):
21 super(VITContrastiveHF, self).__init__()
22 self.model = transformers.AutoModel.from_pretrained(repo_name)
23 self.model.pooler= nn.Identity()
24
25 self.processor = transformers.AutoProcessor.from_pretrained(repo_name)
26 self.processor.do_resize= False
27 # define the correct classifier /// consider to use the `cache_dir`` parameter
28 if classificator_type == 'svm':
29 file_path = hf_hub_download(repo_id=repo_name, filename='sklearn/ocsvm_kernel_poly_gamma_auto_nu_0_1_crop.joblib')
30 self.classifier = joblib.load(file_path)
31
32 elif classificator_type == 'linear':
33 file_path = hf_hub_download(repo_id=repo_name, filename='sklearn/linear_tot_classifier_epoch-32.sav')
34 self.classifier = joblib.load(file_path)
35
36 elif classificator_type == 'knn':
37 file_path = hf_hub_download(repo_id=repo_name, filename='sklearn/knn_tot_classifier_epoch-32.sav')
38 self.classifier = joblib.load(file_path)
39
40 else:
41 raise ValueError('Selected an invalid classifier')
42
43 def forward(self, x, return_feature=False):
44 features = self.model(x)
45 if return_feature:
46 return features
47 features = features.last_hidden_state[:,0,:].cpu().detach().numpy()
48 # features.last_hidden_state[:,0,:].shape
49 predictions = self.classifier.predict(features)
50 return torch.from_numpy(predictions)
51
52# HF inference code
53classificator_type = 'linear'
54model = VITContrastiveHF(repo_name='aimagelab/CoDE', classificator_type=classificator_type)
55
56transform = transforms.Compose([
57 transforms.CenterCrop(224),
58 transforms.ToTensor(),
59 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
60])
61
62model.eval()
63model.model.to(device)
64y_pred= []
65# Put your image to tes
66img = Image.open('206496010652.png').convert('RGB')
67
68with torch.no_grad():
69 # in_tens = model.processor(img, return_tensors='pt')['pixel_values']
70 in_tens = transform(img).unsqueeze(0)
71
72 in_tens= in_tens.to(device)
73 y_pred.extend(model(in_tens).flatten().tolist())
74
75# check the correct label of the predict image
76for el in y_pred:
77 if el == 1:
78 print('Fake')
79 elif el == 0:
80 print('Real')
81 elif el == -1:
82 print('Real')
83 else:
84 print('Error')
85