File size: 1,496 Bytes
cf69c91
 
0f98d7f
cf69c91
0f98d7f
868f784
0f98d7f
cf69c91
 
 
0f98d7f
 
868f784
319e2a1
0f98d7f
cf69c91
 
 
 
 
319e2a1
0f98d7f
 
319e2a1
0f98d7f
 
 
319e2a1
 
 
 
0f98d7f
319e2a1
cf69c91
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
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
from .model_implements.bit import BigTransfer


class Similarity:
    def get_models(self):
        return [
            model.SimilarityModel(name= 'Mobilenet V3', image_size= 224, model_cls = ModelnetV3()),
            model.SimilarityModel(name= 'Big Transfer (BiT)', image_size= 224, model_cls = BigTransfer()),
            model.SimilarityModel(name= 'Vision Transformer', image_size= 224, model_cls = VitBase(), image_input_type='pil'),
            ]        

    def check_similarity(self, img_urls, model):
        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), 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'{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 results