Add endpoint de debug pra visualizar img segmentada.

feature_extractor_single.py

@@ -2,17 +2,42 @@ import os
 import torch
 import numpy as np
 from PIL import Image
-from torchvision import transforms
 from transformers import ConvNextImageProcessor, ConvNextForImageClassification
 from rembg import remove
 
-#Classe para extração de features (ConvNeXt)
+#FUNÇÃO AUXILIAR DE SEGMENTAÇÃO (REUTILIZÁVEL)
+def segment_image(pil_image: Image.Image) -> Image.Image:
+    """
+    Recebe uma imagem PIL, remove o fundo e coloca fundo preto.
+    Retorna a imagem PIL tratada.
+    """
+    try:
+        #1. Remove o fundo (Rembg)
+        img_no_bg = remove(pil_image)
+        
+        #2. Composição em fundo preto
+        #Cria uma imagem totalmente preta do mesmo tamanho
+        fundo_preto = Image.new("RGB", img_no_bg.size, (0, 0, 0))
+        
+        #Usa o canal Alpha como máscara
+        if img_no_bg.mode == 'RGBA':
+            mask = img_no_bg.split()[3]
+            fundo_preto.paste(img_no_bg, mask=mask)
+            return fundo_preto
+        else:
+            return img_no_bg.convert("RGB")
+            
+    except Exception as e:
+        print(f"AVISO: Falha na segmentação ({e}). Retornando original.")
+        return pil_image.convert("RGB")
+
+
+#CLASSE EXTRATORA
 class FeatureExtractor:
     def __init__(self, device=None):
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         print(f"Usando dispositivo: {self.device}")
 
-        # Modelo e pré-processador
         self.processor = ConvNextImageProcessor.from_pretrained(
             "facebook/convnext-large-224-22k-1k"
         )
@@ -20,61 +45,25 @@ class FeatureExtractor:
             "facebook/convnext-large-224-22k-1k"
         ).to(self.device)
 
-        #Remove camada de classificação (ficam só as features)
         self.model.classifier = torch.nn.Identity()
         self.model.eval()
 
     def extract_convnext(self, image_path: str) -> np.ndarray:
-        #1. Abre a imagem original
         print(f"Processando imagem: {os.path.basename(image_path)}")
         input_img = Image.open(image_path).convert("RGB")
         
-        #2. SEGMENTAÇÃO (remoção do fundo)
-        #O rembg devolve uma imagem RGBA (com transparência)
-        try:
-            img_no_bg = remove(input_img)
-            
-            #3. COMPOSIÇÃO EM FUNDO PRETO
-            #Cria uma imagem totalmente preta do mesmo tamanho
-            fundo_preto = Image.new("RGB", img_no_bg.size, (0, 0, 0))
-            
-            #Usa o canal Alpha (transparência) da imagem recortada como máscara
-            #Onde for folha, cola a folha. Onde for transparente, mantém o preto.
-            mask = img_no_bg.split()[3] #Pega o 4º canal (Alpha)
-            fundo_preto.paste(img_no_bg, mask=mask)
-            
-            final_image = fundo_preto
-            print("Fundo removido e substituído por preto com sucesso.")
-            
-        except Exception as e:
-            print(f"AVISO: Falha na segmentação ({e}). Usando imagem original.")
-            final_image = input_img
+        #1. CHAMA A FUNÇÃO DE SEGMENTAÇÃO
+        final_image = segment_image(input_img)
 
-        #4. Passa para o ConvNeXt (que já faz o resize e normalize internamente)
+        #2. Passa para o ConvNeXt
         inputs = self.processor(final_image, return_tensors="pt").to(self.device)
         
         with torch.no_grad():
             features = self.model(**inputs).logits
             
         features_np = features.cpu().numpy().flatten()
-        print(f"Vetor de características extraído com shape: {features_np.shape}")
-        
         return features_np
 
-#Função principal chamada pelo main.py
-def process_single_image(image_path: str, output_dir: str = "processed"):
-    """
-    Pipeline: Segmentação (Rembg) -> Fundo preto -> ConvNeXt
-    """
+def process_single_image(image_path: str):
     extractor = FeatureExtractor()
-    features = extractor.extract_convnext(image_path) 
-    return features
-
-#Execução direta para testes locais
-if __name__ == "__main__":
-    #Teste com uma imagem local
-    image_path = "teste_folha.jpg" #Mudar para um arquivo real se for testar
-    if os.path.exists(image_path):
-        process_single_image(image_path)
-    else:
-        print("Imagem de teste não encontrada.")
+    return extractor.extract_convnext(image_path)

main.py

@@ -5,9 +5,12 @@ import json
 import numpy as np
 from fastapi import FastAPI, File, UploadFile, Depends, HTTPException, status
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from feature_extractor_single import process_single_image
 from datetime import datetime
 import unicodedata
+import io
+from PIL import Image
+from starlette.responses import StreamingResponse
+from feature_extractor_single import process_single_image, segment_image
 
 def normalize_string(s: str) -> str:
     """
@@ -164,4 +167,30 @@ async def extract_features(file: UploadFile = File(...), token: str = Depends(ve
     os.remove(temp_path)
 
     features_list = features_array.tolist()
-    return {"features": features_list}
+    return {"features": features_list}
+
+@app.post("/debug/view_segmentation/")
+async def view_segmentation(file: UploadFile = File(...), token: str = Depends(verify_token)):
+    """
+    Endpoint de debug.
+    Retorna a imagem processada (fundo preto) para verificação visual.
+    Útil para saber o que o modelo está "enxergando".
+    """
+    try:
+        #1. Lê a imagem da memória (sem salvar no disco pra ser rápido)
+        contents = await file.read()
+        pil_image = Image.open(io.BytesIO(contents)).convert("RGB")
+
+        #2. Aplica a mesma lógica de segmentação do modelo
+        processed_image = segment_image(pil_image)
+
+        #3. Salva a imagem processada em um buffer de memória (bytes)
+        img_byte_arr = io.BytesIO()
+        processed_image.save(img_byte_arr, format='JPEG', quality=95)
+        img_byte_arr.seek(0)
+
+        #4. Retorna como uma stream de imagem (O navegador/Swagger exibe isso!)
+        return StreamingResponse(img_byte_arr, media_type="image/jpeg")
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Erro ao processar imagem: {e}")