try vit

src/similarity/model_implements/vit_base.py

@@ -1,16 +1,21 @@
 from transformers import ViTFeatureExtractor, ViTForImageClassification
 from PIL import Image
+import numpy as np 
 
 class VitBase():
 
     def __init__(self):
-        self.model = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
+        self.feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
+        self.model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
     
     def extract_feature(self, imgs):
         features = []
         for img in imgs:
-            feature = self.model(images=img, return_tensors="np")
-            print('type::', type(**feature))
-            features.append(feature)
-            print(features[0].shape)
+            feature = self.feature_extractor(images=img, return_tensors="tf")
+            print('keys: ', feature.keys())
+            f = self.model(feature)
+            print('--> f', type(f))
+            # print('type::', (feature['pixel_values'].shape))
+            features.append(np.squeeze(feature['pixel_values']))
+            print('shape:::',features[0].shape)
         return features

src/similarity/similarity.py

@@ -23,7 +23,7 @@ class Similarity:
         for i, v in enumerate(features):
             if i == 0: continue 
             dist = matrix.cosine(features[0], v)
-            print(f'distance: {dist}')
+            # print(f'distance: {dist}')
 
         return 'oke'