similarity.py

import base64
import cv2
import numpy as np
import requests
import logging
from typing import List
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing.image import img_to_array
from sklearn.metrics.pairwise import cosine_similarity
from skimage.metrics import structural_similarity as ssim
from models import RequestModel, ResponseModel
from PIL import Image
from io import BytesIO

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
mobilenet = MobileNetV2(weights="imagenet", include_top=False, pooling='avg')

def preprocess_image_for_mobilenet(image):
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    elif image.shape[2] == 1:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    else:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    image = cv2.resize(image, (224, 224))
    image = img_to_array(image)
    image = np.expand_dims(image, axis=0)
    image = preprocess_input(image)

    return image

def mobilenet_sim(img1, img2):
    try:
        img1_proc = preprocess_image_for_mobilenet(img1)
        img2_proc = preprocess_image_for_mobilenet(img2)

        feat1 = mobilenet.predict(img1_proc, verbose=0)
        feat2 = mobilenet.predict(img2_proc, verbose=0)

        sim = cosine_similarity(feat1, feat2)[0][0] 
        sim_score = (sim + 1) * 50 
        print(f"MobileNet similarity score is {sim_score}")
        return float(sim_score)
    except Exception as e:
        logging.error("Erro ao calcular similaridade com MobileNet", exc_info=True)
        return 0

def orb_sim(img1, img2):
    score = 0

    try:
        orb = cv2.ORB_create()
        kp_a, desc_a = orb.detectAndCompute(img1, None)
        kp_b, desc_b = orb.detectAndCompute(img2, None)
    
        bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        matches = bf.match(desc_a, desc_b)
        similar_regions = [i for i in matches if i.distance < 20]
        if len(matches) > 0:
            score = (len(similar_regions) / len(matches)) * 100
            if (score > 0):
                logging.info(f"Orb score is {score}")
    except Exception as e:
        logging.error("Erro ao verificar similaridade ORB", exc_info=True)

    return round(score) if 0 < score < 1 else score

def ssim_sim(img1, img2):
    s, _ = ssim(img1, img2, full=True)
    return (s + 1) * 50


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])
        
            similarity_score = ssim_sim(original_image, image)
            similarity_orb_score = orb_sim(original_image, image)
            similarity_mobilenet_score = mobilenet_sim(original_image, image)
        except Exception as e:
            logging.error(f"Error loading image for resource id {images[i].originId} : {e}")
            similarity_score = 0
            similarity_orb_score = 0
            similarity_mobilenet_score = 0

        response = ResponseModel(originId=images[i].originId, source=images[i].source, sequence=images[i].sequence,
                                 assetCode=images[i].assetCode, similarity=similarity_score, similarityOrb=similarity_orb_score, similarityMobileNet=similarity_mobilenet_score)
        results.append(response)

    return results