main.py

from typing import List
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from models import RequestModel
import os
import json
import cv2
import numpy as np
import base64
import requests

import mimetypes
import tempfile
import subprocess


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()
mobilenet = MobileNetV2(weights="imagenet", include_top=False, pooling='avg')

def orb_sim(img1, img2):
    # ORB
    orb = cv2.ORB_create()
    kp_a, desc_a = orb.detectAndCompute(img1, None)
    kp_b, desc_b = orb.detectAndCompute(img2, None)

    # Brute-force matcher
    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:
        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
        print(f"MobileNet similarity score is {sim_score}")
        return float(sim_score)
    except Exception as e:
        print("Erro ao calcular similaridade com MobileNet")
        return 0

def load_image(source):
    Image.MAX_IMAGE_PIXELS = None

    def extract_frame_from_video(video_path_or_url, time_sec):
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_frame:
            frame_path = temp_frame.name

        command = [
            ffmpeg_path,
            "-ss", str(time_sec),
            "-i", video_path_or_url,
            "-frames:v", "1",
            "-q:v", "2",
            "-y",
            frame_path
        ]
        subprocess.run(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

        if not os.path.exists(frame_path):
            raise ValueError("Failed to extract frame from video.")

        frame = cv2.imread(frame_path, cv2.IMREAD_GRAYSCALE)
        os.remove(frame_path)
        return frame

    if source.startswith('http'):
        mime_type, _ = mimetypes.guess_type(source)
        if mime_type and mime_type.startswith('video'):
            return extract_frame_from_video(source, frame_time)

        # Assume imagem
        response = requests.get(source)
        img = np.asarray(bytearray(response.content), dtype=np.uint8)
        img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)
        return img

    else:
        try:
            img_bytes = base64.b64decode(source)
            mime_type = mimetypes.guess_type("data")[0]  # fallback

            # Tenta abrir como imagem
            img = Image.open(BytesIO(img_bytes))
            img = np.array(img)
            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
            return img
        except Exception:
            # Se falhar, assumimos que é vídeo em base64 e extraímos frame
            with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_video:
                temp_video.write(base64.b64decode(source))
                temp_video_path = temp_video.name

            frame = extract_frame_from_video(temp_video_path, frame_time)
            os.remove(temp_video_path)
            return frame

app = FastAPI()

BASE_DIR = "/tmp/data"

@app.post("/save")
async def save(image_data: RequestModel):
    os.makedirs(BASE_DIR, exist_ok=True)
    filename = os.path.join(BASE_DIR, f"{image_data.originId}_{image_data.assetCode}.json")

    img1 = load_image(image_data.originSource)
    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(img1, img2)
        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)
    return True

@app.get("/files")
async def list_files():
    try:
        files_data = []
        for filename in os.listdir(BASE_DIR):
            filepath = os.path.join(BASE_DIR, filename)
            if os.path.isfile(filepath):
                try:
                    with open(filepath, "r") as f:
                        file_content = f.read()  # Lê o conteúdo do ficheiro
                        # Tenta decodificar o conteúdo como JSON, se possível
                        try:
                            file_content_json = json.loads(file_content)
                            files_data.append({"filename": filename, "content": file_content_json})
                        except json.JSONDecodeError:
                            files_data.append({"filename": filename, "content": file_content}) # Se não for JSON, retorna o texto
                except (IOError, OSError) as e:
                    raise HTTPException(status_code=500, detail=f"Erro ao ler o ficheiro {filename}: {e}")

        return JSONResponse({"files_data": files_data})
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="Diretório de dados não encontrado")

@app.get("/files/similar")
async def list_similar_files():
    try:
        files_data = []
        for filename in os.listdir(BASE_DIR):
            filepath = os.path.join(BASE_DIR, filename)
            if os.path.isfile(filepath):
                try:
                    with open(filepath, "r") as f:
                        file_content = f.read()
                        try:
                            file_content_json = json.loads(file_content)
                            # Check for similarityOrb and filter
                            if "similarityOrb" in file_content_json and file_content_json["similarityOrb"] > 0:
                                files_data.append({"filename": filename, "content": file_content_json})
                        except json.JSONDecodeError:
                            pass  # Skip files that are not valid JSON
                except (IOError, OSError) as e:
                    raise HTTPException(status_code=500, detail=f"Erro ao ler o ficheiro {filename}: {e}")
        return JSONResponse({"files_data": files_data})
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="Diretório de dados não encontrado")

@app.get("/files/find/{origin_id}")
async def get_file_by_origin_id(origin_id: int):
    try:
        for filename in os.listdir(BASE_DIR):
            if filename.startswith(f"{origin_id}_") and filename.endswith(".json"):
                filepath = os.path.join(BASE_DIR, filename)
                if os.path.isfile(filepath):
                    try:
                        with open(filepath, "r") as f:
                            file_content = f.read()
                            try:
                                file_content_json = json.loads(file_content)
                                return JSONResponse({"filename": filename, "content": file_content_json})
                            except json.JSONDecodeError:
                                return JSONResponse({"filename": filename, "content": file_content})
                    except (IOError, OSError) as e:
                        raise HTTPException(status_code=500, detail=f"Erro ao ler o ficheiro {filename}: {e}")
        raise HTTPException(status_code=404, detail=f"Ficheiro com originId '{origin_id}' não encontrado")
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="Diretório de dados não encontrado")