similarity.py

import base64
import io
from typing import List
from skimage.metrics import structural_similarity as ssim
import cv2
import numpy as np
import requests
from PIL import Image

from models import RequestModel, ResponseModel


def load_image_url(source):
    if source.startswith('http'):
        img = Image.open(requests.get(source, stream=True).raw)
    else:
        img_data = base64.b64decode(source)
        img = Image.open(io.BytesIO(img_data))

    img = np.array(img.convert('L'))
    return img


def check_similarity(images: List[RequestModel]):
    print(f'checking similarity...')

    original_image = load_image_url(images[0].source)
    original_image_shape = original_image.shape

    results = []

    for i in range(1, len(images)):
        image = load_image_url(images[i].source)
        image = cv2.resize(image, original_image_shape[::-1])
        s, _ = ssim(original_image, image, full=True)
        similarity_score = (s + 1) * 50
        response = ResponseModel(originId=images[i].originId, sequence=images[i].sequence,
                                 assetCode=images[i].assetCode, similarity=similarity_score)
        results.append(response)
    return results