Update app.py

app.py

@@ -1,74 +1,63 @@
-import os
-os.environ["NUMBA_CACHE_DIR"] = "/tmp/numba_cache"
-os.environ["U2NET_HOME"] = "/tmp/u2net"
-
+import spacy
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from io import BytesIO
-from PIL import Image, ImageFilter
-import rembg
-import onnxruntime as ort
-from concurrent.futures import ThreadPoolExecutor
-import asyncio
-import gc
-import base64
+import nltk
+from nltk.tokenize import word_tokenize
+from nltk.corpus import stopwords
+from collections import Counter
+import re
 
-app = FastAPI()
+# Baixar as stopwords do NLTK (se ainda não tiver feito)
+nltk.download('punkt')
+nltk.download('stopwords')
 
-# Configurações do onnxruntime para CPU
-options = ort.SessionOptions()
-options.intra_op_num_threads = 2  # Limita o número de threads para evitar sobrecarga
-options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL  # Execução sequencial para melhor desempenho em CPU
+# Carregar o modelo spaCy para reconhecimento de entidades nomeadas
+nlp = spacy.load("en_core_web_sm")
+
+app = FastAPI()
 
-# Pool de threads para executar tarefas bloqueantes
-executor = ThreadPoolExecutor(max_workers=4)
+class PostText(BaseModel):
+    post: str  # O post (texto) que será enviado para a API
 
-def resize_image(image, max_size=512):
-    """Redimensiona a imagem para uma largura máxima de 512px, mantendo a proporção."""
-    width, height = image.size
-    if width > max_size or height > max_size:
-        ratio = min(max_size / width, max_size / height)
-        new_size = (int(width * ratio), int(height * ratio))
-        image = image.resize(new_size, Image.Resampling.LANCZOS)
-    return image
+def clean_text(text):
+    """Remove caracteres especiais e faz a limpeza do texto."""
+    text = re.sub(r'[^\w\s]', '', text)  # Remove pontuação
+    text = text.lower()  # Converte para minúsculas
+    return text
 
-def process_image(image_data):
-    # Decodifica a imagem base64
-    image = Image.open(BytesIO(base64.b64decode(image_data)))
+def extract_keywords(text):
+    """Extrai palavras-chave usando spaCy e nltk."""
     
-    # Pré-processamento: apenas redimensiona para 512px
-    image = resize_image(image, max_size=512)
+    # Limpeza inicial do texto
+    cleaned_text = clean_text(text)
     
-    # Remove o fundo da imagem usando rembg
-    output = rembg.remove(image, session_options=options)
+    # Tokenização do texto
+    words = word_tokenize(cleaned_text)
     
-    # Pós-processamento: suaviza as bordas
-    output = output.filter(ImageFilter.SMOOTH_MORE)
+    # Remover stopwords
+    stop_words = set(stopwords.words("english"))
+    filtered_words = [word for word in words if word not in stop_words]
+
+    # Contar a frequência das palavras filtradas
+    word_counts = Counter(filtered_words)
     
-    # Converte a imagem de volta para bytes
-    img_byte_arr = BytesIO()
-    output.save(img_byte_arr, format='PNG')
-    img_byte_arr.seek(0)
+    # Processar entidades nomeadas com spaCy (ex.: pessoas, locais, eventos)
+    doc = nlp(text)
+    entities = [ent.text for ent in doc.ents]
     
-    # Codifica a imagem processada em base64
-    return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
-
-class ImageBase64(BaseModel):
-    image: str  # A imagem em base64
+    # Juntar as palavras mais frequentes e as entidades encontradas
+    keywords = set(filtered_words + entities)
+    
+    # Ordenar e retornar as palavras-chave mais relevantes (top 10)
+    return [keyword for keyword, _ in word_counts.most_common(10)] + entities[:10]
 
-@app.post("/remove-background")
-async def remove_background(image_data: ImageBase64):
+@app.get("/generate-keywords")
+async def generate_keywords(post_text: PostText):
     try:
-        # Executa o processamento da imagem em um thread separado
-        loop = asyncio.get_event_loop()
-        processed_image_base64 = await loop.run_in_executor(executor, process_image, image_data.image)
+        # Gerar as palavras-chave a partir do texto do post
+        keywords = extract_keywords(post_text.post)
         
-        # Retorna a imagem processada em base64
-        return {"processed_image": processed_image_base64}
+        return {"keywords": keywords}
     
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
-    
-    finally:
-        # Força a coleta de lixo após cada requisição
-        gc.collect()