similarity/similarity.py

from model import similarity_model as model
from util import image as image_util
from util import matrix
from .model_implements.bit import BigTransfer


class Similarity:
    def get_model(self):
        return model.SimilarityModel(name= 'Big Transfer (BiT)', image_size= 224, model_cls = BigTransfer())

    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)
        for i, v in enumerate(features):
            if i == 0:
                continue 
            dist = matrix.cosine(features[0], v)

        return dist