Tentando segmentar novamente.

Dockerfile

@@ -4,6 +4,7 @@ FROM python:3.10
 #Variáveis de Ambiente para Caches
 ENV HUGGINGFACE_HUB_CACHE="/tmp/huggingface"
 ENV NUMBA_CACHE_DIR="/tmp/numba_cache"
+ENV U2NET_HOME="/tmp/u2net"
 
 #Define um diretório de trabalho limpo
 WORKDIR /app

feature_extractor_single.py

@@ -4,6 +4,7 @@ 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)
 class FeatureExtractor:
@@ -11,7 +12,7 @@ class FeatureExtractor:
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         print(f"Usando dispositivo: {self.device}")
 
-        #Modelo e pré-processador
+        # Modelo e pré-processador
         self.processor = ConvNextImageProcessor.from_pretrained(
             "facebook/convnext-large-224-22k-1k"
         )
@@ -24,35 +25,56 @@ class FeatureExtractor:
         self.model.eval()
 
     def extract_convnext(self, image_path: str) -> np.ndarray:
-        image = Image.open(image_path).convert("RGB")
-        inputs = self.processor(image, return_tensors="pt").to(self.device)
+        #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
+
+        #4. Passa para o ConvNeXt (que já faz o resize e normalize internamente)
+        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
+#Função principal chamada pelo main.py
 def process_single_image(image_path: str, output_dir: str = "processed"):
     """
-    Extração de features ConvNeXt (sem segmentação)
+    Pipeline: Segmentação (Rembg) -> Fundo preto -> ConvNeXt
     """
-    #1) Extração de features
     extractor = FeatureExtractor()
-    #Passa o CAMINHO ORIGINAL da imagem (ex: /tmp/temp_foto.jpg)
     features = extractor.extract_convnext(image_path) 
+    return features
 
-    #NÃO PRECISAMOS MAIS SALVAR O ARQUIVO .NPY
-    #feat_path = os.path.join(output_dir, "features_single.npy")
-    #np.save(feat_path, features)
-    #print(f"Vetor de características salvo em: {feat_path}")
-
-    return features #RETORNA O ARRAY DE DADOS
-
-#Execução direta
+#Execução direta para testes locais
 if __name__ == "__main__":
-    #Caminho da imagem
-    image_path = r"C:\Users\Julia\Downloads\clb2_600px.jpg"
-
-    #Executa pipeline completo
-    process_single_image(image_path)
+    #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.")

requirements.txt

@@ -20,4 +20,8 @@ umap-learn==0.5.6
 joblib==1.4.2
 
 #Utilitários gerais
-requests==2.32.3
+requests==2.32.3
+
+#Segmentação
+rembg
+onnxruntime

svm_model.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b14bfa70c436e570e1c2a7c567561b9b57c7231db953e44047a34437d9a8b697
-size 10940308
+oid sha256:7220864ee09603bf8abd9bd233a5bd0bf31e9eca496465973449887b4abcd8c6
+size 3399533

umap_reducer.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:96c669ea481469c49a9b1f3401fb3d802020c6f9add7e67d62bff5c86e6c2326
-size 270523735
+oid sha256:f454d8c21b3982fce91bcd2fb679c1ec9d0c9c326ea5059cccee799b58853b20
+size 255570333