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
import uuid


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 = "/tmp/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
        return float(sim_score)
    except Exception as e:
        print("Erro ao calcular similaridade com MobileNet")
        return 0

def load_image(source, assetCode, contentType=None, ffmpeg_path='ffmpeg', frame_time=1):
    Image.MAX_IMAGE_PIXELS = None

    def extract_frame_from_video(video_path_or_url, time_sec):
        print(f"[INFO] A extrair frame do vídeo: {video_path_or_url} no segundo {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
        ]

        print(f"[DEBUG] Comando ffmpeg: {' '.join(command)}")

        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        if result.returncode != 0:
            print(f"[ERRO] ffmpeg falhou com código {result.returncode}")
            print(f"[ERRO] stderr: {result.stderr.decode('utf-8')}")
            raise RuntimeError("Erro ao extrair frame com ffmpeg.")

        if not os.path.exists(frame_path):
            print("[ERRO] Frame não criado. Verifica se o caminho do vídeo está correto e acessível.")
            raise ValueError("Frame não encontrado após execução do ffmpeg.")

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

        if frame is None:
            print("[ERRO] Falha ao ler frame extraído com OpenCV.")
            raise ValueError("Erro ao carregar frame extraído.")

        print(f"[SUCESSO] Frame extraído com sucesso de {video_path_or_url}")
        return frame

    try:
        if source.startswith('http'):
            print(f"[INFO] Content-Type de {assetCode} é {contentType}")

            if contentType and contentType.startswith('video'):
                return extract_frame_from_video(source, frame_time)

            print(f"[INFO] A carregar imagem {assetCode} a partir de URL")
            response = requests.get(source)
            img = np.asarray(bytearray(response.content), dtype=np.uint8)
            img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)
            return img

        else:
            print(f"[INFO] A tentar carregar base64 de {assetCode} como imagem ou vídeo.")

            try:
                img_bytes = base64.b64decode(source)

                if contentType and contentType.startswith('image'):
                    print(f"[INFO] Base64 de {assetCode} identificado como imagem")
                    img = Image.open(BytesIO(img_bytes))
                    img = np.array(img)
                    img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
                    return img
                else:
                    print(f"[INFO] Base64 de {assetCode} identificado como vídeo")
                    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_video:
                        temp_video.write(img_bytes)
                        temp_video_path = temp_video.name

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

            except Exception as e:
                print(f"[ERRO] Falha ao processar base64 de {assetCode}: {e}")
                raise

    except Exception as e:
        print(f"[ERRO] Falha ao carregar imagem para {assetCode}: {e}")
        return None

@app.post("/save")
async def save(image_data: RequestModel):
    data_to_save = image_data.dict()
    print("Recebido:", data_to_save)
    os.makedirs(BASE_DIR, exist_ok=True)
    filename = os.path.join(BASE_DIR, f"{image_data.originId}_{image_data.assetCode}_{uuid.uuid4().hex[:8]}.json")

    img1 = load_image(image_data.originSource, f"origin {image_data.originSource}")
    img2 = load_image(image_data.source, image_data.assetCode, image_data.contentType)

    similarity_orb = None
    similarity_mobilenet = None

    if img1 is not None and img2 is not None:
        similarity_orb = orb_sim(img1, img2)
        similarity_mobilenet = mobilenet_similarity(img1, img2)

    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/{search}")
async def search_file(search: str):
    try:
        files_data = []
        for filename in os.listdir(BASE_DIR):
            if f"{search}" in filename 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)
                                files_data.append({"filename": filename, "content": file_content_json})
                            except json.JSONDecodeError:
                                files_data.append({"filename": filename, "content": file_content})
                    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")