Update similarity.py

similarity.py

@@ -1,16 +1,52 @@
 import base64
-from typing import List
-from skimage.metrics import structural_similarity as ssim
 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
-import logging
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
+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
 
@@ -31,6 +67,10 @@ def orb_sim(img1, img2):
 
     return 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
@@ -59,16 +99,17 @@ def check_similarity(images: List[RequestModel]):
         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
+        
+            similarity_score = ssim_sim(original_image, image)
             similarity_orb_score = orb_sim(original_image, image) * 100
+            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
 
         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)
+                                 assetCode=images[i].assetCode, similarity=similarity_score, similarityOrb=similarity_orb_score, similarityMobileNet=similarity_mobilenet_score)
         results.append(response)
 
     return results