Update main.py

main.py

@@ -10,6 +10,11 @@ import base64
 import requests
 from PIL import Image
 from io import BytesIO
+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
 
 BASE_DIR = "saved_data"
 app = FastAPI()
@@ -28,6 +33,39 @@ def orb_sim(img1, img2):
         return 0
     return len(similar_regions) / len(matches)
 
+def preprocess_image_for_mobilenet(image):
+    # Garantir que a imagem tem 3 canais
+    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)
+
+    # Redimensionar e preparar imagem
+    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_similarity(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]  # Valor entre -1 e 1
+        sim_score = (sim + 1) * 50  # Escalar para 0-100
+        logging.info(f"MobileNet similarity score is {sim_score}")
+        return sim_score
+    except Exception as e:
+        logging.error("Erro ao calcular similaridade com MobileNet", exc_info=True)
+        return 0
+
 def load_image(source):
     Image.MAX_IMAGE_PIXELS = None
 
@@ -56,13 +94,18 @@ async def save(image_data: RequestModel):
     img2 = load_image(image_data.source)
 
     similarity_orb = None
+    similarity_mobilenet = None
+
     if img1 is not None and img2 is not None:
         similarity_orb = orb_sim(img1, img2)
         print(f"Similaridade ORB entre {image_data.originSource} e {image_data.source}: {similarity_orb}")
+        similarity_mobilenet = mobilenet_similarity(original_image, image)
+        print(f"Similaridade Mobilenet entre {image_data.originSource} e {image_data.source}: {similarity_mobilenet}")
 
     data_to_save = image_data.dict()
     if similarity_orb is not None:
         data_to_save["similarityOrb"] = similarity_orb
+        data_to_save["similarityMobilenet"] = similarity_mobilenet
 
     with open(filename, "w") as f:
         json.dump(data_to_save, f, indent=4)