add new model

app.py

@@ -21,7 +21,7 @@ def check(img_main, img_1, img_2, model_idx):
         for i in range(3)
     ]
     similarity.check_similarity([img_main, img_1, img_2], models[model_idx])
-    return images
+    return []
 
 # def greet(name):
 #     return "Hello " + name + "!!"

src/similarity/model_implements/mobilenet_v3.py

@@ -1,4 +1,14 @@
+import tensorflow_hub as hub
+import numpy as np 
+
 class ModelnetV3():
-    def extract_feature(self, img):
+    def __init__(self):
+        module_handle = "https://tfhub.dev/google/imagenet/mobilenet_v3_large_100_224/feature_vector/5" 
+        self.module = hub.load(module_handle)
+
+    def extract_feature(self, imgs):
         print('getting with ModelnetV3...')
-        return []
+        features = []
+        for img in imgs:
+            features.append(np.squeeze(self.module(img)))
+        return features

src/similarity/model_implements/vit_base.py

@@ -0,0 +1,16 @@
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+
+class VitBase():
+
+    def __init__(self):
+        self.model = ViTFeatureExtractor.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)
+        return features

src/similarity/similarity.py

@@ -1,10 +1,16 @@
 from src.model import simlarity_model as model 
 from src.util import image as image_util
+from src.util import matrix
 from .model_implements.mobilenet_v3 import ModelnetV3
+from .model_implements.vit_base import VitBase
+
 
 class Similarity:
     def get_models(self):
-        return [model.SimilarityModel(name= 'mobilenet_v3', image_size= 224, model_cls = ModelnetV3())]        
+        return [
+            model.SimilarityModel(name= 'Mobilenet V3', image_size= 224, model_cls = ModelnetV3()),
+            model.SimilarityModel(name= 'Vision Transformer', image_size= 224, model_cls = VitBase()),
+            ]        
 
     def check_similarity(self, img_urls, model):
         # model = self.get_models()[model_idx]
@@ -12,7 +18,13 @@ class Similarity:
         for url in img_urls:
             if url == "": continue
             imgs.append(image_util.load_image_url(url, required_size=(model.image_size, model.image_size)))
-        model.model_cls.extract_feature(imgs[0])
+        
+        features = model.model_cls.extract_feature(imgs)
+        for i, v in enumerate(features):
+            if i == 0: continue 
+            dist = matrix.cosine(features[0], v)
+            print(f'distance: {dist}')
+
         return 'oke'
 
     

src/util/image.py

@@ -6,5 +6,5 @@ def load_image_url(url, required_size = (224,224)):
     img = Image.open(requests.get(url, stream=True).raw)
     img = Image.fromarray(np.array(img))
     img = img.resize(required_size)
-    img = (np.expand_dims(np.array(img), 0)/255).astype(np.float32)
+    #img = (np.expand_dims(np.array(img), 0)/255).astype(np.float32)
     return img

src/util/matrix.py

@@ -0,0 +1,5 @@
+from numpy.linalg import norm
+import numpy as np
+
+def cosine(x, y):
+    return np.dot(x,y)/(norm(x)*norm(y))