similarity_dbg / main.py
MarioPrzBasto's picture
Update main.py
de0f224 verified
raw
history blame
4.58 kB
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 httpx
BASE_DIR = "saved_data"
app = FastAPI()
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 < 50]
if len(matches) == 0:
return 0
return len(similar_regions) / len(matches)
async def load_image(image_path: str):
"""Carrega uma imagem a partir de um caminho de arquivo ou URL."""
try:
if image_path.startswith("http://") or image_path.startswith("https://"):
async with httpx.AsyncClient() as client:
response = await client.get(image_path)
response.raise_for_status()
image_bytes = np.frombuffer(response.content, np.uint8)
img = cv2.imdecode(image_bytes, cv2.IMREAD_COLOR)
return img
else:
img = cv2.imread(image_path)
return img
except Exception as e:
print(f"Erro ao carregar a imagem {image_path}: {e}")
return None
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 = await load_image(image_data.originSource)
img2 = await load_image(image_data.source)
similarity_orb = 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}")
data_to_save = image_data.dict()
if similarity_orb is not None:
data_to_save["similarityOrb"] = similarity_orb
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/{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")