similarity.py

import base64
from typing import List
from skimage.metrics import structural_similarity as ssim
import cv2
import numpy as np
import requests
from models import RequestModel, ResponseModel
from PIL import Image
from io import BytesIO
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def load_image_url(source):
    Image.MAX_IMAGE_PIXELS = None

    if source.startswith('http'):
        response = requests.get(source)
        img = np.asarray(bytearray(response.content), dtype=np.uint8)
        img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)
    else:
        img = base64.b64decode(source)
        img = Image.open(BytesIO(img))
        img = np.array(img)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

    return img

def check_similarity(images: List[RequestModel]):
    logging.info(f"Checking similarity for main source with resource id {images[0].originId}")

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

    results = []

    for i in range(1, len(images)):
        try:
            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
        except Exception as e:
            logging.error(f"Error loading image for resource id {images[i].originId} : {e}")
            similarity_score = 0

        response = ResponseModel(originId=images[i].originId, sequence=images[i].sequence,
                                 assetCode=images[i].assetCode, similarity=similarity_score)
        results.append(response)

    return results