Update similarity.py

similarity.py

@@ -1,5 +1,4 @@
 import base64
-import io
 from typing import List
 from skimage.metrics import structural_similarity as ssim
 import cv2
@@ -7,21 +6,28 @@ import numpy as np
 import requests
 from models import RequestModel, ResponseModel
 from PIL import Image
-Image.MAX_IMAGE_PIXELS = None
+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'):
-        img = Image.open(requests.get(source, stream=True).raw)
+        response = requests.get(source)
+        img = np.asarray(bytearray(response.content), dtype=np.uint8)
+        img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)
     else:
-        img_data = base64.b64decode(source)
-        img = Image.open(io.BytesIO(img_data))
+        img = base64.b64decode(source)
+        img = Image.open(BytesIO(img))
+        img = np.array(img)
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
 
-    img = np.array(img.convert('L'))
     return img
 
-
 def check_similarity(images: List[RequestModel]):
-    print(f'checking similarity...')
+    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
@@ -29,11 +35,17 @@ def check_similarity(images: List[RequestModel]):
     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
+        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