finish vit

app.py

@@ -7,27 +7,9 @@ from src.similarity.similarity import Similarity
 similarity = Similarity()
 models = similarity.get_models()
 
-def check(img_main, img_1, img_2, model_idx):
-    images = [
-        (random.choice(
-            [
-                "https://images.unsplash.com/photo-1507003211169-0a1dd7228f2d?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=387&q=80",
-                "https://images.unsplash.com/photo-1554151228-14d9def656e4?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=386&q=80",
-                "https://images.unsplash.com/photo-1542909168-82c3e7fdca5c?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxzZWFyY2h8MXx8aHVtYW4lMjBmYWNlfGVufDB8fDB8fA%3D%3D&w=1000&q=80",
-                "https://images.unsplash.com/photo-1546456073-92b9f0a8d413?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=387&q=80",
-                "https://images.unsplash.com/photo-1601412436009-d964bd02edbc?ixlib=rb-1.2.1&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=464&q=80",
-            ]
-        ), f"label {i}" if i != 0 else "label" * 50)
-        for i in range(3)
-    ]
-    similarity.check_similarity([img_main, img_1, img_2], models[model_idx])
-    return []
-
-# def greet(name):
-#     return "Hello " + name + "!!"
-
-# iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-# iface.launch()
+def check(img_main, img_1, img_2, model_idx):    
+    result = similarity.check_similarity([img_main, img_1, img_2], models[model_idx])
+    return result
 
 with gr.Blocks() as demo:
     gr.Markdown('Checking Image Similarity')
@@ -41,7 +23,7 @@ with gr.Blocks() as demo:
     model = gr.Dropdown([m.name for m in models], label='Model', type='index')
     
     gallery = gr.Gallery(
-            label="Generated images", show_label=True, elem_id="gallery"
+            label="Generated images", show_label=False, elem_id="gallery"
         ).style(grid=[2], height="auto")
 
     submit_btn = gr.Button('Check Similarity')

src/model/simlarity_model.py

@@ -5,4 +5,5 @@ from .similarity_interface import SimilarityInterface
 class SimilarityModel:
     name: str 
     image_size: int
-    model_cls: SimilarityInterface
+    model_cls: SimilarityInterface
+    image_input_type: str = 'array'

src/similarity/model_implements/vit_base.py

@@ -1,21 +1,20 @@
-from transformers import ViTFeatureExtractor, ViTForImageClassification
+from transformers import ViTFeatureExtractor, ViTModel
 from PIL import Image
 import numpy as np 
+import torch
 
 class VitBase():
 
     def __init__(self):
-        self.feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
-        self.model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
+        self.feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
+        self.model = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
     
     def extract_feature(self, imgs):
         features = []
         for img in imgs:
-            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)
+            inputs = self.feature_extractor(images=img, return_tensors="pt")
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+            last_hidden_states =  outputs.last_hidden_state            
+            features.append(np.squeeze(last_hidden_states.numpy()).flatten())            
         return features

src/similarity/similarity.py

@@ -9,7 +9,7 @@ class Similarity:
     def get_models(self):
         return [
             model.SimilarityModel(name= 'Mobilenet V3', image_size= 224, model_cls = ModelnetV3()),
-            model.SimilarityModel(name= 'Vision Transformer', image_size= 224, model_cls = VitBase()),
+            model.SimilarityModel(name= 'Vision Transformer', image_size= 224, model_cls = VitBase(), image_input_type='pil'),
             ]        
 
     def check_similarity(self, img_urls, model):
@@ -17,14 +17,18 @@ class Similarity:
         imgs = []
         for url in img_urls:
             if url == "": continue
-            imgs.append(image_util.load_image_url(url, required_size=(model.image_size, model.image_size)))
+            imgs.append(image_util.load_image_url(url, required_size=(model.image_size, model.image_size), image_type=model.image_input_type))
         
         features = model.model_cls.extract_feature(imgs)
+        results = []
         for i, v in enumerate(features):
             if i == 0: continue 
             dist = matrix.cosine(features[0], v)
-            # print(f'distance: {dist}')
+            print(f'{i} -- distance: {dist}')
+            # results.append((imgs[i], f'similarity: {int(dist*100)}%'))
+            original_img = image_util.load_image_url(img_urls[i], required_size=None, image_type='pil')
+            results.append((original_img, f'similarity: {int(dist*100)}%'))
 
-        return 'oke'
+        return results
 
     

src/util/image.py

@@ -2,9 +2,12 @@ from PIL import Image
 import numpy as np
 import requests
 
-def load_image_url(url, required_size = (224,224)):
+def load_image_url(url, required_size = (224,224), image_type = 'array'):
+    print(f'downloading.. {url}, type: {image_type}')
     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)
+    if required_size is not None:
+        img = img.resize(required_size)
+    if image_type == 'array':
+        img = (np.expand_dims(np.array(img), 0)/255).astype(np.float32)
     return img